diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index a7852a52b4..5221425e4b 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -267,7 +267,6 @@ def _apply_fp8_training(transformer, on_event) -> bool: LoRA modules. Never fatal: on any failure the run continues in bf16 with a warning.""" try: from torchao.float8 import Float8LinearConfig, convert_to_float8_training - convert_to_float8_training( transformer, module_filter_fn = _fp8_module_filter, @@ -285,25 +284,41 @@ def _pick_auto_precision(prequant, device, free_gb, dense_gb, capability, has_fp free VRAM at decision time. bf16 + regional compile is the measured speed winner (2.3-2.6x over nf4 on B200); fp8 stays an explicit opt-in because torchao float8's dynamic-scaling overhead made it SLOWER than compiled bf16 at LoRA-training shapes on - the same hardware. ``capability``/``has_fp8`` remain parameters so the policy can be - revisited per GPU generation without changing callers.""" + the same hardware. int8 must still materialise the full bf16 transformer before + ``quantize_`` shrinks it module-by-module, so its band requires the dense-load + transient (1.15x dense) to fit -- what int8 buys in that band is steady-state + headroom for activations and the latent cache, not load-time memory. + ``capability``/``has_fp8`` remain parameters so the policy can be revisited per GPU + generation without changing callers.""" _ = capability, has_fp8 if prequant or device != "cuda" or not free_gb or not dense_gb: return "nf4" - headroom = 1.5 - if free_gb > dense_gb * headroom: + if free_gb > dense_gb * 1.5: return "bf16" - if free_gb > dense_gb * 0.55 * headroom: + if free_gb > dense_gb * 1.15: return "int8" return "nf4" def _resolve_base_precision(cfg, spec, device) -> str: """Resolve "auto" against the live GPU (free VRAM measured BEFORE anything loads); - explicit modes pass through (normalized() already validated them).""" + explicit modes pass through (normalized() already validated them against the repo and + compute dtype) but are re-checked against the live device here: the dense modes are + CUDA-only, and /info never advertises them on a host without a GPU, so an explicit + request from a stale or direct client fails fast instead of loading a full dense + transformer onto the CPU.""" mode = (cfg.base_precision or "nf4").strip().lower() if mode != "auto": + if mode in ("bf16", "int8", "fp8") and device != "cuda": + raise ValueError( + f"base_precision={mode!r} needs a CUDA GPU; this host has none. " + f"Use base_precision='nf4' or 'auto'." + ) return mode + # auto may only resolve to the dense modes when the run uses bf16 compute, mirroring + # the normalized() rule for explicit dense modes; otherwise stay on the nf4 floor. + if getattr(cfg, "mixed_precision", "bf16") != "bf16": + return "nf4" prequant = repo_is_prequantized(cfg.base_model) free_gb = None capability = None @@ -367,7 +382,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 +510,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 +616,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 +789,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 @@ -796,7 +826,12 @@ def _sample_cached_latents(cache, idxs, variant_rng, device): return lat_a + lat_b * torch.randn_like(lat_a) -def _should_compile(cfg, base_is_bnb, device, base_precision = "nf4") -> bool: +def _should_compile( + cfg, + base_is_bnb, + device, + base_precision = "nf4", +) -> bool: mode = (cfg.compile_transformer or "auto").strip().lower() if device != "cuda" or mode == "off": return False @@ -813,7 +848,12 @@ def _should_compile(cfg, base_is_bnb, device, base_precision = "nf4") -> bool: def _maybe_compile_transformer( - transformer, cfg, base_is_bnb, device, on_event, base_precision = "nf4" + transformer, + cfg, + base_is_bnb, + device, + on_event, + base_precision = "nf4", ) -> bool: """Regionally compile the transformer blocks (diffusers compile_repeated_blocks) after the LoRA is attached. Never fatal: a wrap failure falls back to eager with a warning @@ -831,7 +871,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 +961,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 +1017,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 +1143,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 +1220,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..56b1020cd5 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -140,7 +140,6 @@ def train_precision_modes() -> tuple[list[str], str]: recommended = "nf4" try: import torch - if torch.cuda.is_available(): modes += ["bf16", "int8"] major, minor = torch.cuda.get_device_capability() @@ -434,11 +433,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 +458,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 +482,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 +520,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_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index 6c59a84fee..af6789b199 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -105,9 +105,13 @@ def test_pick_auto_precision_policy_table(): assert p(False, "cuda", 140, 23.8, (8, 0), True) == "bf16" assert p(False, "cuda", 140, 23.8, (10, 0), False) == "bf16" - # Middle band (25 > 23.8 * 0.55 * 1.5 = 19.6, but not > 23.8 * 1.5 = 35.7) -> int8. - assert p(False, "cuda", 25, 23.8, (10, 0), True) == "int8" - # Too little free VRAM for even int8 -> nf4. + # Middle band (30 > 23.8 * 1.15 = 27.4, but not > 23.8 * 1.5 = 35.7) -> int8. + assert p(False, "cuda", 30, 23.8, (10, 0), True) == "int8" + # int8 still materialises the full bf16 transformer before quantize_ shrinks it, so + # free VRAM below the dense-load transient (25 < 27.4) must fall back to nf4 even + # though the QUANTIZED weights would have fit. + assert p(False, "cuda", 25, 23.8, (10, 0), True) == "nf4" + # Too little free VRAM for any dense load -> nf4. assert p(False, "cuda", 10, 23.8, (10, 0), True) == "nf4" @@ -118,8 +122,22 @@ def test_resolve_base_precision_passes_explicit_through(): spec = dit._SPECS["flux.1"] cfg = _cfg(base_precision = "bf16") assert dit._resolve_base_precision(cfg, spec, "cuda") == "bf16" - # Device is irrelevant for an explicit mode: still bf16 on cpu. - assert dit._resolve_base_precision(cfg, spec, "cpu") == "bf16" + + # The dense modes are CUDA-only: an explicit request on a GPU-less host fails fast + # (before any model load) instead of silently proceeding; /info never advertised it. + with pytest.raises(ValueError, match = "CUDA"): + dit._resolve_base_precision(cfg, spec, "cpu") + # nf4 stays a passthrough on any device (the bnb load path owns its own errors). + assert dit._resolve_base_precision(_cfg(base_precision = "nf4"), spec, "cpu") == "nf4" + + +def test_resolve_auto_requires_bf16_compute(): + # auto may resolve to bf16/int8 which train in bf16 compute, so a non-bf16 + # mixed_precision pins auto to the nf4 floor BEFORE any GPU probe (pure, no CUDA + # needed here) -- mirroring the normalized() rule for explicit dense modes. + spec = dit._SPECS["flux.1"] + cfg = _cfg(base_precision = "auto", mixed_precision = "fp16") + assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4" # ── _fp8_module_filter ──────────────────────────────────────────────────────── @@ -154,7 +172,6 @@ def test_train_precision_modes_no_cuda(monkeypatch): # Patch the torch module attribute the function imports so it observes a CPU-only box: # no CUDA -> the nf4-only floor with nf4 recommended, and it never raises. import torch - monkeypatch.setattr(torch.cuda, "is_available", lambda: False) assert train_precision_modes() == (["nf4"], "nf4") 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