ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add grad norm chart, clearer completion state, Windows caption keys, GGUF compute copy

Browse source 
- 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
This commit is contained in:
Daniel Han 2026-07-04 04:31:04 +00:00
commit 8ad8a58742
11 changed files with 126 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

5
studio/backend/core/inference/diffusion_eager_patches.py
View file

@ -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:

6
studio/backend/core/training/diffusion_dit_trainer.py
View file

@ -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,
)

5
studio/backend/core/training/diffusion_lora_trainer.py
View file

@ -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,
)

5
studio/backend/core/training/diffusion_train_common.py
View file

@ -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

45
studio/backend/core/training/diffusion_training_service.py
View file

@ -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

9
studio/backend/models/training.py
View file

@ -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

9
studio/backend/routes/training.py
View file

@ -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"

1
studio/frontend/src/features/images/api.ts
View file

@ -291,6 +291,7 @@ export interface DiffusionMetricHistory {
steps: number[];
loss: number[];
lr: Array<number | null>;
grad_norm: Array<number | null>;
}
// A snapshot of the current diffusion training job (GET /api/train/diffusion/status).

4
studio/frontend/src/features/images/images-page.tsx
View file

@ -481,7 +481,7 @@ function AdvancedSelect({
return (
<div className="flex flex-col gap-1">
<div className="flex items-center justify-between gap-2">
<span className="flex items-center gap-1 text-xs font-medium text-muted-foreground">
<span className="flex shrink-0 items-center gap-1 whitespace-nowrap text-xs font-medium text-muted-foreground">
{label}
{hint && <InfoHint>{hint}</InfoHint>}
</span>
@ -1872,7 +1872,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
{!status?.loaded || status.model_kind === "gguf" ? (
<AdvancedSelect
label="GGUF compute"
desc="Off runs the GGUF as-is. INT8/FP8/FP4 dequantise the transformer onto low-precision tensor cores for a faster step, at the cost of a larger download and more VRAM."
desc="Off runs the GGUF as-is. INT8/FP8/FP4 instead download the base model's bf16 transformer and quantise it directly onto low-precision tensor cores (the GGUF is not requantised): a faster step, at the cost of a larger download and more VRAM."
hint="Optional speed-up for GGUF models. Off runs the GGUF as-is. FP8/INT8/FP4 instead load the FULL base model and quantise its transformer onto low-precision tensor cores: faster per step, but a larger download and more VRAM, and it falls back to the GGUF if it can't fit. Needs CUDA."
value={transformerQuant}
onValueChange={(v) => setTransformerQuant(v as typeof transformerQuant)}

36
studio/frontend/src/features/images/train/diffusion-charts.tsx
View file

@ -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<number>();
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 && (
<GradNormChartCard
data={gradNormData}
domain={gradNormDomain}
visibleStepDomain={stepDomain}
xAxisTicks={xAxisTicks}
scale="linear"
/>
)}
</div>
);
}

25
studio/frontend/src/features/images/train/diffusion-train-panel.tsx
View file

@ -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({
<>
<div className="bg-card corner-squircle flex flex-col gap-3 rounded-3xl p-5 ring-1 ring-foreground/10">
<div className="flex items-center justify-between">
<span className="text-sm font-semibold capitalize">{status.status}</span>
{/* A finished run should be unmistakable at a glance, so completed swaps
the plain status word for a celebratory line in the success color. */}
<span
className={
status.status === "completed"
? "text-sm font-semibold text-emerald-600 dark:text-emerald-400"
: "text-sm font-semibold capitalize"
}
>
{status.status === "completed" ? "Training complete \u{1F389}" : status.status}
</span>
<span className="text-xs text-muted-foreground">
{status.total_steps > 0 ? `${status.step}/${status.total_steps} steps` : ""}
</span>
@ -821,7 +838,11 @@ export function DiffusionTrainPanel({
)}
</div>
<DiffusionCharts lossHistory={lossHistory} lrHistory={lrHistory} />
<DiffusionCharts
lossHistory={lossHistory}
lrHistory={lrHistory}
gradNormHistory={gradNormHistory}
/>
{completed && (
<div className="bg-card corner-squircle flex flex-col gap-2 rounded-3xl p-5 ring-1 ring-foreground/10">