From 8ad8a587429acb8f8dda5911d32a484e1d5a2178 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 04:31:04 +0000 Subject: [PATCH] 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 && (