diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index ce815f8268..1cbdd93c53 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -778,8 +778,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 @@ -788,7 +791,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() @@ -818,10 +823,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 = [], [] @@ -832,9 +839,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( @@ -1136,7 +1143,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 7c7ca79d51..d8480d5ab1 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -177,11 +177,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( @@ -189,7 +192,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() @@ -222,8 +225,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 = [], [], [] @@ -237,7 +242,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 @@ -443,7 +448,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_training_service.py b/studio/backend/core/training/diffusion_training_service.py index fa44cceabe..4fe072105d 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -17,6 +17,7 @@ runs a scripted target on a thread. from __future__ import annotations +import math import multiprocessing as mp import threading import time @@ -31,6 +32,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 @@ -98,17 +114,12 @@ def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None 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 - flr: Optional[float] - try: - flr = float(lr) if lr is not None else None - except (TypeError, ValueError): - flr = None + # lr may be None (sparse LR series) or non-finite; a non-finite lr is nulled, not dropped, + # so a bad lr point never taints the (loss-driven) history. + flr = _finite_or_none(lr) steps = state["metric_steps"] losses = state["metric_loss"] lrs = state["metric_lr"] @@ -308,13 +319,23 @@ 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"] + ) 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"]), + loss = loss, + avg_loss = avg_loss, + learning_rate = learning_rate, message = "Training...", ) # Fold optional perf fields (emitted by the trainers) so the UI can show diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index f3b7d2b58a..6e2f9dea32 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -189,6 +189,45 @@ 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"), + } + ) + snap = svc.status() + assert snap["loss"] is None + assert snap["avg_loss"] is None + assert snap["learning_rate"] 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