From e143e1ce33b7c7cffcd25b0de85459698de2619e Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Thu, 23 Jul 2026 17:14:20 +0800 Subject: [PATCH] feat(studio): Mac-aware training controls for MLX (optimizers, LoftQ, packing) (#7358) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(studio): offer MLX-supported optimizers on Apple Silicon The training form's optimizer dropdown only listed CUDA/bitsandbytes optimizers (adamw_8bit, paged variants, torch fused). On Apple Silicon the MLX trainer supports a different set (adamw, adam, lion, muon, sgd, adafactor) and remaps every bitsandbytes/torch name to plain AdamW, so the dropdown misrepresented what actually runs. Offer the MLX optimizer list when the device is a Mac, and derive the displayed value so the control is never blank: the shared CUDA default and the other bitsandbytes/torch options render as AdamW (exactly how the MLX backend normalizes them), while any other value is shown as-is so an unrecognized or non-canonical imported optimizer is never mislabeled. Non-Mac behavior is unchanged. The run-summary optimizer label now resolves from both lists. * feat(studio): show an MLX-appropriate optimizer tooltip on Apple Silicon The optimizer tooltip described "8-bit variants" and recommended "Fused" for vision models, neither of which is offered when training runs on MLX. On Apple Silicon, show a tooltip that matches the MLX optimizer set and notes that Lion typically needs a lower learning rate than AdamW. Copy-only: no change to the selected optimizer or the learning rate, and the non-Mac tooltip is unchanged. The new string is added to the English locale; other locales fall back to English until translated, matching how new keys are handled elsewhere. * fix(studio): label Mac CUDA-alias optimizers as AdamW in the run summary On Apple Silicon the run-configuration summary looked up the stored optimizer name directly, so a run that kept a CUDA/bitsandbytes default such as adamw_8bit was labeled "AdamW 8-bit" even though the picker shows "AdamW" and the MLX backend runs plain AdamW. Mirror the training form's derivation so those aliases are labeled AdamW in the summary too. Display-only: no change to the stored or submitted optimizer, and non-Mac summaries are unchanged. * feat(studio): disable LoftQ and sequence packing on Apple Silicon Neither LoftQ nor sequence packing is supported on MLX — the backend rejects LoftQ and the trainer silently forces packing off — yet the training form still offered both on Apple Silicon. Disable the LoftQ LoRA-init option (greyed and unclickable, with an inline "Not supported on Apple Silicon" note) and the "Enable packing" checkbox (greyed, with a tooltip explaining why), matching how the unsupported "Enable streaming" control is presented. Clearing effects reset a stale loftq/packing value to its default on Mac so the disabled controls never submit it. Non-Mac behavior is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/frontend/src/config/training.ts | 13 +++ .../studio/sections/params-section.tsx | 79 +++++++++++++++++-- .../studio/sections/progress-section.tsx | 17 +++- studio/frontend/src/i18n/locales/en.ts | 2 + tests/studio/playwright_ui_font_scale.py | 18 ++--- tests/studio/test_ui_font_scale_contract.py | 12 ++- 6 files changed, 115 insertions(+), 26 deletions(-) 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", + )}{" "}