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;