diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 873e9aa203..10f7e93e5e 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -92,6 +92,19 @@ export const OPTIMIZER_OPTIONS: ReadonlyArray<{ value: string; label: string }> { value: "adamw_torch_fused", label: "AdamW (PyTorch Fused)" }, ]; +// Optimizers the MLX trainer actually supports on Apple Silicon. Values must +// match SUPPORTED_MLX_OPTIMIZERS in unsloth-zoo's mlx/trainer.py; on MLX the +// bitsandbytes/torch names above have no meaning and are remapped to plain +// AdamW, so Studio offers this list instead when running on a Mac. +export const MLX_OPTIMIZER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [ + { value: "adamw", label: "AdamW" }, + { value: "adam", label: "Adam" }, + { value: "lion", label: "Lion" }, + { value: "muon", label: "Muon" }, + { value: "sgd", label: "SGD" }, + { value: "adafactor", label: "Adafactor" }, +]; + export const LR_SCHEDULER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [ { value: "linear", label: "Linear" }, { value: "cosine", label: "Cosine" }, diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 029fef4e19..3270eb4e3d 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -36,6 +36,7 @@ import { CONTEXT_LENGTHS, CPT_TARGET_MODULES, LR_SCHEDULER_OPTIONS, + MLX_OPTIMIZER_OPTIONS, OPTIMIZER_OPTIONS, TARGET_MODULES, } from "@/config/training"; @@ -204,6 +205,42 @@ export function ParamsSection(): ReactElement { setCtxInput(String(store.contextLength)); }, [store.contextLength]); + // On Apple Silicon the MLX trainer supports a different optimizer set than + // the CUDA/bitsandbytes list, so offer the MLX names there. + const isMac = platformDeviceType === "mac"; + const optimizerOptions = isMac ? MLX_OPTIMIZER_OPTIONS : OPTIMIZER_OPTIONS; + + // On Mac, the MLX backend normalizes every CUDA/bitsandbytes optimizer in + // OPTIMIZER_OPTIONS (including the shared default) to plain AdamW, so show + // AdamW for those to keep the control truthful and non-blank. Any other + // value -- an MLX optimizer the user picked, or an unrecognized/non-canonical + // imported one -- is shown as-is rather than mislabeled as AdamW, since the + // backend would run or reject it on its own terms. Non-Mac display unchanged. + const isCudaAliasOptimizer = OPTIMIZER_OPTIONS.some( + (o) => o.value === store.optimizerType, + ); + const selectedOptimizer = + isMac && isCudaAliasOptimizer ? "adamw" : store.optimizerType; + + // LoftQ is not supported on MLX (the backend rejects it), so clear a stale + // selection to lora on Apple Silicon -- whether persisted, applied from a + // model default, or imported -- so the backend never receives it. + const setLoraVariant = store.setLoraVariant; + useEffect(() => { + if (isMac && store.loraVariant === "loftq") { + setLoraVariant("lora"); + } + }, [isMac, store.loraVariant, setLoraVariant]); + + // Packing is not supported on MLX (the backend forces it off), so clear it on + // Apple Silicon -- the checkbox is disabled and the flag is never sent. + const setPacking = store.setPacking; + useEffect(() => { + if (isMac && store.packing) { + setPacking(false); + } + }, [isMac, store.packing, setPacking]); + const trySetContextLength = (input: string): number | null => { const n = Number(input); if (Number.isInteger(n) && n > 0) { @@ -706,8 +743,9 @@ export function ParamsSection(): ReactElement { ))} @@ -765,7 +805,11 @@ export function ParamsSection(): ReactElement { label={t("studio.params.optimizer")} tooltip={ <> - {t("studio.params.optimizerTooltip")}{" "} + {t( + isMac + ? "studio.params.optimizerTooltipMlx" + : "studio.params.optimizerTooltip", + )}{" "}