feat: enhance training flow with new runtime hints, adjustable steps/epochs

- Added halfway/completed training hints with actionable links. - Introduced sliders for adjusting max steps and epochs dynamically. - Refined tooltip explanations for configuration parameters. - Enabled custom overlay styling for `AlertDialogContent`.
2026-02-17 20:08:02 +01:00 · 2026-02-17 20:08:02 +01:00 · f0a4f77028
commit f0a4f77028
parent f02cd32c56
4 changed files with 132 additions and 50 deletions
--- a/studio/frontend/src/components/ui/alert-dialog.tsx
+++ b/studio/frontend/src/components/ui/alert-dialog.tsx
@ -42,19 +42,21 @@ function AlertDialogOverlay({
  );
 }

-function AlertDialogContent({
-  className,
-  size = "default",
-  ...props
-}: React.ComponentProps<typeof AlertDialogPrimitive.Content> & {
-  size?: "default" | "sm";
-}) {
-  return (
-    <AlertDialogPortal>
-      <AlertDialogOverlay />
-      <AlertDialogPrimitive.Content
-        data-slot="alert-dialog-content"
-        data-size={size}
+function AlertDialogContent({
+  className,
+  size = "default",
+  overlayClassName,
+  ...props
+}: React.ComponentProps<typeof AlertDialogPrimitive.Content> & {
+  size?: "default" | "sm";
+  overlayClassName?: string;
+}) {
+  return (
+    <AlertDialogPortal>
+      <AlertDialogOverlay className={overlayClassName} />
+      <AlertDialogPrimitive.Content
+        data-slot="alert-dialog-content"
+        data-size={size}
        className={cn(
          "data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 bg-background ring-foreground/5 gap-6 rounded-4xl p-6 ring-1 duration-100 data-[size=default]:max-w-xs data-[size=sm]:max-w-xs data-[size=default]:sm:max-w-md group/alert-dialog-content fixed top-1/2 left-1/2 z-50 grid w-full -translate-x-1/2 -translate-y-1/2 outline-none",
          className,
--- a/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx
+++ b/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx
@ -28,6 +28,8 @@ import { useShallow } from "zustand/react/shallow";
 export function HyperparametersStep() {
  const {
    trainingMethod,
+    maxSteps,
+    setMaxSteps,
    epochs,
    setEpochs,
    contextLength,
@ -43,6 +45,8 @@ export function HyperparametersStep() {
  } = useTrainingConfigStore(
    useShallow((s) => ({
      trainingMethod: s.trainingMethod,
+      maxSteps: s.maxSteps,
+      setMaxSteps: s.setMaxSteps,
      epochs: s.epochs,
      setEpochs: s.setEpochs,
      contextLength: s.contextLength,
@ -60,6 +64,8 @@ export function HyperparametersStep() {

  const showLoraParams =
    trainingMethod === "lora" || trainingMethod === "qlora";
+  const maxStepsSliderMax = Math.max(500, maxSteps, 30);
+  const epochsSliderMax = Math.max(10, epochs, 1);

  return (
    <FieldGroup>
@ -68,7 +74,7 @@ export function HyperparametersStep() {
        <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">
-              Epochs
+              Max Steps
              <Tooltip>
                <TooltipTrigger asChild={true}>
                  <button
@ -82,7 +88,7 @@ export function HyperparametersStep() {
                  </button>
                </TooltipTrigger>
                <TooltipContent>
-                  Number of times to iterate over the entire dataset.{" "}
+                  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"
@ -96,21 +102,21 @@ export function HyperparametersStep() {
            </FieldLabel>
            <div className="flex items-center gap-3">
              <Slider
-                value={[epochs]}
-                onValueChange={([v]) => setEpochs(v)}
-                min={1}
-                max={10}
+                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={epochs}
-                onChange={(e) => setEpochs(Number(e.target.value))}
-                min={1}
-                max={10}
+                value={maxSteps}
+                onChange={(e) => setMaxSteps(Number(e.target.value))}
+                min={0}
+                max={maxStepsSliderMax}
                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"
+                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>
@ -196,6 +202,56 @@ export function HyperparametersStep() {
              className="w-32 font-mono"
            />
          </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>

--- a/studio/frontend/src/features/studio/sections/params-section.tsx
+++ b/studio/frontend/src/features/studio/sections/params-section.tsx
@ -112,6 +112,8 @@ export function ParamsSection(): ReactElement {
  const showVisionLora = store.isVisionModel && store.isDatasetMultimodal === true;
  const [loraOpen, setLoraOpen] = useState(false);
  const [hyperOpen, setHyperOpen] = useState(false);
+  const maxStepsSliderMax = Math.max(500, store.maxSteps, 30);
+  const epochsSliderMax = Math.max(20, store.epochs, 1);

  return (
    <div data-tour="studio-params" className="lg:col-span-4">
@ -123,11 +125,11 @@ export function ParamsSection(): ReactElement {
        className="min-h-[450px]"
      >
        <div className="flex flex-col gap-4">
-        {/* Epochs */}
+        {/* Max Steps */}
        <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">
-              Epochs
+              Max Steps
              <Tooltip>
                <TooltipTrigger asChild={true}>
                  <button
@ -141,8 +143,7 @@ export function ParamsSection(): ReactElement {
                  </button>
                </TooltipTrigger>
                <TooltipContent>
-                  How many times the model sees the entire dataset during
-                  training.{" "}
+                  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"
@ -156,23 +157,23 @@ export function ParamsSection(): ReactElement {
            </span>
            <input
              type="number"
-              value={store.epochs}
-              onChange={(e) => store.setEpochs(Number(e.target.value))}
-              min={1}
-              max={20}
+              value={store.maxSteps}
+              onChange={(e) => store.setMaxSteps(Number(e.target.value))}
+              min={0}
+              max={maxStepsSliderMax}
              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"
+              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>
          <Slider
-            value={[store.epochs]}
-            onValueChange={([v]) => store.setEpochs(v)}
-            min={1}
-            max={20}
+            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">
-            Number of full passes over the training dataset
+            Total optimizer steps. Use 0 to run by epochs.
          </p>
        </div>

@ -602,11 +603,12 @@ export function ParamsSection(): ReactElement {
                  max={100}
                  step={1}
                />
-                <Row
-                  label="Max Steps"
+                <SliderRow
+                  label="Epochs"
                  tooltip={
                    <>
-                      Override total steps. 0 means use epochs instead.{" "}
+                      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"
@ -617,14 +619,12 @@ export function ParamsSection(): ReactElement {
                      </a>
                    </>
                  }
-                >
-                  <Input
-                    type="number"
-                    value={store.maxSteps}
-                    onChange={(e) => store.setMaxSteps(Number(e.target.value))}
-                    className="w-28 font-mono"
-                  />
-                </Row>
+                  value={store.epochs}
+                  onChange={store.setEpochs}
+                  min={0}
+                  max={epochsSliderMax}
+                  step={1}
+                />
                <Row
                  label="Save Steps"
                  tooltip={
--- a/studio/frontend/src/features/studio/sections/progress-section.tsx
+++ b/studio/frontend/src/features/studio/sections/progress-section.tsx
@ -31,6 +31,7 @@ import {
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { useState, type ReactElement, type ReactNode } from "react";
+import { Link } from "@tanstack/react-router";
 import { useShallow } from "zustand/react/shallow";
 import { useGpuUtilization } from "@/hooks";
 import { formatDuration, formatNumber, phaseColors, phaseLabel } from "./progress-section-lib";
@ -101,6 +102,9 @@ export function ProgressSection(): ReactElement {
    elapsed != null && elapsed > 0
      ? runtime.currentStep / elapsed
      : null;
+  const showHalfwayHint =
+    runtime.phase === "training" && pct >= 50 && pct < 100;
+  const showCompletedHint = runtime.phase === "completed";

  const stoppedLoss = getDisplayMetric(
    runtime.isTrainingRunning,
@ -198,7 +202,7 @@ export function ProgressSection(): ReactElement {
            >
              <HugeiconsIcon icon={StopIcon} className="size-3" /> Stop
            </Button>
-            <AlertDialogContent>
+            <AlertDialogContent overlayClassName="bg-background/40 supports-backdrop-filter:backdrop-blur-[1px]">
              <AlertDialogHeader>
                <AlertDialogTitle>Stop Training</AlertDialogTitle>
                <AlertDialogDescription>
@ -252,6 +256,26 @@ export function ProgressSection(): ReactElement {
            </div>
          </div>

+          {(showHalfwayHint || showCompletedHint) && (
+            <div className="rounded-xl border border-emerald-500/25 bg-emerald-500/8 p-3">
+              <p className="text-xs font-medium text-emerald-900 dark:text-emerald-200">
+                {showCompletedHint
+                  ? "Training done. Next step: compare base vs fine-tuned outputs."
+                  : "Halfway done. Training is past 50%."}
+              </p>
+              {showCompletedHint && (
+                <div className="mt-2 flex flex-wrap gap-2">
+                  <Button asChild={true} size="xs">
+                    <Link to="/chat">Compare in Chat</Link>
+                  </Button>
+                  <Button asChild={true} size="xs" variant="outline">
+                    <Link to="/export">Export Model</Link>
+                  </Button>
+                </div>
+              )}
+            </div>
+          )}
+
          {runtime.error && (
            <p className="text-xs text-red-500 leading-relaxed">{runtime.error}</p>
          )}