studio: add max steps and epochs toggle switch (#4296)

* feat: add Epochs toggle for Max Steps

* refactor: dedupe max-steps/epochs toggle logic and fix input bug

* fix(studio): max-steps input validation and prevSaveSteps seed in epochs mode

---------

Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
This commit is contained in:
Lee Jackson 2026-03-15 21:33:51 +00:00 committed by GitHub
commit 1e3aa4ff92
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 322 additions and 165 deletions

View file

@ -23,7 +23,7 @@ import {
TooltipTrigger,
} from "@/components/ui/tooltip";
import { CONTEXT_LENGTHS } from "@/config/training";
import { useTrainingConfigStore } from "@/features/training";
import { useMaxStepsEpochsToggle, useTrainingConfigStore } from "@/features/training";
import { InformationCircleIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useShallow } from "zustand/react/shallow";
@ -35,6 +35,8 @@ export function HyperparametersStep() {
setMaxSteps,
epochs,
setEpochs,
saveSteps,
setSaveSteps,
contextLength,
setContextLength,
learningRate,
@ -52,6 +54,8 @@ export function HyperparametersStep() {
setMaxSteps: s.setMaxSteps,
epochs: s.epochs,
setEpochs: s.setEpochs,
saveSteps: s.saveSteps,
setSaveSteps: s.setSaveSteps,
contextLength: s.contextLength,
setContextLength: s.setContextLength,
learningRate: s.learningRate,
@ -67,6 +71,15 @@ export function HyperparametersStep() {
const showLoraParams =
trainingMethod === "lora" || trainingMethod === "qlora";
const { useEpochs, toggleUseEpochs } = useMaxStepsEpochsToggle({
maxSteps,
epochs,
saveSteps,
setMaxSteps,
setEpochs,
setSaveSteps,
});
const maxStepsSliderMax = Math.max(500, maxSteps, 30);
const epochsSliderMax = Math.max(10, epochs, 1);
@ -75,52 +88,84 @@ export function HyperparametersStep() {
<FieldSet>
<FieldLegend variant="label">Training</FieldLegend>
<div className="flex flex-col gap-4">
<div className="flex items-center justify-between">
<FieldLabel className="flex items-center gap-1.5 !text-sm text-muted-foreground">
Max Steps
<Tooltip>
<TooltipTrigger asChild={true}>
<button
type="button"
className="text-muted-foreground/50 hover:text-muted-foreground"
>
<HugeiconsIcon
icon={InformationCircleIcon}
className="size-3.5"
/>
</button>
</TooltipTrigger>
<TooltipContent>
Override total steps. Set 0 to use epochs instead.{" "}
<a
href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
target="_blank"
rel="noopener noreferrer"
className="text-primary underline"
>
Read more
</a>
</TooltipContent>
</Tooltip>
</FieldLabel>
<div className="flex items-center gap-3">
<Slider
value={[Math.min(maxStepsSliderMax, Math.max(0, maxSteps))]}
onValueChange={([v]) => setMaxSteps(v)}
min={0}
max={maxStepsSliderMax}
step={1}
className="w-40"
/>
<input
type="number"
value={maxSteps}
onChange={(e) => setMaxSteps(Number(e.target.value))}
min={0}
max={maxStepsSliderMax}
step={1}
className="w-16 text-right font-mono text-xs font-medium bg-muted/50 border border-border rounded-lg px-1.5 py-0.5 focus:outline-none focus:ring-1 focus:ring-primary/30 [&::-webkit-inner-spin-button]:appearance-none"
/>
<div
key={useEpochs ? "epochs" : "steps"}
className="flex flex-col gap-2 animate-in fade-in-1 slide-in-from-bottom-1 duration-200"
>
<div className="flex items-center justify-between">
<FieldLabel className="flex items-center gap-1.5 !text-sm text-muted-foreground">
{useEpochs ? "Epochs" : "Max Steps"}
<Tooltip>
<TooltipTrigger asChild={true}>
<button
type="button"
className="text-muted-foreground/50 hover:text-muted-foreground"
>
<HugeiconsIcon
icon={InformationCircleIcon}
className="size-3.5"
/>
</button>
</TooltipTrigger>
<TooltipContent>
{useEpochs
? "Number of full passes over the dataset."
: "Override total optimizer steps."}{" "}
<a
href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
target="_blank"
rel="noopener noreferrer"
className="text-primary underline"
>
Read more
</a>
</TooltipContent>
</Tooltip>
</FieldLabel>
<div className="flex items-center gap-3">
<button
type="button"
onClick={toggleUseEpochs}
className="text-xs text-primary underline cursor-pointer"
>
{useEpochs ? "Use Max Steps" : "Use Epochs"}
</button>
<Slider
value={[
useEpochs
? Math.min(epochsSliderMax, Math.max(1, epochs))
: Math.min(maxStepsSliderMax, Math.max(1, maxSteps)),
]}
onValueChange={([v]) =>
useEpochs ? setEpochs(v) : setMaxSteps(v)
}
min={1}
max={useEpochs ? epochsSliderMax : maxStepsSliderMax}
step={1}
className="w-40"
/>
<input
type="number"
value={useEpochs ? epochs : maxSteps}
onChange={(e) => {
const raw = e.target.value;
if (raw === "") return;
const value = Number(raw);
if (!Number.isFinite(value) || value < 1) return;
if (useEpochs) {
setEpochs(value);
} else {
setMaxSteps(value);
}
}}
min={1}
max={useEpochs ? epochsSliderMax : maxStepsSliderMax}
step={1}
className="w-16 text-right font-mono text-xs font-medium bg-muted/50 border border-border rounded-lg px-1.5 py-0.5 focus:outline-none focus:ring-1 focus:ring-primary/30 [&::-webkit-inner-spin-button]:appearance-none"
/>
</div>
</div>
</div>
@ -206,55 +251,6 @@ export function HyperparametersStep() {
/>
</div>
<div className="flex items-center justify-between">
<FieldLabel className="flex items-center gap-1.5 !text-sm text-muted-foreground">
Epochs
<Tooltip>
<TooltipTrigger asChild={true}>
<button
type="button"
className="text-muted-foreground/50 hover:text-muted-foreground"
>
<HugeiconsIcon
icon={InformationCircleIcon}
className="size-3.5"
/>
</button>
</TooltipTrigger>
<TooltipContent>
Number of full passes over the dataset. Set 0 to run by max
steps.{" "}
<a
href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
target="_blank"
rel="noopener noreferrer"
className="text-primary underline"
>
Read more
</a>
</TooltipContent>
</Tooltip>
</FieldLabel>
<div className="flex items-center gap-3">
<Slider
value={[Math.min(epochsSliderMax, Math.max(0, epochs))]}
onValueChange={([v]) => setEpochs(v)}
min={0}
max={epochsSliderMax}
step={1}
className="w-40"
/>
<input
type="number"
value={epochs}
onChange={(e) => setEpochs(Number(e.target.value))}
min={0}
max={epochsSliderMax}
step={1}
className="w-12 text-right font-mono text-xs font-medium bg-muted/50 border border-border rounded-lg px-1.5 py-0.5 focus:outline-none focus:ring-1 focus:ring-primary/30 [&::-webkit-inner-spin-button]:appearance-none"
/>
</div>
</div>
</div>
</FieldSet>

View file

@ -29,7 +29,7 @@ import {
OPTIMIZER_OPTIONS,
TARGET_MODULES,
} from "@/config/training";
import { useTrainingConfigStore } from "@/features/training";
import { useMaxStepsEpochsToggle, useTrainingConfigStore } from "@/features/training";
import type { GradientCheckpointing } from "@/types/training";
import {
ArrowDown01Icon,
@ -120,6 +120,15 @@ export function ParamsSection(): ReactElement {
const showVisionLora = store.isVisionModel && store.isDatasetImage === true;
const [loraOpen, setLoraOpen] = useState(false);
const [hyperOpen, setHyperOpen] = useState(false);
const { useEpochs, toggleUseEpochs } = useMaxStepsEpochsToggle({
maxSteps: store.maxSteps,
epochs: store.epochs,
saveSteps: store.saveSteps,
setMaxSteps: store.setMaxSteps,
setEpochs: store.setEpochs,
setSaveSteps: store.setSaveSteps,
});
const maxStepsSliderMax = Math.max(500, store.maxSteps, 30);
const epochsSliderMax = Math.max(20, store.epochs, 1);
@ -133,56 +142,92 @@ export function ParamsSection(): ReactElement {
className="md:min-h-[470px]"
>
<div className="flex flex-col gap-4">
{/* Max Steps */}
{/* Max Steps / Epochs */}
<div className="flex flex-col gap-2">
<div className="flex items-center justify-between">
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
Max Steps
<Tooltip>
<TooltipTrigger asChild={true}>
<button
type="button"
className="text-foreground/70 hover:text-foreground"
>
<HugeiconsIcon
icon={InformationCircleIcon}
className="size-3"
/>
</button>
</TooltipTrigger>
<TooltipContent>
Override total steps. Set 0 to use epochs instead.{" "}
<a
href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
target="_blank"
rel="noopener noreferrer"
className="text-primary underline"
>
Read more
</a>
</TooltipContent>
</Tooltip>
</span>
<input
type="number"
value={store.maxSteps}
onChange={(e) => store.setMaxSteps(Number(e.target.value))}
min={0}
max={maxStepsSliderMax}
<div
key={useEpochs ? "epochs" : "steps"}
className="flex flex-col gap-2 animate-in fade-in-0 slide-in-from-bottom-1 duration-200"
>
<div className="flex items-center justify-between">
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
{useEpochs ? "Epochs" : "Max Steps"}
<Tooltip>
<TooltipTrigger asChild={true}>
<button
type="button"
className="text-foreground/70 hover:text-foreground"
>
<HugeiconsIcon
icon={InformationCircleIcon}
className="size-3"
/>
</button>
</TooltipTrigger>
<TooltipContent>
{useEpochs
? "Number of full passes over the dataset."
: "Override total optimizer steps."}{" "}
<a
href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
target="_blank"
rel="noopener noreferrer"
className="text-primary underline"
>
Read more
</a>
</TooltipContent>
</Tooltip>
</span>
<div className="flex items-center gap-3">
<button
type="button"
onClick={toggleUseEpochs}
className="text-xs text-primary underline cursor-pointer"
>
{useEpochs ? "Use Max Steps" : "Use Epochs"}
</button>
<input
type="number"
value={useEpochs ? store.epochs : store.maxSteps}
onChange={(e) => {
const raw = e.target.value;
if (raw === "") return;
const value = Number(raw);
if (!Number.isFinite(value) || value < 1) return;
if (useEpochs) {
store.setEpochs(value);
} else {
store.setMaxSteps(value);
}
}}
min={1}
max={useEpochs ? epochsSliderMax : maxStepsSliderMax}
step={1}
className="w-16 text-right font-mono text-xs font-medium bg-muted/50 border border-border rounded-lg px-1.5 py-0.5 focus:outline-none focus:ring-1 focus:ring-primary/30 [&::-webkit-inner-spin-button]:appearance-none"
/>
</div>
</div>
<Slider
value={[
useEpochs
? Math.min(epochsSliderMax, Math.max(1, store.epochs))
: Math.min(maxStepsSliderMax, Math.max(1, store.maxSteps)),
]}
onValueChange={([v]) =>
useEpochs ? store.setEpochs(v) : store.setMaxSteps(v)
}
min={1}
max={useEpochs ? epochsSliderMax : maxStepsSliderMax}
step={1}
className="w-16 text-right font-mono text-xs font-medium bg-muted/50 border border-border rounded-lg px-1.5 py-0.5 focus:outline-none focus:ring-1 focus:ring-primary/30 [&::-webkit-inner-spin-button]:appearance-none"
/>
<p className="text-[10px] text-muted-foreground">
{useEpochs
? "Each epoch is one full pass over your dataset."
: "Limits training to a fixed number of optimizer steps."}
</p>
</div>
<Slider
value={[Math.min(maxStepsSliderMax, Math.max(0, store.maxSteps))]}
onValueChange={([v]) => store.setMaxSteps(v)}
min={0}
max={maxStepsSliderMax}
step={1}
/>
<p className="text-[10px] text-muted-foreground">
Total optimizer steps. Use 0 to run by epochs.
</p>
</div>
{/* Context length */}
@ -681,28 +726,30 @@ export function ParamsSection(): ReactElement {
max={100}
step={1}
/>
<SliderRow
label="Epochs"
tooltip={
<>
Number of full passes over the dataset. Set 0 to run by
max steps.{" "}
<a
href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
target="_blank"
rel="noopener noreferrer"
className="text-primary underline"
>
Read more
</a>
</>
}
value={store.epochs}
onChange={store.setEpochs}
min={0}
max={epochsSliderMax}
step={1}
/>
{!useEpochs && (
<SliderRow
label="Epochs"
tooltip={
<>
Number of full passes over the dataset. Set 0 to run by
max steps.{" "}
<a
href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
target="_blank"
rel="noopener noreferrer"
className="text-primary underline"
>
Read more
</a>
</>
}
value={store.epochs}
onChange={store.setEpochs}
min={0}
max={epochsSliderMax}
step={1}
/>
)}
<Row
label="Save Steps"
tooltip={

View file

@ -0,0 +1,113 @@
// 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 { useCallback, useEffect, useState } from "react";
const PREV_MAX_STEPS_KEY = "unsloth_prev_max_steps";
const PREV_SAVE_STEPS_KEY = "unsloth_prev_save_steps";
const DEFAULT_MAX_STEPS = 30;
const DEFAULT_EPOCHS = 3;
function readStoredNumber(key: string, fallback: number): number {
if (typeof window === "undefined") return fallback;
try {
const value = window.localStorage.getItem(key);
if (value === null) return fallback;
const parsed = Number(value);
return Number.isFinite(parsed) ? parsed : fallback;
} catch {
return fallback;
}
}
function writeStoredNumber(key: string, value: number): void {
if (typeof window === "undefined") return;
try {
window.localStorage.setItem(key, String(value));
} catch {
// Best effort only; ignore storage errors in restricted environments.
}
}
function normalizePrevMaxSteps(value: number): number {
return Number.isFinite(value) && value > 0 ? value : DEFAULT_MAX_STEPS;
}
function normalizePrevSaveSteps(value: number): number {
return Number.isFinite(value) && value >= 0 ? value : 0;
}
type UseMaxStepsEpochsToggleParams = {
maxSteps: number;
epochs: number;
saveSteps: number;
setMaxSteps: (value: number) => void;
setEpochs: (value: number) => void;
setSaveSteps: (value: number) => void;
defaultEpochs?: number;
};
type UseMaxStepsEpochsToggleResult = {
useEpochs: boolean;
toggleUseEpochs: () => void;
};
export function useMaxStepsEpochsToggle({
maxSteps,
epochs,
saveSteps,
setMaxSteps,
setEpochs,
setSaveSteps,
defaultEpochs = DEFAULT_EPOCHS,
}: UseMaxStepsEpochsToggleParams): UseMaxStepsEpochsToggleResult {
const useEpochs = maxSteps === 0;
const [prevMaxSteps, setPrevMaxSteps] = useState(() =>
normalizePrevMaxSteps(readStoredNumber(PREV_MAX_STEPS_KEY, DEFAULT_MAX_STEPS)),
);
const [prevSaveSteps, setPrevSaveSteps] = useState(() => {
if (maxSteps === 0 && saveSteps > 0) {
return normalizePrevSaveSteps(saveSteps);
}
return normalizePrevSaveSteps(readStoredNumber(PREV_SAVE_STEPS_KEY, 0));
});
useEffect(() => {
if (maxSteps > 0) {
const normalized = normalizePrevMaxSteps(maxSteps);
setPrevMaxSteps(normalized);
writeStoredNumber(PREV_MAX_STEPS_KEY, normalized);
}
}, [maxSteps]);
useEffect(() => {
if (!useEpochs) {
const normalized = normalizePrevSaveSteps(saveSteps);
setPrevSaveSteps(normalized);
writeStoredNumber(PREV_SAVE_STEPS_KEY, normalized);
}
}, [saveSteps, useEpochs]);
const toggleUseEpochs = useCallback(() => {
if (useEpochs) {
setMaxSteps(normalizePrevMaxSteps(prevMaxSteps));
setSaveSteps(normalizePrevSaveSteps(prevSaveSteps));
return;
}
setMaxSteps(0);
setEpochs(epochs || defaultEpochs);
setSaveSteps(0);
}, [
defaultEpochs,
epochs,
prevMaxSteps,
prevSaveSteps,
setEpochs,
setMaxSteps,
setSaveSteps,
useEpochs,
]);
return { useEpochs, toggleUseEpochs };
}

View file

@ -8,6 +8,7 @@ export {
} from "./stores/training-runtime-store";
export { useTrainingActions } from "./hooks/use-training-actions";
export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle";
export { useMaxStepsEpochsToggle } from "./hooks/use-max-steps-epochs-toggle";
export { HfDatasetSubsetSplitSelectors } from "./components/hf-dataset-subset-split-selectors";
export { useDatasetPreviewDialogStore } from "./stores/dataset-preview-dialog-store";
export { uploadTrainingDataset } from "./api/datasets-api";