Merge branch 'diffusion-train-perf' into diffusion-train-precision

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(
                     [

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
                     )

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

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