Show only loss and learning-rate charts for diffusion training

The Train tab reused the LLM charts section, which also rendered an empty
Grad Norm card and an Eval Loss card showing an Evaluation not configured
placeholder with a red smear. Neither applies to diffusion LoRA training.
Add a diffusion-only two-card view that reuses the loss and learning-rate
cards directly with fixed presentation defaults, and note under the loss
chart that per-step loss is noisy by design so users read the smoothed
line for the trend rather than the raw jitter.
This commit is contained in:
Daniel Han 2026-07-02 20:55:57 +00:00
commit 15d1de6e05
2 changed files with 154 additions and 14 deletions

View file

@ -0,0 +1,143 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { type ReactElement, useMemo } from "react";
import type { TrainingSeriesPoint } from "@/features/training";
// The loss + LR cards are pure presentational (props only), so reuse them directly. We do
// NOT reuse ChartsSection/ChartsContent: those also render Grad Norm and an Eval Loss card,
// 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 { 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";
// eslint-disable-next-line no-restricted-imports
import {
MAX_RENDER_POINTS,
buildStepTicks,
buildYDomain,
compressSeries,
ema,
} from "@/features/studio/sections/charts/utils";
// Fixed presentation defaults (the LLM tab exposes these via a settings sheet; here we pick
// sensible constants): EMA smoothing on, linear scale, raw + smoothed + average lines shown,
// no outlier trimming (diffusion loss is naturally noisy, not spiky-with-outliers).
const SMOOTHING = 0.8;
function toLossItems(series: TrainingSeriesPoint[]): { step: number; loss: number }[] {
return series
.filter((p) => Number.isFinite(p.value))
.map((p) => ({ step: p.step, loss: p.value }));
}
// The x-domain that spans all points (the LLM tab supports a scrollable window; a training
// run here is short enough to always show the whole thing).
function fullStepDomain(steps: number[]): [number, number] {
if (steps.length === 0) return [0, 1];
const min = steps[0];
const max = steps[steps.length - 1];
if (min === max) return [min, min + 4];
if (max - min < 6) return [Math.max(0, max - 6), max];
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.
export function DiffusionCharts({
lossHistory,
lrHistory,
}: {
lossHistory: TrainingSeriesPoint[];
lrHistory: TrainingSeriesPoint[];
}): ReactElement | null {
const lossItems = useMemo(() => toLossItems(lossHistory), [lossHistory]);
const smoothed = useMemo(
() => (lossItems.length > 0 ? ema(lossItems, SMOOTHING) : []),
[lossItems],
);
const reducedLoss = useMemo(
() => compressSeries(smoothed, MAX_RENDER_POINTS),
[smoothed],
);
const lossData = useMemo(
() =>
reducedLoss.map((p) => ({
...p,
displayLoss: p.loss,
displaySmoothed: p.smoothed,
})),
[reducedLoss],
);
const lrData = useMemo(
() =>
compressSeries(
lrHistory
.filter((p) => Number.isFinite(p.value))
.map((p) => ({ step: p.step, lr: p.value, displayLr: p.value })),
MAX_RENDER_POINTS,
),
[lrHistory],
);
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);
return Array.from(set).sort((a, b) => a - b);
}, [lossData, lrData]);
const stepDomain = useMemo(() => fullStepDomain(steps), [steps]);
const xAxisTicks = useMemo(
() => buildStepTicks(stepDomain[0], stepDomain[1]),
[stepDomain],
);
const lossDomain = useMemo(
() => buildYDomain(lossData.flatMap((p) => [p.displayLoss, p.displaySmoothed])),
[lossData],
);
const lrDomain = useMemo(
() => buildYDomain(lrData.map((p) => p.displayLr)),
[lrData],
);
const avgRaw =
lossItems.length > 0
? +(lossItems.reduce((s, p) => s + p.loss, 0) / lossItems.length).toFixed(4)
: 0;
if (lossItems.length === 0 && lrData.length === 0) return null;
return (
<div className="grid grid-cols-1 gap-4 lg:grid-cols-2">
<div className="flex flex-col gap-1">
<TrainingLossChartCard
data={lossData}
domain={lossDomain}
visibleStepDomain={stepDomain}
xAxisTicks={xAxisTicks}
avgRaw={avgRaw}
avgDisplay={avgRaw}
showRaw={true}
showSmoothed={true}
showAvgLine={true}
scale="linear"
/>
<p className="px-1 text-[11px] leading-snug text-muted-foreground">
Per-step loss is noisy by design: every step samples a random noise level. Watch
the smoothed line for the trend, not the raw jitter.
</p>
</div>
<LearningRateChartCard
data={lrData}
domain={lrDomain}
visibleStepDomain={stepDomain}
xAxisTicks={xAxisTicks}
scale="linear"
/>
</div>
);
}

View file

@ -6,10 +6,6 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
// Deep import (not the studio feature index) so the Images bundle does not statically pull
// in the heavy StudioPage; charts-section itself lazy-loads its recharts content.
// eslint-disable-next-line no-restricted-imports
import { ChartsSection } from "@/features/studio/sections/charts-section";
import type { TrainingSeriesPoint } from "@/features/training";
// eslint-disable-next-line no-restricted-imports -- matches images-page.tsx's token access
import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store";
@ -27,6 +23,8 @@ import {
uploadDiffusionDataset,
} from "../api";
import { DatasetLabelingGrid, LabelingGridToggle } from "./dataset-labeling-grid";
import { DatasetShowcase } from "./dataset-showcase";
import { DiffusionCharts } from "./diffusion-charts";
import { ExampleDatasetCards } from "./example-dataset-cards";
// The families the Train tab can train, in the popularity order the user asked for. This is
@ -549,6 +547,14 @@ export function DiffusionTrainPanel({
) : (
selectedDataset && (
<>
{selectedDataset.image_count > 0 && !gridOpen && (
<DatasetShowcase
dataset={dataset}
imageCount={selectedDataset.image_count}
refreshKey={gridRefresh}
onBrowse={() => setGridOpen(true)}
/>
)}
<LabelingGridToggle
count={selectedDataset.image_count}
open={gridOpen}
@ -707,16 +713,7 @@ export function DiffusionTrainPanel({
)}
</div>
<ChartsSection
currentStep={status.step}
totalSteps={status.total_steps}
isTraining={running}
evalEnabled={false}
lossHistory={lossHistory}
lrHistory={lrHistory}
gradNormHistory={[]}
evalLossHistory={[]}
/>
<DiffusionCharts lossHistory={lossHistory} lrHistory={lrHistory} />
{completed && (
<div className="bg-card corner-squircle flex flex-col gap-2 rounded-3xl p-5 ring-1 ring-foreground/10">