diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 44597ce561..b4aa754402 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -279,7 +279,9 @@ def _apply_fp8_training(transformer, on_event) -> bool: return False -def _pick_auto_precision(prequant, device, free_gb, dense_gb, capability, has_fp8) -> str: +def _pick_auto_precision( + prequant, device, free_gb, dense_gb, capability, has_fp8, has_torchao = True +) -> str: """Pure policy for base_precision="auto": nf4 for a prequant base or no CUDA; else the fastest dense mode whose weights + headroom (activations, optimizer, cache) fit the free VRAM at decision time. bf16 + regional compile is the measured speed winner @@ -288,16 +290,18 @@ def _pick_auto_precision(prequant, device, free_gb, dense_gb, capability, has_fp 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.""" + headroom for activations and the latent cache, not load-time memory. int8 also needs + torchao at runtime (``_int8_quantize_base`` has no fallback, unlike fp8), so auto only + picks it when torchao is importable and drops to nf4 otherwise. ``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" if free_gb > dense_gb * 1.5: return "bf16" if free_gb > dense_gb * 1.15: - return "int8" + return "int8" if has_torchao else "nf4" return "nf4" @@ -324,6 +328,11 @@ def _resolve_base_precision(cfg, spec, device) -> str: free_gb = None capability = None has_fp8 = False + # int8 quantization has no runtime fallback, so gate the auto pick on torchao being + # importable (find_spec avoids the cost/side-effects of an actual import). + import importlib.util + + has_torchao = importlib.util.find_spec("torchao") is not None if device == "cuda": try: import torch @@ -333,7 +342,9 @@ def _resolve_base_precision(cfg, spec, device) -> str: has_fp8 = hasattr(torch, "float8_e4m3fn") except Exception: # noqa: BLE001 -- probe failure -> the safe mode pass - return _pick_auto_precision(prequant, device, free_gb, spec.dense_bf16_gb, capability, has_fp8) + return _pick_auto_precision( + prequant, device, free_gb, spec.dense_bf16_gb, capability, has_fp8, has_torchao + ) # ── FLUX.1-dev ──────────────────────────────────────────────────────────────── @@ -768,8 +779,11 @@ def _load_pixel_tensor_planned(path, resolution, center_crop, u_left, u_top, fli def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_event, check_stop): """Precompute the per-image latent posterior cache: for each planned crop/flip variant, - encode once and store the affine (A, B) pair on CPU (pinned when possible) in the - training dtype. Returns None if the build was interrupted by a stop request.""" + encode once and store the affine (A, B) pair on CPU (pinned when possible) in fp32. The + stats stay fp32 so the per-step sample happens in fp32 and only the RESULT is cast to + weight_dtype, matching the in-loop path (encode fp32 -> sample/normalise fp32 -> + .to(weight_dtype)); fp32 doubles the cache RAM over bf16 but the cache is tiny (a handful + of latents per image). Returns None if the build was interrupted by a stop request.""" plan = _plan_cache_variants( len(image_paths), cfg.cache_variants, cfg.center_crop, cfg.random_flip, cfg.seed @@ -778,7 +792,9 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev def _hold(t): if t is None: return None - t = t.to(weight_dtype).cpu() + import torch + + t = t.to(torch.float32).cpu() if device == "cuda": try: t = t.pin_memory() @@ -808,10 +824,12 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev return cache -def _sample_cached_latents(cache, idxs, variant_rng, device): +def _sample_cached_latents(cache, idxs, variant_rng, device, weight_dtype): """Draw one latent per index from the cache: pick a variant, then sample the posterior (A + B * randn) when the family is stochastic. Fresh noise per step, exactly like an - in-loop ``latent_dist.sample()``.""" + in-loop ``latent_dist.sample()``. The cached stats are fp32, so the sample is drawn in + fp32 and only the RESULT is cast to weight_dtype (matching the in-loop path's + ``encode_latents(...).to(weight_dtype)``).""" import torch parts_a, parts_b = [], [] @@ -822,9 +840,9 @@ def _sample_cached_latents(cache, idxs, variant_rng, device): parts_b.append(b) lat_a = torch.cat(parts_a).to(device, non_blocking = True) if parts_b[0] is None: - return lat_a + return lat_a.to(dtype = weight_dtype) lat_b = torch.cat(parts_b).to(device, non_blocking = True) - return lat_a + lat_b * torch.randn_like(lat_a) + return (lat_a + lat_b * torch.randn_like(lat_a)).to(dtype = weight_dtype) def _should_compile( @@ -1130,7 +1148,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto for _ in range(cfg.gradient_accumulation_steps): idxs = [rng.randrange(n_images) for _ in range(batch_size)] if latent_cache is not None: - latents = _sample_cached_latents(latent_cache, idxs, variant_rng, device) + latents = _sample_cached_latents( + latent_cache, idxs, variant_rng, device, weight_dtype + ) else: px = torch.stack( [ diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 09d1c5618b..b89ad350a7 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -179,11 +179,14 @@ def _build_sdxl_latent_cache( vae, vae_scale, image_paths, cfg, device, weight_dtype, on_event, check_stop ): """Precompute the per-image latent posterior cache: for each planned crop/flip variant, - encode once and store ``(A, B, time_ids)`` on CPU in the training dtype. ``A`` and ``B`` - are the affine posterior parameters (mean/std with the VAE scale folded in) so a per-step - sample is ``A + B * randn`` -- distribution-identical to an in-loop ``latent_dist.sample()`` - -- and ``time_ids`` is the SDXL micro-conditioning for the crop. Returns None if the build - was interrupted by a stop request. ``vae_scale`` is read before the VAE is freed.""" + encode once and store ``(A, B, time_ids)`` on CPU in fp32. ``A`` and ``B`` are the affine + posterior parameters (mean/std with the VAE scale folded in) so a per-step sample is + ``A + B * randn`` -- distribution-identical to an in-loop ``latent_dist.sample()`` -- and + ``time_ids`` is the SDXL micro-conditioning for the crop. The stats stay fp32 so the + per-step sample happens in fp32 and only the RESULT is cast to weight_dtype, matching the + in-loop path (encode fp32 -> sample fp32 -> scale -> .to(weight_dtype)); fp32 doubles the + cache RAM over bf16 but the cache is tiny (a handful of latents per image). Returns None if + the build was interrupted by a stop request. ``vae_scale`` is read before the VAE is freed.""" import torch plan = _plan_cache_variants( @@ -191,7 +194,7 @@ def _build_sdxl_latent_cache( ) def _hold(t): - t = t.to(weight_dtype).cpu() + t = t.to(torch.float32).cpu() if device == "cuda": try: t = t.pin_memory() @@ -224,8 +227,10 @@ def _build_sdxl_latent_cache( def _sample_sdxl_cached_latents(cache, idxs, variant_rng, device, weight_dtype): """Draw one latent + its time_ids per index from the cache: pick a variant, then sample the posterior (A + B * randn) with fresh noise per step, exactly like an in-loop - ``latent_dist.sample() * vae_scale``. Returns ``(latents, batch_time_ids)`` already on - ``device`` in the training dtype (scale + dtype are folded into the cache).""" + ``latent_dist.sample() * vae_scale``. The cached stats are fp32, so the sample is drawn in + fp32 and only the RESULT is cast to weight_dtype (matching the in-loop path). Returns + ``(latents, batch_time_ids)`` already on ``device`` in the training dtype (scale is folded + into the cache).""" import torch parts_a, parts_b, tid_rows = [], [], [] @@ -239,7 +244,7 @@ def _sample_sdxl_cached_latents(cache, idxs, variant_rng, device, weight_dtype): tid_rows.append(time_ids) lat_a = torch.cat(parts_a).to(device, non_blocking = True) lat_b = torch.cat(parts_b).to(device, non_blocking = True) - latents = lat_a + lat_b * torch.randn_like(lat_a) + latents = (lat_a + lat_b * torch.randn_like(lat_a)).to(dtype = weight_dtype) batch_time_ids = torch.tensor(tid_rows, device = device, dtype = weight_dtype) return latents, batch_time_ids @@ -449,7 +454,8 @@ def run_diffusion_lora_training( for _ in range(cfg.gradient_accumulation_steps): idx, img_paths, captions = _next_batch() if latent_cache is not None: - # Scale + dtype are already folded into the cache, so do not re-apply. + # Scale is folded into the cache; the sampler draws in fp32 and casts the + # result to weight_dtype (matching the in-loop path below). latents, batch_time_ids = _sample_sdxl_cached_latents( latent_cache, idx, variant_rng, device, weight_dtype ) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 05faaa5f3b..e0936a1dd4 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -326,7 +326,11 @@ class DiffusionLoraConfig: base_precision = str(self.base_precision or "nf4").strip().lower() if base_precision not in ("nf4", "bf16", "int8", "fp8", "auto"): raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / auto") - if base_precision in ("bf16", "int8", "fp8"): + # base_precision is a DiT-only lever (nf4/bf16/int8/fp8/auto for the transformer + # load); SDXL uses its own mixed_precision path and ignores base_precision entirely, + # so the dense-mode gates (prequant base / non-bf16 compute) apply only to the DiT + # families. The mode-name validity check above still runs for every family. + if resolved_family != "sdxl" and base_precision in ("bf16", "int8", "fp8"): if repo_is_prequantized(self.base_model): raise ValueError( f"base_precision={base_precision!r} needs a dense base repo, but " diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index fb1527b8bf..b33b061f1d 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -18,6 +18,7 @@ runs a scripted target on a thread. from __future__ import annotations import json +import math import multiprocessing as mp import re import threading @@ -34,6 +35,21 @@ _CTX = mp.get_context("spawn") _TERMINAL = ("complete", "error") +def _finite_or_none(value: Any) -> Optional[float]: + """Coerce a numeric progress field to a finite float, or None. A divergent run (or a + grad clip that returns inf) can push loss / grad_norm to NaN or +/-Infinity, and those + are invalid in strict JSON -- FastAPI's encoder would emit the JS-only NaN/Infinity + tokens that break a strict client parse. Nulling them here (the single service ingestion + point both trainers feed) keeps every status snapshot and persisted record JSON-safe.""" + if value is None: + return None + try: + f = float(value) + except (TypeError, ValueError): + return None + return f if math.isfinite(f) else None + + def _run_diffusion_child(*, event_queue: Any, stop_queue: Any, config: dict) -> None: # Imported lazily so this module (and the route layer) stays torch-free at import. from .diffusion_lora_trainer import run_diffusion_training_process @@ -157,21 +173,14 @@ def _append_metric( return if istep <= 0 or loss is None: return - try: - floss = float(loss) - except (TypeError, ValueError): + floss = _finite_or_none(loss) + if floss is None: # non-numeric or non-finite (NaN/Inf): skip, keep the curve JSON-safe return - if floss != floss: # NaN guard - return - - def _opt_float(v: Any) -> Optional[float]: - try: - return float(v) if v is not None else None - except (TypeError, ValueError): - return None - - flr = _opt_float(lr) - fgn = _opt_float(grad_norm) + # lr / grad_norm may be None (sparse series) or non-finite; non-finite values are + # nulled, not dropped, so a bad point never taints the (loss-driven) history while the + # arrays stay index-aligned with steps. + flr = _finite_or_none(lr) + fgn = _finite_or_none(grad_norm) steps = state["metric_steps"] losses = state["metric_loss"] lrs = state["metric_lr"] @@ -431,14 +440,25 @@ class DiffusionTrainingService: # training state, surface the text. s["message"] = str(ev.get("message", "warning")) elif etype == "progress": + # Null any non-finite float (NaN/Inf from a divergent step or an inf grad + # norm) so the JSON status stays strict-parseable; a missing key keeps the + # last value, a present-but-non-finite one becomes None. + loss = _finite_or_none(ev["loss"]) if "loss" in ev else s["loss"] + avg_loss = _finite_or_none(ev["avg_loss"]) if "avg_loss" in ev else s["avg_loss"] + learning_rate = ( + _finite_or_none(ev["learning_rate"]) + if "learning_rate" in ev + else s["learning_rate"] + ) + grad_norm = _finite_or_none(ev["grad_norm"]) if "grad_norm" in ev else s["grad_norm"] s.update( status = "running", step = ev.get("step", s["step"]), total_steps = ev.get("total_steps", s["total_steps"]), - loss = ev.get("loss", s["loss"]), - avg_loss = ev.get("avg_loss", s["avg_loss"]), - learning_rate = ev.get("learning_rate", s["learning_rate"]), - grad_norm = ev.get("grad_norm", s["grad_norm"]), + loss = loss, + avg_loss = avg_loss, + learning_rate = learning_rate, + grad_norm = grad_norm, message = "Training...", ) # Fold optional perf fields (emitted by the trainers) so the UI can show diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index af6789b199..1f8ca9318f 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -32,6 +32,9 @@ from models.training import DiffusionTrainingStartRequest # family from their names alone, so normalized() runs without a network call. _FLUX_DENSE = "black-forest-labs/FLUX.1-dev" _Z_PREQUANT = "unsloth/Z-Image-Turbo-unsloth-bnb-4bit" +# An SDXL base whose name LOOKS prequant (bnb-4bit): SDXL ignores base_precision, so the +# dense-mode gates must not fire for it even with a dense mode + fp16 compute. +_SDXL_PREQUANT_NAME = "some/sdxl-model-bnb-4bit" def _cfg(base_model = _FLUX_DENSE, **kw) -> DiffusionLoraConfig: @@ -66,6 +69,32 @@ def test_base_precision_validation(): assert _cfg(base_model = _Z_PREQUANT, base_precision = "auto").normalized().base_precision == "auto" +def test_base_precision_gates_skip_sdxl(): + # SDXL ignores base_precision, so the dense-mode gates (prequant base / non-bf16 compute) + # must not fire for it: a prequant-looking SDXL name with base_precision="bf16" does not + # raise, and the mode is still stored lowered. + norm = _cfg(base_model = _SDXL_PREQUANT_NAME, base_precision = "bf16").normalized() + assert norm.resolved_family == "sdxl" + assert norm.base_precision == "bf16" + + # The non-bf16-compute gate is also skipped for SDXL (fp16 is a valid SDXL mixed + # precision), even with a dense base_precision requested. + norm2 = _cfg( + base_model = "stabilityai/stable-diffusion-xl-base-1.0", + base_precision = "int8", + mixed_precision = "fp16", + ).normalized() + assert norm2.resolved_family == "sdxl" + + # The mode-name validity check still runs for SDXL: an unknown mode is rejected. + with pytest.raises(ValueError, match = "base_precision"): + _cfg(base_model = _SDXL_PREQUANT_NAME, base_precision = "banana").normalized() + + # The gates STILL fire for a DiT family: a prequant DiT base with a dense mode raises. + with pytest.raises(ValueError, match = "dense base repo"): + _cfg(base_model = _Z_PREQUANT, base_precision = "bf16").normalized() + + # ── repo_is_prequantized heuristic + trainer alias ──────────────────────────── @pytest.mark.parametrize( "repo, expected", @@ -107,6 +136,10 @@ def test_pick_auto_precision_policy_table(): # 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 needs torchao at runtime (no fallback), so the int8 band drops to nf4 when + # torchao is not importable while the bf16 band is unaffected. + assert p(False, "cuda", 30, 23.8, (10, 0), True, False) == "nf4" + assert p(False, "cuda", 140, 23.8, (10, 0), True, False) == "bf16" # 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. @@ -140,6 +173,47 @@ def test_resolve_auto_requires_bf16_compute(): assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4" +def test_resolve_auto_int8_band_gates_on_torchao(monkeypatch): + # The int8 auto band needs torchao at runtime; when torchao is not importable + # _resolve_base_precision must fall to nf4 instead of picking an int8 that would crash + # in _int8_quantize_base. Drive the probe into the int8 band and toggle torchao. + import importlib.util as _ilu + import torch + + spec = dit._SPECS["flux.1"] # dense_bf16_gb = 23.8 + cfg = _cfg(base_precision = "auto", mixed_precision = "bf16") + real_find_spec = _ilu.find_spec + + def _no_torchao(name, *args, **kwargs): + if name == "torchao": + return None # simulate torchao not installed + return real_find_spec(name, *args, **kwargs) + + def _has_torchao(name, *args, **kwargs): + if name == "torchao": + return object() # simulate torchao installed + return real_find_spec(name, *args, **kwargs) + + class _FakeCuda: + # Free VRAM in the int8 band (30 > 23.8 * 1.15) but below the bf16 band. + @staticmethod + def mem_get_info(): + return (int(30 * 1e9), int(80 * 1e9)) + + @staticmethod + def get_device_capability(): + return (10, 0) + + monkeypatch.setattr(torch, "cuda", _FakeCuda) + + monkeypatch.setattr(_ilu, "find_spec", _no_torchao) + assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4" + + # With torchao importable the same band picks int8. + monkeypatch.setattr(_ilu, "find_spec", _has_torchao) + assert dit._resolve_base_precision(cfg, spec, "cuda") == "int8" + + # ── _fp8_module_filter ──────────────────────────────────────────────────────── def test_fp8_module_filter(): lin = nn.Linear(64, 64) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 6320b7fc42..e4bd456f09 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -201,6 +201,48 @@ def test_apply_event_transitions(): assert svc.status()["status"] == "error" and svc.status()["message"] == "boom" +def test_progress_nulls_non_finite_floats_for_strict_json(): + # A divergent step (or an inf grad norm) can push loss / avg_loss / learning_rate to + # NaN or Infinity, which strict JSON forbids. The service must null those so the status + # snapshot and the metric history stay strict-JSON serializable. + import json + import math + + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) + svc._apply_event( + { + "type": "progress", + "step": 3, + "total_steps": 10, + "loss": float("nan"), + "avg_loss": float("inf"), + "learning_rate": float("-inf"), + "grad_norm": float("inf"), + } + ) + snap = svc.status() + assert snap["loss"] is None + assert snap["avg_loss"] is None + assert snap["learning_rate"] is None + # The reviewer's exact case: an inf pre-clip grad norm must not reach the status JSON. + assert snap["grad_norm"] is None + # The non-finite point is skipped in the history, so the loss series stays clean. + assert snap["metric_loss"] == [] + assert snap["metric_steps"] == [] + # strict JSON (allow_nan=False) round-trips without a ValueError from NaN/Infinity. + json.dumps(snap, allow_nan = False) + + # A finite point after the bad one is recorded and preserved verbatim. + svc._apply_event( + {"type": "progress", "step": 4, "total_steps": 10, "loss": 0.5, "learning_rate": 1e-4} + ) + snap2 = svc.status() + assert snap2["loss"] == 0.5 + assert snap2["metric_loss"] == [0.5] and snap2["metric_steps"] == [4] + assert math.isfinite(snap2["learning_rate"]) + json.dumps(snap2, allow_nan = False) + + def test_terminal_events_clear_model_load_flag(): # A stop or error during model load emits complete/error WITHOUT a preceding # model_load_completed, so the terminal update must reset in_model_load or the