From 8ad8a587429acb8f8dda5911d32a484e1d5a2178 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 04:31:04 +0000 Subject: [PATCH 1/6] Add grad norm chart, clearer completion state, Windows caption keys, GGUF compute copy - Trainers emit the pre-clip gradient norm; the service keeps a bounded grad_norm history and the Train tab renders a Grad Norm chart next to Loss and LR - Completed runs show 'Training complete' with a celebratory marker in the success color instead of a plain status word - metadata.jsonl caption keys now match on Windows (as_posix relative paths) in both the trainer discovery and the dataset image records - RMSNorm eager patch skips installation on torch builds without F.rms_norm instead of failing at forward time - GGUF compute description no longer says the GGUF is dequantised: the INT8/FP8/FP4 modes load the base model's bf16 transformer and quantise that directly; label no longer wraps in the Advanced panel --- .../core/inference/diffusion_eager_patches.py | 5 +++ .../core/training/diffusion_dit_trainer.py | 6 ++- .../core/training/diffusion_lora_trainer.py | 5 ++- .../core/training/diffusion_train_common.py | 5 ++- .../training/diffusion_training_service.py | 45 ++++++++++++++----- studio/backend/models/training.py | 9 +++- studio/backend/routes/training.py | 9 ++++ studio/frontend/src/features/images/api.ts | 1 + .../src/features/images/images-page.tsx | 4 +- .../images/train/diffusion-charts.tsx | 36 +++++++++++++-- .../images/train/diffusion-train-panel.tsx | 25 ++++++++++- 11 files changed, 126 insertions(+), 24 deletions(-) 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/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index d4ddf9d6d4..4f112b0a27 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -640,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 @@ -662,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 deef4d03ce..db2eca1011 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -338,8 +338,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() @@ -364,6 +366,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 fe077020ac..ef63735185 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -338,9 +338,10 @@ 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..f53fd4623e 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,19 @@ 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 +108,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 +304,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 +313,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 2479571357..965a854042 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -725,12 +725,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): @@ -745,6 +747,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 e61042f1f6..54324dd544 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) @@ -1503,7 +1504,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/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 e712f24c21..1b1d2c366d 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}} @@ -1872,7 +1872,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { {!status?.loaded || status.model_kind === "gguf" ? ( setTransformerQuant(v as typeof transformerQuant)} 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 45df1d50a6..0f67a3aebe 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -359,6 +359,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 ?? []); @@ -784,7 +791,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` : ""} @@ -821,7 +838,11 @@ export function DiffusionTrainPanel({ )}
- + {completed && (
From 1d3aa53d1f48069cf6224826db0bee9bb465c06e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 05:01:58 +0000 Subject: [PATCH 2/6] Validate the training config before importing diffusers and pin the arbiter test's device The fp16-on-bf16-family refusal in run_dit_lora_training now fires before the heavy imports, so a host without diffusers gets the real validation error instead of ModuleNotFoundError. test_in_progress_returns_409_after_validation_passes pins the resolved device to cuda because the load route only takes the GPU arbiter for non-CPU loads, which made the ownership assert host-dependent. --- .../core/training/diffusion_dit_trainer.py | 29 ++++++++++--------- studio/backend/tests/test_diffusion_routes.py | 11 +++++++ 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 4f112b0a27..50bb0f5d3b 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -487,6 +487,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 @@ -494,11 +509,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) @@ -515,15 +525,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). 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"}, From e6d775d0fc80e677292bb29ab856245f5e27fc39 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 05:08:39 +0000 Subject: [PATCH 3/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_train_common.py | 4 +++- .../backend/core/training/diffusion_training_service.py | 8 +++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index ef63735185..62a16eef8c 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -341,7 +341,9 @@ def discover_image_caption_pairs( # 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(img.relative_to(root).as_posix()) + 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 f53fd4623e..5c6a184a86 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -87,7 +87,13 @@ def _idle_state() -> dict[str, Any]: } -def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any, grad_norm: Any = None) -> None: +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``. From d146209f886a207d15a796baf1863d1acfa11ebd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 06:17:45 +0000 Subject: [PATCH 4/6] Rename the GGUF compute control to Dtype and simplify the empty-state copy The always-visible description under the select is gone (the hint tooltip keeps the full detail) and the no-model gallery placeholder now reads 'Select a diffusion model to load'. --- studio/frontend/src/features/images/images-page.tsx | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 1b1d2c366d..418f0e4104 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1871,8 +1871,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)} @@ -1888,7 +1887,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { ) : (
- GGUF compute + Dtype GGUF models only
@@ -2596,7 +2595,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"}

)} From 7bf80f6a4ec7573e25772a25fc8f4bc79695ad47 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 08:51:10 +0000 Subject: [PATCH 5/6] Deny fp8/mxfp8/nvfp4 dense quant for the Qwen DiT (black frames, measured) A 28-pair accuracy gate on a B200 (same-seed vs the dense bf16 reference) found per-row fp8 dynamic quant renders EVERY qwen-image frame black (mean luma 0.0000, SSIM 0.016), reproduced identically with on-the-fly quantize_ on the dense transformer, so it is the model's activation range, not a checkpoint artifact. mxfp8 shows real semantic damage at 1024px (CLIP delta mean 0.0146, worst cases 0.064/0.102) and nvfp4 measures LPIPS mean 0.51. int8 dynamic (per-token scales) is excellent on Qwen: LPIPS mean 0.069, SSIM 0.958. The per-scheme smoke probe only proves the GEMM kernel runs, so it cannot catch model-level breakage. Add _FAMILY_SCHEME_DENY consulted by select_transformer_quant_scheme: auto skips denied schemes (Qwen lands on int8) and an explicit denied request returns None, the same GGUF-fallback contract as an unsupported scheme. Family is threaded from the three diffusion.py call sites; existing behavior is unchanged for every other family. 4 new tests; 529 diffusion tests green; CI-sim green. --- studio/backend/core/inference/diffusion.py | 13 +++- .../inference/diffusion_transformer_quant.py | 40 +++++++++++- .../backend/tests/test_diffusion_backend.py | 12 ++-- .../tests/test_diffusion_transformer_quant.py | 62 ++++++++++++++++++- 4 files changed, 112 insertions(+), 15 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 28eafe6829..2d1a547166 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -418,7 +418,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( @@ -1292,7 +1294,9 @@ 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 @@ -1331,7 +1335,10 @@ 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_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 39aa751778..363098bcdb 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,25 @@ 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 +251,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 +418,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 +430,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/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index d36e62f3cf..8a337557d4 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1764,7 +1764,7 @@ 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 +1828,7 @@ 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 +1958,7 @@ 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 +2005,7 @@ 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 +2016,10 @@ 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_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index 60296bdfc5..8fa4c71e7d 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -433,7 +433,7 @@ 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 +458,13 @@ 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 +480,59 @@ 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 == {} From ab56d819356021d809553c372abba2348476b54d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 08:51:43 +0000 Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/diffusion.py | 12 ++++---- .../inference/diffusion_transformer_quant.py | 4 ++- .../backend/tests/test_diffusion_backend.py | 24 +++++++++++---- .../tests/test_diffusion_transformer_quant.py | 29 ++++++++++++------- 4 files changed, 46 insertions(+), 23 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 2d1a547166..fda2a27a4b 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1294,9 +1294,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, family = getattr(fam, "name", None) - ) + 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 @@ -1336,8 +1334,12 @@ class DiffusionBackend: ) pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device) scheme = quantize_transformer( - pipe, target, mode = mode, family = getattr(fam, "name", None), - fast_accum = fast_accum, logger = logger, + 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") diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 363098bcdb..4d31ae6f25 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -226,7 +226,9 @@ def dense_transformer_supported(target: Any) -> bool: def select_transformer_quant_scheme( - target: Any, requested: Optional[str], family: Optional[str] = None + target: Any, + requested: Optional[str], + family: Optional[str] = None, ) -> Optional[str]: """The concrete scheme to apply, or None to fall back to GGUF. diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 8a337557d4..cd25103a24 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, family = None: 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, family = None: "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, family = None: 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, family = None: "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, family = None: 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, family = None: "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_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index 8fa4c71e7d..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, family = None: 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, family = None: 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, family = None: 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") @@ -492,10 +498,7 @@ def test_family_deny_auto_skips_fp8_for_qwen(monkeypatch): _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 - ) + assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image-edit") == TQ_INT8 def test_family_deny_refuses_explicit_fp8_for_qwen(monkeypatch): @@ -525,14 +528,18 @@ def test_quantize_transformer_threads_family(monkeypatch): pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) called = {} tqz = types.ModuleType("torchao.quantization") - def _quantize(module, config, filter_fn = None): + + 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 quantize_transformer(pipe, _target(), mode = "fp8", family = "qwen-image") is None assert called == {}