diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts
index 9ad09ba813..da60328d40 100644
--- a/studio/frontend/src/config/training.ts
+++ b/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,
diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx
index cf25d5ba54..edc21e0e0e 100644
--- a/studio/frontend/src/features/studio/sections/params-section.tsx
+++ b/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 {
+
+ How the learning rate changes over training. Linear decays
+ steadily; cosine decays in a curve.{" "}
+
+ Read more
+
+ >
+ }
+ >
+
+
()(
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()(
},
{
name: "unsloth_training_config_v1",
- version: 4,
+ version: 5,
migrate: (persisted, version) => {
const s = persisted as Record;
if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
@@ -360,6 +361,9 @@ export const useTrainingConfigStore = create()(
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,
diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts
index 8557739b06..b268a08b79 100644
--- a/studio/frontend/src/features/training/types/config.ts
+++ b/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;