From f0a4f77028bbf30c217e8eeb18b1997d9d40460e Mon Sep 17 00:00:00 2001
From: Shine1i <wasimysdev@gmail.com>
Date: Tue, 17 Feb 2026 20:08:02 +0100
Subject: [PATCH] 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`.
---
 .../src/components/ui/alert-dialog.tsx        | 28 +++----
 .../components/steps/hyperparameters-step.tsx | 78 ++++++++++++++++---
 .../studio/sections/params-section.tsx        | 50 ++++++------
 .../studio/sections/progress-section.tsx      | 26 ++++++-
 4 files changed, 132 insertions(+), 50 deletions(-)

diff --git a/studio/frontend/src/components/ui/alert-dialog.tsx b/studio/frontend/src/components/ui/alert-dialog.tsx
index b39a310c49..61327e525e 100644
--- 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,
diff --git a/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx b/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx
index 45f9c82878..9ad9c31e0a 100644
--- 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>
 
diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx
index 141a38286d..2d7968e6c5 100644
--- 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={
diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx
index 76282af813..2d075e870b 100644
--- 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>
           )}