Merge branch 'diffusion-krea2' into diffusion-train-perf2
# Conflicts: # studio/backend/core/training/diffusion_dit_trainer.py # studio/backend/core/training/diffusion_train_common.py # studio/frontend/src/features/images/train/diffusion-train-panel.tsx
This commit is contained in:
Daniel Han
commit
14b2a68025
7 changed files with 274 additions and 46 deletions
|
|
@ -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(
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 <select value> would then DISPLAY the first
|
||||
|
|
@ -387,6 +412,22 @@ export function DiffusionTrainPanel({
|
|||
? baseChoice
|
||||
: family?.base_repos[0] ?? CUSTOM_BASE;
|
||||
|
||||
// The resolved base repo/path the request will carry, and whether it looks prequantized
|
||||
// (bnb-4bit etc.). The dense base precisions are invalid for such a repo, so we gate them.
|
||||
const resolvedBase = (effectiveBase === CUSTOM_BASE ? customBase : effectiveBase).trim();
|
||||
const basePrequantized = isDiT && repoIsPrequantized(resolvedBase);
|
||||
|
||||
// A prequantized base cannot serve the dense precisions; auto-flip a dense selection back
|
||||
// to "auto" (which resolves to nf4 for such a repo) so the run does not fail at the backend
|
||||
// validator. Reuses the precisionDirty ref so a later family change still re-seeds from the
|
||||
// recommendation. The dense options are also disabled in the select below.
|
||||
useEffect(() => {
|
||||
if (basePrequantized && DENSE_PRECISIONS.has(basePrecision)) {
|
||||
precisionDirty.current = false;
|
||||
setBasePrecision("auto");
|
||||
}
|
||||
}, [basePrequantized, basePrecision]);
|
||||
|
||||
const poll = useCallback(async () => {
|
||||
try {
|
||||
setStatus(await getDiffusionTrainingStatus());
|
||||
|
|
@ -436,7 +477,11 @@ export function DiffusionTrainPanel({
|
|||
);
|
||||
|
||||
// Notify the parent exactly once per run that produced an adapter (full completion or
|
||||
// stop-and-save) so it rescans the LoRA picker.
|
||||
// stop-and-save) so it rescans the LoRA picker. The flag is re-armed both here (when a
|
||||
// new run is observed as "running") and in onStart (the moment a start is requested), so
|
||||
// a second run still notifies even if the poll never catches the intermediate "running"
|
||||
// state; onStart also guards the double-fire when the poll re-observes the same terminal
|
||||
// status before the new run has begun.
|
||||
const notifiedComplete = useRef(false);
|
||||
useEffect(() => {
|
||||
const producedAdapter =
|
||||
|
|
@ -580,6 +625,10 @@ export function DiffusionTrainPanel({
|
|||
// read-time clamp (running && stopRequestedLocal) re-arms the moment the new run goes
|
||||
// active, rendering a permanently disabled "Stopping..." button.
|
||||
setStopRequestedLocal(false);
|
||||
// Re-arm the completion notification for this run. Resetting here (not only when the
|
||||
// poll later sees "running") means a second run still notifies even if its "running"
|
||||
// phase is never observed, and prevents the prior run's terminal status re-firing it.
|
||||
notifiedComplete.current = false;
|
||||
// A history view must not shadow the new live run.
|
||||
setViewRun(null);
|
||||
try {
|
||||
|
|
@ -824,7 +873,11 @@ export function DiffusionTrainPanel({
|
|||
aria-label="Base precision"
|
||||
>
|
||||
{precisionModes.map((m) => (
|
||||
<option key={m} value={m}>
|
||||
<option
|
||||
key={m}
|
||||
value={m}
|
||||
disabled={basePrequantized && DENSE_PRECISIONS.has(m)}
|
||||
>
|
||||
{precisionLabel(m)}
|
||||
</option>
|
||||
))}
|
||||
|
|
@ -833,6 +886,13 @@ export function DiffusionTrainPanel({
|
|||
How the frozen base weights are quantised. nf4 (4-bit) uses the least VRAM;
|
||||
bf16 is fastest but needs the most. Auto picks this family's recommended
|
||||
mode.
|
||||
{basePrequantized && (
|
||||
<>
|
||||
{" "}
|
||||
This base is already 4-bit quantised, so only nf4/auto apply; pick a dense
|
||||
(bf16) base repo for the other modes.
|
||||
</>
|
||||
)}
|
||||
</p>
|
||||
</div>
|
||||
) : (
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue