diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index a986b31bdd..8120f544ad 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -423,7 +423,9 @@ class DiffusionBackend: target = self._resolve_device_target(fam) if not dense_transformer_supported(target): return False - scheme = select_transformer_quant_scheme(target, mode) + scheme = select_transformer_quant_scheme( + target, mode, family = getattr(fam, "name", None) + ) if scheme is None: return False source = resolve_prequant_source( @@ -1329,7 +1331,7 @@ class DiffusionBackend: BEFORE the loader compiles the repeated block, so the order stays quantize -> compile -> placement.""" # 1. Pre-quantized checkpoint, when one is configured for the resolved scheme. - scheme = select_transformer_quant_scheme(target, mode) + scheme = select_transformer_quant_scheme(target, mode, family = getattr(fam, "name", None)) if scheme is None: # Bail BEFORE the (multi-GB) dense download: an explicit unsupported scheme # (e.g. fp8 on Ampere, nvfp4 off Blackwell) would otherwise materialise the @@ -1368,7 +1370,14 @@ class DiffusionBackend: base, subfolder = "transformer", torch_dtype = dtype, token = hf_token ) pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device) - scheme = quantize_transformer(pipe, target, mode = mode, fast_accum = fast_accum, logger = logger) + scheme = quantize_transformer( + pipe, + target, + mode = mode, + family = getattr(fam, "name", None), + fast_accum = fast_accum, + logger = logger, + ) if scheme is None: raise RuntimeError("transformer quant unsupported for this device/scheme") return pipe, scheme diff --git a/studio/backend/core/inference/diffusion_eager_patches.py b/studio/backend/core/inference/diffusion_eager_patches.py index 593f545984..fff77eb139 100644 --- a/studio/backend/core/inference/diffusion_eager_patches.py +++ b/studio/backend/core/inference/diffusion_eager_patches.py @@ -181,6 +181,11 @@ def install_compile_safe_patches() -> int: for cls, new_fn in _specs(): if cls is None: continue + # torch < 2.4 has no F.rms_norm: leave diffusers' original RMSNorm.forward in + # place rather than installing a patch whose fast path would AttributeError. + if cls is _RMSNorm and not hasattr(F, "rms_norm"): + logger.info("eager-patch: skipping RMSNorm (this torch has no F.rms_norm)") + continue # Capture the live original BEFORE patching so the RMSNorm fast path can fall back # to it for the uncommon (NPU / bias / fp32-weight / tuple-dim) cases. if cls is _RMSNorm: diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 39aa751778..4d31ae6f25 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -101,6 +101,30 @@ _AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = ( ((8, 0), (TQ_INT8,)), # Ampere sm_80 / sm_86 ) +# Families whose activation ranges break specific dense-quant schemes at the MODEL +# level. The kernel smoke probe below cannot see this (it only proves the GEMM runs); +# these were measured with the 28-pair prequant accuracy gate on a B200 +# (scripts/prequant_accuracy_gate.py) and reproduced with on-the-fly quantisation: +# qwen-image + fp8 -> every frame black (mean luma 0.0000, SSIM 0.016 vs bf16). The +# same per-row fp8 that matches bf16 on Z-Image / FLUX: Qwen's +# activation outliers exceed even per-row fp8's dynamic range. +# qwen-image + mxfp8 -> real semantic damage at 1024px (CLIP delta mean 0.0146, worst +# cases 0.064 / 0.102 -- 2x the per-case bound). +# qwen-image + nvfp4 -> LPIPS mean 0.51 vs bf16: unusable. +# int8 dynamic (per-token) is excellent on Qwen (LPIPS mean 0.069 / SSIM 0.958), so the +# auto ladder falls through to it. The deny also applies to an EXPLICIT request: a +# scheme that renders black frames has no legitimate use, and returning None gives the +# caller the same fallback contract as an unsupported scheme (GGUF build). +_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = { + "qwen-image": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), + "qwen-image-edit": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), # same DiT + activations +} + + +def _family_denied(family, scheme: str) -> bool: + return scheme in _FAMILY_SCHEME_DENY.get(str(family or "").strip().lower(), ()) + + # Cache of (scheme, device) -> bool so the quantise+matmul smoke test runs once. _SMOKE_CACHE: dict[tuple[str, str], bool] = {} @@ -201,18 +225,27 @@ def dense_transformer_supported(target: Any) -> bool: return False -def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Optional[str]: +def select_transformer_quant_scheme( + target: Any, + requested: Optional[str], + family: Optional[str] = None, +) -> Optional[str]: """The concrete scheme to apply, or None to fall back to GGUF. ``auto`` walks the per-arch ladder and returns the first scheme that passes a real quantise+matmul smoke test, so on a box where the Blackwell fp4 / mx kernels are unavailable it lands on fp8 / int8 with no error. An explicit scheme is honored only - if supported (else None -> GGUF), never silently swapped for a different one.""" + if supported (else None -> GGUF), never silently swapped for a different one. + ``family`` additionally applies the measured model-level deny list + (``_FAMILY_SCHEME_DENY``): schemes that produce black frames or out-of-bar drift on + that family are skipped by ``auto`` and refused when explicit.""" requested = normalize_transformer_quant(requested) if requested is None or not dense_transformer_supported(target): return None device = str(getattr(target, "device", "cuda")) if requested != TQ_AUTO: + if _family_denied(family, requested): + return None return requested if _scheme_supported(requested, device) else None cap = _capability() if cap is None: @@ -220,6 +253,8 @@ def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Op for floor, schemes in _AUTO_LADDER: if cap >= floor: for scheme in _prefer_consumer_scheme(schemes, device): + if _family_denied(family, scheme): + continue if _scheme_supported(scheme, device): return scheme return None @@ -385,6 +420,7 @@ def quantize_transformer( target: Any, *, mode: Optional[str], + family: Optional[str] = None, min_features: int = DEFAULT_MIN_LINEAR_FEATURES, fast_accum: Optional[bool] = None, logger: Any = None, @@ -396,7 +432,7 @@ def quantize_transformer( ``fast_accum`` (fp8 only) overrides the per-GPU-class accumulate choice: None auto-detects (fast on consumer, precise on data-center), True/False force it.""" - scheme = select_transformer_quant_scheme(target, mode) + scheme = select_transformer_quant_scheme(target, mode, family = family) if scheme is None: return None transformer = getattr(pipe, "transformer", None) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 03f886155d..80e40abcaa 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -486,6 +486,21 @@ def run_dit_lora_training( should_stop: Optional[StopCb] = None, ) -> str: """Train a flow-matching DiT LoRA (FLUX.1-dev / Qwen-Image / Z-Image) and export it.""" + cfg = config.normalized() + spec = _SPECS.get(cfg.resolved_family) + if spec is None: + raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}") + + # DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that + # explicitly asks for fp16 on a bf16-only family is refused rather than silently + # upgraded, so the choice is never misrepresented. Validation runs before the heavy + # imports so a host without diffusers still sees the real error. + if cfg.mixed_precision == "fp16" and spec.force_bf16: + raise ValueError( + f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / " + f"embedder internals. Set mixed precision to bf16." + ) + import torch import torch.nn.functional as F from diffusers import FlowMatchEulerDiscreteScheduler @@ -493,11 +508,6 @@ def run_dit_lora_training( from peft import LoraConfig from peft.utils import get_peft_model_state_dict - cfg = config.normalized() - spec = _SPECS.get(cfg.resolved_family) - if spec is None: - raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}") - rng = random.Random(cfg.seed) torch.manual_seed(cfg.seed) @@ -514,15 +524,6 @@ def run_dit_lora_training( save_on_stop = False return True - # DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that - # explicitly asks for fp16 on a bf16-only family is refused rather than silently - # upgraded, so the choice is never misrepresented. - if cfg.mixed_precision == "fp16" and spec.force_bf16: - raise ValueError( - f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / " - f"embedder internals. Set mixed precision to bf16." - ) - device = "cuda" if torch.cuda.is_available() else "cpu" # The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is # unsupported for real runs but keeps import/unit tests architecture-agnostic). @@ -639,8 +640,11 @@ def run_dit_lora_training( (loss / cfg.gradient_accumulation_steps).backward() step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps + grad_norm: Optional[float] = None if cfg.max_grad_norm and cfg.max_grad_norm > 0: - torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm) + # clip_grad_norm_ returns the PRE-clip total norm: the signal the grad-norm + # chart wants (spikes stay visible even when clipping flattens the update). + grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)) optimizer.step() running_loss += step_loss @@ -661,6 +665,7 @@ def run_dit_lora_training( loss = round(step_loss, 5), avg_loss = round(running_loss / done, 5), learning_rate = cfg.learning_rate, + grad_norm = round(grad_norm, 5) if grad_norm is not None else None, samples_per_second = sps, peak_memory_gb = peak_gb or None, ) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index be006061a7..b96dc1fa07 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -340,8 +340,10 @@ def run_diffusion_lora_training( # max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that); # passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning). + grad_norm: Optional[float] = None if cfg.max_grad_norm and cfg.max_grad_norm > 0: - torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm) + # clip_grad_norm_ returns the PRE-clip total norm (the grad-norm chart signal). + grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)) optimizer.step() lr_sched.step() @@ -366,6 +368,7 @@ def run_diffusion_lora_training( loss = round(step_loss, 5), avg_loss = round(running_loss / done, 5), learning_rate = lr_sched.get_last_lr()[0], + grad_norm = round(grad_norm, 5) if grad_norm is not None else None, samples_per_second = samples_per_second, peak_memory_gb = peak_gb or None, ) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index d854cc4173..0040650f0c 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -384,9 +384,12 @@ def discover_image_caption_pairs( if sidecar.is_file(): caption = sidecar.read_text(encoding = "utf-8").strip() break - # 2. metadata row keyed by file name (basename or the name as written). + # 2. metadata row keyed by file name (basename or the relative path; as_posix so a + # Windows backslash path still matches the jsonl's forward-slash keys). if caption is None: - caption = meta_caption.get(img.name) or meta_caption.get(str(img.relative_to(root))) + caption = meta_caption.get(img.name) or meta_caption.get( + img.relative_to(root).as_posix() + ) # 3. dreambooth instance prompt. if caption is None and instance_prompt: caption = instance_prompt diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 2ad30e8347..5c6a184a86 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -67,6 +67,7 @@ def _idle_state() -> dict[str, Any]: "loss": None, "avg_loss": None, "learning_rate": None, + "grad_norm": None, "num_images": None, "in_model_load": False, "output_dir": None, @@ -82,16 +83,25 @@ def _idle_state() -> dict[str, Any]: "metric_steps": [], "metric_loss": [], "metric_lr": [], + "metric_grad_norm": [], } -def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None: - """Append one (step, loss, lr) point to the bounded history arrays on ``state``. +def _append_metric( + state: dict[str, Any], + step: Any, + loss: Any, + lr: Any, + grad_norm: Any = None, +) -> None: + """Append one (step, loss, lr, grad_norm) point to the bounded history arrays on + ``state``. Only records finite, positive-step points (mirrors the LLM trainer, which logs history only for step > 0 with a real loss). When the arrays hit ``_METRIC_CAP`` they are decimated in place (keep every other point) so appends stay bounded without losing the - curve's shape. lr may be None (kept as None so the LR series can be sparse).""" + curve's shape. lr / grad_norm may be None (kept as None so those series can be + sparse while staying index-aligned with ``steps``).""" try: istep = int(step) except (TypeError, ValueError): @@ -104,22 +114,34 @@ def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None 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 + + 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) steps = state["metric_steps"] losses = state["metric_loss"] lrs = state["metric_lr"] + gns = state["metric_grad_norm"] if len(steps) >= _METRIC_CAP: state["metric_steps"] = steps[::2] state["metric_loss"] = losses[::2] state["metric_lr"] = lrs[::2] - steps, losses, lrs = state["metric_steps"], state["metric_loss"], state["metric_lr"] + state["metric_grad_norm"] = gns[::2] + steps, losses, lrs, gns = ( + state["metric_steps"], + state["metric_loss"], + state["metric_lr"], + state["metric_grad_norm"], + ) steps.append(istep) losses.append(floss) lrs.append(flr) + gns.append(fgn) class DiffusionTrainingService: @@ -288,6 +310,7 @@ class DiffusionTrainingService: 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"]), message = "Training...", ) # Fold optional perf fields (emitted by the trainers) so the UI can show @@ -296,8 +319,14 @@ class DiffusionTrainingService: s["samples_per_second"] = ev.get("samples_per_second") if ev.get("peak_memory_gb") is not None: s["peak_memory_gb"] = ev.get("peak_memory_gb") - # Retain a bounded (step, loss, lr) history for the live loss chart. - _append_metric(s, ev.get("step"), ev.get("loss"), ev.get("learning_rate")) + # Retain a bounded (step, loss, lr, grad_norm) history for the live charts. + _append_metric( + s, + ev.get("step"), + ev.get("loss"), + ev.get("learning_rate"), + ev.get("grad_norm"), + ) elif etype == "complete": # Reset in_model_load: a stop during model load emits complete without a # preceding model_load_completed, which would otherwise leave a stale diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index c9606eb5b1..522d5d3521 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -734,12 +734,14 @@ class DiffusionTrainingStartResponse(BaseModel): class DiffusionMetricHistory(BaseModel): - """Paired step-indexed history arrays for the live training charts. ``lr`` entries may - be null so a sparse learning-rate series still aligns with ``steps`` by index.""" + """Paired step-indexed history arrays for the live training charts. ``lr`` and + ``grad_norm`` entries may be null so those sparse series still align with ``steps`` + by index.""" steps: List[int] = Field(default_factory = list) loss: List[float] = Field(default_factory = list) lr: List[Optional[float]] = Field(default_factory = list) + grad_norm: List[Optional[float]] = Field(default_factory = list) class DiffusionTrainingStatusResponse(BaseModel): @@ -754,6 +756,9 @@ class DiffusionTrainingStatusResponse(BaseModel): loss: Optional[float] = None avg_loss: Optional[float] = None learning_rate: Optional[float] = None + # Pre-clip gradient norm from the trainer's progress events (None when clipping is + # disabled), feeding the grad-norm chart. + grad_norm: Optional[float] = None num_images: Optional[int] = None in_model_load: bool = False output_dir: Optional[str] = None diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 8868eb1e2a..a82491c04f 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1267,6 +1267,7 @@ async def diffusion_training_status(current_subject: str = Depends(get_current_s steps = snap.pop("metric_steps", []), loss = snap.pop("metric_loss", []), lr = snap.pop("metric_lr", []), + grad_norm = snap.pop("metric_grad_norm", []), ) return DiffusionTrainingStatusResponse(**snap, metric_history = metric_history) @@ -1515,7 +1516,15 @@ def _image_record( caption = None break if caption is None: + # Basename first, then the relative path as written in the jsonl (as_posix so a + # Windows backslash path still matches forward-slash keys) -- the same lookup + # order discover_image_caption_pairs uses. meta = meta_captions.get(image_path.name) + if meta is None: + try: + meta = meta_captions.get(image_path.relative_to(folder).as_posix()) + except ValueError: + meta = None if meta is not None: caption = meta source = "metadata" diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index cf5f8ac07b..a30d3ffadf 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1764,7 +1764,9 @@ def _stub_dense_quant(monkeypatch, *, scheme = "fp8"): monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) # Resolve the scheme without the real GPU smoke probe, and configure no pre-quant # checkpoint so the dense materialise+quantise branch is the one exercised. - monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: scheme) + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: scheme + ) monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) def _quantize(pipe, target, *, mode, **kw): @@ -1828,7 +1830,9 @@ def test_transformer_quant_prequant_path_engaged(fake_runtime, tmp_path, monkeyp backend = DiffusionBackend() _force_cuda_target(backend, monkeypatch) monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) - monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8") + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8" + ) monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object()) prequant_obj = object() loaded: dict = {"n": 0} @@ -1958,7 +1962,9 @@ def test_transformer_quant_unsupported_scheme_skips_dense_download( backend = DiffusionBackend() _force_cuda_target(backend, monkeypatch) monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) - monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: None) + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: None + ) monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) @classmethod @@ -2005,7 +2011,9 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): _force_cuda_target(backend, monkeypatch) fam = detect_family("unsloth/Z-Image-Turbo-GGUF") monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) - monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8") + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8" + ) monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is True @@ -2016,10 +2024,14 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False # Unsupported scheme bails before the dense path (and so must the prefetch). monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) - monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: None) + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: None + ) assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False # Device without dense support (e.g. non-CUDA) never widens. - monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8") + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8" + ) monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: False) assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index f6a9ea2766..567cc39c26 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -659,6 +659,17 @@ def test_in_progress_returns_409_after_validation_passes(client, monkeypatch): backend = _FakeBackend() backend.begin_load = _busy monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + # Pin the resolved device to cuda: the route only takes the arbiter for non-CPU + # loads, so on a CPU-only host the ownership assert below would never hold. + import types as _types + + import core.inference.diffusion_device as devmod + + monkeypatch.setattr( + devmod, + "resolve_diffusion_device_target", + lambda: _types.SimpleNamespace(device = "cuda"), + ) resp = client.post( "/api/inference/images/load", json = {"model_path": "unsloth/Z-Image-Turbo-GGUF", "gguf_filename": "q.gguf"}, diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index 60296bdfc5..a621bd3874 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -433,7 +433,9 @@ def test_fp8_config_uses_per_row_granularity(): def test_quantize_transformer_applies_and_marks(monkeypatch): - monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_FP8) + monkeypatch.setattr( + tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8 + ) seen: dict = {} def _mk(scheme, fast_accum = None): @@ -458,13 +460,17 @@ def test_quantize_transformer_applies_and_marks(monkeypatch): def test_quantize_transformer_none_when_unsupported(monkeypatch): - monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: None) + monkeypatch.setattr( + tq, "select_transformer_quant_scheme", lambda target, mode, family = None: None + ) pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) assert quantize_transformer(pipe, _target(), mode = "auto") is None def test_quantize_transformer_tolerates_failure(monkeypatch): - monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_INT8) + monkeypatch.setattr( + tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_INT8 + ) monkeypatch.setattr(tq, "_make_quant_config", lambda scheme: "cfg") tqz = types.ModuleType("torchao.quantization") @@ -480,3 +486,60 @@ def test_quantize_transformer_tolerates_failure(monkeypatch): pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) # A quantise failure returns None (caller falls back to GGUF), never raises. assert quantize_transformer(pipe, _target(), mode = "int8") is None + + +# ── family scheme deny (measured model-level breakage) ──────────────────────── + + +def test_family_deny_auto_skips_fp8_for_qwen(monkeypatch): + # B200 with every scheme available: auto must NOT pick fp8 / nvfp4 / mxfp8 for the + # Qwen DiT (per-row fp8 renders black frames on it; see _FAMILY_SCHEME_DENY) and + # falls through the ladder to int8, which measures excellent on Qwen. + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image") == TQ_INT8 + assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image-edit") == TQ_INT8 + + +def test_family_deny_refuses_explicit_fp8_for_qwen(monkeypatch): + # An explicit fp8 request on qwen-image returns None (same contract as an + # unsupported scheme: the caller builds the GGUF pipeline instead). int8 stays + # honored on qwen, and fp8 stays honored on families outside the deny table. + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "fp8", family = "qwen-image") is None + assert select_transformer_quant_scheme(_target(), "int8", family = "qwen-image") == TQ_INT8 + assert select_transformer_quant_scheme(_target(), "fp8", family = "z-image") == TQ_FP8 + + +def test_family_deny_no_family_keeps_ladder(monkeypatch): + # Without a family (or an unknown one) the ladder is unchanged: fp8 first on B200. + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + assert select_transformer_quant_scheme(_target(), "auto", family = "sdxl") == TQ_FP8 + + +def test_quantize_transformer_threads_family(monkeypatch): + # quantize_transformer passes the family down to the selector, so a denied + # (family, scheme) pair never reaches torchao. + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) + called = {} + tqz = types.ModuleType("torchao.quantization") + + def _quantize( + module, + config, + filter_fn = None, + ): + called["scheme"] = True + + tqz.quantize_ = _quantize + tqz.Int8DynamicActivationInt8WeightConfig = lambda: "int8-cfg" + tqz.Float8DynamicActivationFloat8WeightConfig = lambda **kw: "fp8-cfg" + tqz.PerRow = lambda: "per-row" + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + assert quantize_transformer(pipe, _target(), mode = "fp8", family = "qwen-image") is None + assert called == {} diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 911a2ff426..6ff2c72b77 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -291,6 +291,7 @@ export interface DiffusionMetricHistory { steps: number[]; loss: number[]; lr: Array; + grad_norm: Array; } // A snapshot of the current diffusion training job (GET /api/train/diffusion/status). diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 0bf4179a25..b518ad2cb9 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -481,7 +481,7 @@ function AdvancedSelect({ return (
- + {label} {hint && {hint}} @@ -1899,8 +1899,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { to GGUF (or nothing loaded) and otherwise show why it is unavailable. */} {!status?.loaded || status.model_kind === "gguf" ? ( setTransformerQuant(v as typeof transformerQuant)} @@ -1916,7 +1915,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { ) : (
- GGUF compute + Dtype GGUF models only
@@ -2624,7 +2623,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {

{status?.loaded ? "Enter a prompt and hit Generate." - : "Select a model quant to load, then generate."} + : "Select a diffusion model to load"}

)} diff --git a/studio/frontend/src/features/images/train/diffusion-charts.tsx b/studio/frontend/src/features/images/train/diffusion-charts.tsx index c4378156cf..a59dc3f7fb 100644 --- a/studio/frontend/src/features/images/train/diffusion-charts.tsx +++ b/studio/frontend/src/features/images/train/diffusion-charts.tsx @@ -9,6 +9,8 @@ import type { TrainingSeriesPoint } from "@/features/training"; // which are meaningless for diffusion LoRA training and showed as an empty card and an // "Evaluation not configured" placeholder. This is a diffusion-only two-card layout. // eslint-disable-next-line no-restricted-imports +import { GradNormChartCard } from "@/features/studio/sections/charts/grad-norm-chart-card"; +// eslint-disable-next-line no-restricted-imports import { LearningRateChartCard } from "@/features/studio/sections/charts/learning-rate-chart-card"; // eslint-disable-next-line no-restricted-imports import { TrainingLossChartCard } from "@/features/studio/sections/charts/training-loss-chart-card"; @@ -43,14 +45,17 @@ function fullStepDomain(steps: number[]): [number, number] { return [min, max]; } -// A diffusion-only metrics view: just Training Loss and Learning Rate, side by side, with a -// note under the loss card explaining why per-step loss looks noisy. +// A diffusion-only metrics view: Training Loss and Learning Rate side by side, plus Grad +// Norm (the pre-clip total gradient norm; spikes flag instability that raw loss noise +// hides), with a note under the loss card explaining why per-step loss looks noisy. export function DiffusionCharts({ lossHistory, lrHistory, + gradNormHistory = [], }: { lossHistory: TrainingSeriesPoint[]; lrHistory: TrainingSeriesPoint[]; + gradNormHistory?: TrainingSeriesPoint[]; }): ReactElement | null { const lossItems = useMemo(() => toLossItems(lossHistory), [lossHistory]); const smoothed = useMemo( @@ -82,12 +87,24 @@ export function DiffusionCharts({ [lrHistory], ); + const gradNormData = useMemo( + () => + compressSeries( + gradNormHistory + .filter((p) => Number.isFinite(p.value)) + .map((p) => ({ step: p.step, gradNorm: p.value, displayGradNorm: p.value })), + MAX_RENDER_POINTS, + ), + [gradNormHistory], + ); + const steps = useMemo(() => { const set = new Set(); for (const p of lossData) set.add(p.step); for (const p of lrData) set.add(p.step); + for (const p of gradNormData) set.add(p.step); return Array.from(set).sort((a, b) => a - b); - }, [lossData, lrData]); + }, [lossData, lrData, gradNormData]); const stepDomain = useMemo(() => fullStepDomain(steps), [steps]); const xAxisTicks = useMemo( @@ -103,6 +120,10 @@ export function DiffusionCharts({ () => buildYDomain(lrData.map((p) => p.displayLr)), [lrData], ); + const gradNormDomain = useMemo( + () => buildYDomain(gradNormData.map((p) => p.displayGradNorm)), + [gradNormData], + ); const avgRaw = lossItems.length > 0 @@ -138,6 +159,15 @@ export function DiffusionCharts({ xAxisTicks={xAxisTicks} scale="linear" /> + {gradNormData.length > 0 && ( + + )}
); } 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 75b1d00ad1..04fdeded5c 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -366,6 +366,13 @@ export function DiffusionTrainPanel({ .map((step, i) => ({ step, value: h.lr[i] })) .filter((p): p is TrainingSeriesPoint => p.value != null); }, [status?.metric_history]); + const gradNormHistory: TrainingSeriesPoint[] = useMemo(() => { + const h = status?.metric_history; + if (!h) return []; + return h.steps + .map((step, i) => ({ step, value: h.grad_norm?.[i] ?? null })) + .filter((p): p is TrainingSeriesPoint => p.value != null); + }, [status?.metric_history]); const onUpload = useCallback(async () => { const files = Array.from(fileInputRef.current?.files ?? []); @@ -791,7 +798,17 @@ export function DiffusionTrainPanel({ <>
- {status.status} + {/* A finished run should be unmistakable at a glance, so completed swaps + the plain status word for a celebratory line in the success color. */} + + {status.status === "completed" ? "Training complete \u{1F389}" : status.status} + {status.total_steps > 0 ? `${status.step}/${status.total_steps} steps` : ""} @@ -828,7 +845,11 @@ export function DiffusionTrainPanel({ )}
- + {completed && (