diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 17b5aff61e..df19453448 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -337,7 +337,9 @@ def _apply_mxfp8_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 @@ -346,16 +348,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" @@ -382,6 +386,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 @@ -391,7 +400,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 ──────────────────────────────────────────────────────────────── @@ -939,8 +950,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 @@ -949,7 +963,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() @@ -979,10 +995,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 = [], [] @@ -993,9 +1011,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( @@ -1307,7 +1325,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 d22b7f737e..70f1b420cf 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -181,11 +181,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( @@ -193,7 +196,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() @@ -226,8 +229,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 = [], [], [] @@ -241,7 +246,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 @@ -462,7 +467,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 c867fde7b1..fd49c33423 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -335,8 +335,14 @@ class DiffusionLoraConfig: raise ValueError("compile_transformer must be one of off / on / auto") base_precision = str(self.base_precision or "nf4").strip().lower() if base_precision not in ("nf4", "bf16", "int8", "fp8", "mxfp8", "auto"): - raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto") - if base_precision in ("bf16", "int8", "fp8", "mxfp8"): + raise ValueError( + "base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto" + ) + # base_precision is a DiT-only lever (the transformer load precision); 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", "mxfp8"): 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 ff83dfc2b8..075ca19edb 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 diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index d3f089a42c..966fd31811 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -92,6 +92,21 @@ const FAMILY_PRESETS: FamilyPreset[] = [ const CUSTOM_BASE = "__custom__"; const UPLOAD_DATASET = "__upload__"; +// The dense DiT base precisions: they load a dense (bf16) base and quantise/cast it, so the +// backend rejects them for an already-quantised bnb-4bit repo. "nf4"/"auto" stay valid. +const DENSE_PRECISIONS = new Set(["bf16", "int8", "fp8", "mxfp8"]); +// Mirror the backend's repo_is_prequantized heuristic: a repo whose name marks a +// bitsandbytes 4-bit build already ships a quantised transformer and cannot serve the dense +// base precisions. Kept in sync with diffusion_train_common.repo_is_prequantized. +function repoIsPrequantized(baseModel: string): boolean { + const name = baseModel.toLowerCase(); + return ( + name.includes("bnb-4bit") || + name.includes("-4bit") || + name.includes("int4") || + name.includes("nf4") + ); +} // Dataset-select option value prefix for a not-yet-imported example; picking it imports. const EXAMPLE_PREFIX = "example:"; const DATASET_FILE_ACCEPT = ".png,.jpg,.jpeg,.webp,.bmp,.txt,.caption,.jsonl"; @@ -377,6 +392,16 @@ export function DiffusionTrainPanel({ } }, [family, loadedBaseRepo, reportedFamily?.recommended_precision]); + // mixed_precision is an SDXL-only lever (its UI control is hidden for DiT families). A + // dense DiT base precision (bf16/int8/fp8) requires bf16 compute, and every DiT family + // trains in bf16, so reset precision to bf16 when the family changes to a DiT. Without + // this, an fp16/no value left over from SDXL rides along in the DiT start payload and the + // backend rejects it (dense modes need mixed_precision=bf16). Kept in its own effect so it + // does not re-trigger the base/settings reseed above. + useEffect(() => { + if (isDiT) setPrecision("bf16"); + }, [isDiT]); + // The base actually used everywhere (request, deploy, select value). baseChoice can // briefly hold another family's repo between a family switch and the reseed effect // (or if that effect is skipped); a raw