added lr_scheduler type to the frontend

studio/frontend/src/config/training.ts

@@ -82,11 +82,17 @@ export const OPTIMIZER_OPTIONS: ReadonlyArray<{ value: string; label: string }>
   { value: "adamw_torch_fused", label: "AdamW (PyTorch Fused)" },
 ];
 
+export const LR_SCHEDULER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [
+  { value: "linear", label: "Linear" },
+  { value: "cosine", label: "Cosine" },
+];
+
 export const DEFAULT_HYPERPARAMS = {
   epochs: 3,
   contextLength: 2048,
   learningRate: 2e-4,
   optimizerType: "adamw_8bit",
+  lrSchedulerType: "linear",
   loraRank: 16,
   loraAlpha: 32,
   loraDropout: 0.05,

studio/frontend/src/features/studio/sections/params-section.tsx

@@ -22,6 +22,7 @@ import {
 } from "@/components/ui/tooltip";
 import {
   CONTEXT_LENGTHS,
+  LR_SCHEDULER_OPTIONS,
   OPTIMIZER_OPTIONS,
   TARGET_MODULES,
 } from "@/config/training";
@@ -549,6 +550,43 @@ export function ParamsSection(): ReactElement {
                     </SelectContent>
                   </Select>
                 </Row>
+                <Row
+                  label="LR scheduler"
+                  tooltip={
+                    <>
+                      How the learning rate changes over training. Linear decays
+                      steadily; cosine decays in a curve.{" "}
+                      <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>
+                    </>
+                  }
+                >
+                  <Select
+                    value={store.lrSchedulerType}
+                    onValueChange={(v) => store.setLrSchedulerType(v)}
+                  >
+                    <SelectTrigger className="w-48 font-mono">
+                      <SelectValue />
+                    </SelectTrigger>
+                    <SelectContent>
+                      {LR_SCHEDULER_OPTIONS.map((opt) => (
+                        <SelectItem
+                          key={opt.value}
+                          value={opt.value}
+                          className="font-mono"
+                        >
+                          {opt.label}
+                        </SelectItem>
+                      ))}
+                    </SelectContent>
+                  </Select>
+                </Row>
                 <SliderRow
                   label="Batch Size"
                   tooltip={

studio/frontend/src/features/training/api/mappers.ts

@@ -41,7 +41,7 @@ export function buildTrainingStartPayload(
     random_seed: config.randomSeed,
     packing: config.packing,
     optim: config.optimizerType,
-    lr_scheduler_type: "linear",
+    lr_scheduler_type: config.lrSchedulerType,
     use_lora: adapterMethod,
     lora_r: config.loraRank,
     lora_alpha: config.loraAlpha,

studio/frontend/src/features/training/api/models-api.ts

@@ -9,6 +9,8 @@ interface BackendTrainingDefaults {
   max_seq_length?: number;
   num_epochs?: number;
   learning_rate?: number | string;
+  optim?: string;
+  lr_scheduler_type?: string;
   batch_size?: number;
   gradient_accumulation_steps?: number;
   warmup_steps?: number;

studio/frontend/src/features/training/lib/model-defaults.ts

@@ -7,6 +7,8 @@ type ModelDefaultsPatch = Partial<
     | "epochs"
     | "contextLength"
     | "learningRate"
+    | "optimizerType"
+    | "lrSchedulerType"
     | "loraRank"
     | "loraAlpha"
     | "loraDropout"
@@ -86,6 +88,12 @@ export function mapBackendModelConfigToTrainingPatch(
   const learningRate = toNumber(training?.learning_rate);
   if (learningRate !== undefined) patch.learningRate = learningRate;
 
+  const optim = toStringValue(training?.optim);
+  if (optim !== undefined) patch.optimizerType = optim;
+
+  const lrSchedulerType = toStringValue(training?.lr_scheduler_type);
+  if (lrSchedulerType !== undefined) patch.lrSchedulerType = lrSchedulerType;
+
   const batchSize = toNumber(training?.batch_size);
   if (batchSize !== undefined) patch.batchSize = batchSize;
 

studio/frontend/src/features/training/stores/training-config-store.ts

@@ -306,6 +306,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         setContextLength: (contextLength) => set({ contextLength }),
         setLearningRate: (learningRate) => set({ learningRate }),
         setOptimizerType: (optimizerType) => set({ optimizerType }),
+        setLrSchedulerType: (lrSchedulerType) => set({ lrSchedulerType }),
         setLoraRank: (loraRank) => set({ loraRank }),
         setLoraAlpha: (loraAlpha) => set({ loraAlpha }),
         setLoraDropout: (loraDropout) => set({ loraDropout }),
@@ -347,7 +348,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
     },
     {
       name: "unsloth_training_config_v1",
-      version: 4,
+      version: 5,
       migrate: (persisted, version) => {
         const s = persisted as Record<string, unknown>;
         if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
@@ -360,6 +361,9 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         if (version < 4 && s.optimizerType == null) {
           s.optimizerType = DEFAULT_HYPERPARAMS.optimizerType;
         }
+        if (version < 5 && s.lrSchedulerType == null) {
+          s.lrSchedulerType = DEFAULT_HYPERPARAMS.lrSchedulerType;
+        }
         return s as unknown as TrainingConfigStore;
       },
       partialize: partializePersistedState,

studio/frontend/src/features/training/types/config.ts

@@ -31,6 +31,7 @@ export interface TrainingConfigState {
   contextLength: number;
   learningRate: number;
   optimizerType: string;
+  lrSchedulerType: string;
   loraRank: number;
   loraAlpha: number;
   loraDropout: number;
@@ -87,6 +88,7 @@ export interface TrainingConfigActions {
   setContextLength: (length: number) => void;
   setLearningRate: (rate: number) => void;
   setOptimizerType: (value: string) => void;
+  setLrSchedulerType: (value: string) => void;
   setLoraRank: (rank: number) => void;
   setLoraAlpha: (alpha: number) => void;
   setLoraDropout: (dropout: number) => void;