diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index a7852a52b4..2b506d0ec2 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -367,7 +367,12 @@ def _flux_encode_latent_stats(vae, pixel_values): return (dist.mean - vae.config.shift_factor) * scale, dist.std * scale -def _flux_collate(entries, device, weight_dtype, pad_to = None): +def _flux_collate( + entries, + device, + weight_dtype, + pad_to = None, +): import torch # FLUX embeds are fixed-length (encode_prompt pads to max_sequence_length), so a plain @@ -490,7 +495,12 @@ def _qwen_encode_latent_stats(vae, pixel_values): return (dist.mean - mean) / std, dist.std / std -def _qwen_collate(entries, device, weight_dtype, pad_to = None): +def _qwen_collate( + entries, + device, + weight_dtype, + pad_to = None, +): import torch import torch.nn.functional as F @@ -591,7 +601,12 @@ def _zimage_encode_latent_stats(vae, pixel_values): return _zimage_encode_latents(vae, pixel_values), None -def _zimage_collate(entries, device, weight_dtype, pad_to = None): +def _zimage_collate( + entries, + device, + weight_dtype, + pad_to = None, +): caps = [e[0].to(device = device, dtype = weight_dtype) for e in entries] return (caps,) @@ -759,7 +774,7 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev total = len(image_paths) for i, path in enumerate(image_paths): variants = [] - for (u_left, u_top, flip) in plan[i]: + for u_left, u_top, flip in plan[i]: px = ( _load_pixel_tensor_planned( path, cfg.resolution, cfg.center_crop, u_left, u_top, flip @@ -831,7 +846,9 @@ def _maybe_compile_transformer( fn = getattr(transformer, "compile_repeated_blocks", None) if not callable(fn): - _emit(on_event, "warning", message = "torch.compile unavailable for this model; running eager.") + _emit( + on_event, "warning", message = "torch.compile unavailable for this model; running eager." + ) return False try: dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None) @@ -919,7 +936,14 @@ def run_dit_lora_training( perf_snap = _apply_perf_flags(cfg, device) try: return _train_dit( - cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_stop, + cfg, + spec, + pairs, + rng, + device, + weight_dtype, + on_event, + _check_stop, lambda: save_on_stop, ) finally: @@ -968,8 +992,12 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto ) if latent_cache is None: # stopped during the cache build; nothing trained yet _emit( - on_event, "complete", output_dir = str(out_dir), lora_path = None, - stopped = True, steps_run = 0, + on_event, + "complete", + output_dir = str(out_dir), + lora_path = None, + stopped = True, + steps_run = 0, ) return str(out_dir) try: @@ -1090,7 +1118,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto noisy = (1.0 - sigmas) * latents + sigmas * noise embeds = spec.collate( - [caption_embeds[captions[i]] for i in idxs], device, weight_dtype, + [caption_embeds[captions[i]] for i in idxs], + device, + weight_dtype, pad_to = qwen_pad_to, ) with autocast: @@ -1165,6 +1195,7 @@ def _make_optimizer(params, lr): regression for LoRA -- else torch AdamW, fused on CUDA (with a fallback when this build/device lacks the fused kernel).""" import torch + try: import bitsandbytes as bnb return bnb.optim.AdamW8bit(params, lr = lr) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index ceb3802ea7..7c7ca79d51 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -201,7 +201,7 @@ def _build_sdxl_latent_cache( total = len(image_paths) for i, path in enumerate(image_paths): variants = [] - for (u_left, u_top, flip) in plan[i]: + for u_left, u_top, flip in plan[i]: tensor, time_ids = _load_image_tensor_planned( path, cfg.resolution, cfg.center_crop, u_left, u_top, flip ) @@ -390,7 +390,13 @@ def run_diffusion_lora_training( latent_cache = None if use_cache: latent_cache = _build_sdxl_latent_cache( - vae, vae_scale, [p for p, _ in pairs], cfg, device, weight_dtype, on_event, + vae, + vae_scale, + [p for p, _ in pairs], + cfg, + device, + weight_dtype, + on_event, _check_stop, ) if latent_cache is None: # stopped during the cache build; nothing trained yet diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index c2bcab4bcf..91901ba382 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -434,11 +434,7 @@ def _emit(on_event: Optional[EventCb], type_: str, **kw: Any) -> None: def _plan_cache_variants( - num_images: int, - cache_variants: int, - center_crop: bool, - random_flip: bool, - seed: int, + num_images: int, cache_variants: int, center_crop: bool, random_flip: bool, seed: int ) -> list[list[tuple[float, float, bool]]]: """Seed-deterministic crop/flip plan for the latent cache: per image, up to ``cache_variants`` draws of (u_left, u_top, flip) with the crop as unit fractions the @@ -463,10 +459,13 @@ def _plan_cache_variants( def _apply_perf_flags( - cfg: "DiffusionLoraConfig", device: str, cudnn_benchmark: bool = False + cfg: "DiffusionLoraConfig", + device: str, + cudnn_benchmark: bool = False, ) -> dict: """Set the run-scoped torch backend knobs: TF32 matmuls + high fp32 matmul precision - (under ``cfg.enable_tf32``), plus cudnn autotuning when the caller opts in. Autotune is + when ``cfg.enable_tf32`` is on, strict fp32 (all TF32 flags cleared) when it is off, + plus cudnn autotuning when the caller opts in. Autotune is for the conv-heavy SDXL U-Net only: measured on B200, it DOUBLES peak VRAM (fp32 VAE conv workspaces) while the DiT loop -- pure matmuls once the latent cache is built -- gains nothing from it. Returns a snapshot for ``_restore_perf_flags``. Best-effort: @@ -484,6 +483,12 @@ def _apply_perf_flags( torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.set_float32_matmul_precision("high") + else: + # The opt-out is a strict-fp32 A/B mode, so actively clear the flags rather + # than inherit ambient state (cudnn TF32 defaults to ON in torch). + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.set_float32_matmul_precision("highest") if cudnn_benchmark: torch.backends.cudnn.benchmark = True # The cuDNN SDPA backend's TRAINING graph is broken for the FLUX attention shapes @@ -516,8 +521,15 @@ def _restore_perf_flags(snap: Optional[dict]) -> None: if snap.get("matmul_precision"): torch.set_float32_matmul_precision(snap["matmul_precision"]) - if snap.get("cudnn_sdp"): - torch.backends.cuda.enable_cudnn_sdp(True) + # Restore the exact pre-run cudnn SDPA state; None means the flag was unreadable + # (or absent) at apply time and was never touched. + cuda_backends = getattr(torch.backends, "cuda", None) + if ( + snap.get("cudnn_sdp") is not None + and cuda_backends is not None + and hasattr(cuda_backends, "enable_cudnn_sdp") + ): + cuda_backends.enable_cudnn_sdp(bool(snap["cudnn_sdp"])) except Exception: # noqa: BLE001 -- best-effort restore pass diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py index 27a184a2fa..d430601abd 100644 --- a/studio/backend/tests/test_diffusion_train_perf.py +++ b/studio/backend/tests/test_diffusion_train_perf.py @@ -82,9 +82,7 @@ def test_plan_cache_variants_deterministic_and_deduped(): # ── per-family collate fns ──────────────────────────────────────────────────── def test_flux_collate_shapes(): # FLUX embeds are fixed length: 3 entries batch by a plain cat; text_ids are shared. - entries = [ - (torch.randn(1, 512, 32), torch.randn(1, 16), torch.randn(512, 3)) for _ in range(3) - ] + entries = [(torch.randn(1, 512, 32), torch.randn(1, 16), torch.randn(512, 3)) for _ in range(3)] pe, pooled, text_ids = _flux_collate(entries, "cpu", torch.float32) assert pe.shape == (3, 512, 32) assert pooled.shape == (3, 16) @@ -242,9 +240,7 @@ def test_service_stop_save_flag(): # ── preparing / warning events + stopped completion messages ────────────────── def test_apply_event_preparing_and_warning(): svc = DiffusionTrainingService() - svc._apply_event( - {"type": "preparing", "stage": "cache_latents", "done": 4, "total": 8} - ) + svc._apply_event({"type": "preparing", "stage": "cache_latents", "done": 4, "total": 8}) st = svc.status() assert st["status"] == "running" assert st["in_model_load"] is True @@ -335,3 +331,30 @@ def test_perf_flags_cpu_roundtrip(): snap = _apply_perf_flags(_cfg(), "cpu") assert isinstance(snap, dict) _restore_perf_flags(snap) # no exception + + +def test_perf_flags_tf32_off_clears_flags(): + # enable_tf32=False is the strict-fp32 A/B mode: it must actively clear the TF32 flags + # (cudnn TF32 defaults ON in torch) rather than inherit ambient state, and restore must + # put the ambient values back. The flag attributes are plain Python state, present and + # settable on CPU-only torch builds, so this runs without a GPU. + import torch + + before = ( + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + torch.get_float32_matmul_precision(), + ) + snap = _apply_perf_flags(_cfg(enable_tf32 = False), "cuda") + try: + assert torch.backends.cuda.matmul.allow_tf32 is False + assert torch.backends.cudnn.allow_tf32 is False + assert torch.get_float32_matmul_precision() == "highest" + finally: + _restore_perf_flags(snap) + after = ( + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + torch.get_float32_matmul_precision(), + ) + assert after == before