feat(studio): Mac-aware training controls for MLX (optimizers, LoftQ, packing) (#7358)
* 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>
This commit is contained in:
Long Yixing
GitHub
parent
5aedfd0b46
commit
e143e1ce33
6 changed files with 115 additions and 26 deletions
|
|
@ -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" },
|
||||
|
|
|
|||
|
|
@ -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 {
|
|||
<button
|
||||
key={opt.value}
|
||||
type="button"
|
||||
disabled={isMac && opt.value === "loftq"}
|
||||
onClick={() => store.setLoraVariant(opt.value)}
|
||||
className={`flex-1 corner-squircle rounded-xl border px-3 py-2 text-left transition-colors cursor-pointer ${
|
||||
className={`flex-1 corner-squircle rounded-xl border px-3 py-2 text-left transition-colors cursor-pointer disabled:cursor-not-allowed disabled:opacity-60 ${
|
||||
store.loraVariant === opt.value
|
||||
? "border-ring-strong bg-primary/5"
|
||||
: "border-border hover:border-foreground/20"
|
||||
|
|
@ -715,7 +753,9 @@ export function ParamsSection(): ReactElement {
|
|||
>
|
||||
<p className="text-xs font-medium">{opt.label}</p>
|
||||
<p className="text-ui-10 text-muted-foreground">
|
||||
{opt.desc}
|
||||
{isMac && opt.value === "loftq"
|
||||
? "Not supported on Apple Silicon"
|
||||
: opt.desc}
|
||||
</p>
|
||||
</button>
|
||||
))}
|
||||
|
|
@ -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",
|
||||
)}{" "}
|
||||
<a
|
||||
href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
|
||||
target="_blank"
|
||||
|
|
@ -778,14 +822,14 @@ export function ParamsSection(): ReactElement {
|
|||
}
|
||||
>
|
||||
<Select
|
||||
value={store.optimizerType}
|
||||
value={selectedOptimizer}
|
||||
onValueChange={(v) => store.setOptimizerType(v)}
|
||||
>
|
||||
<SelectTrigger className="w-48">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{OPTIMIZER_OPTIONS.map((opt) => (
|
||||
{optimizerOptions.map((opt) => (
|
||||
<SelectItem key={opt.value} value={opt.value}>
|
||||
{formatOptimizerLabel(opt.value, opt.label, t)}
|
||||
</SelectItem>
|
||||
|
|
@ -1105,14 +1149,37 @@ export function ParamsSection(): ReactElement {
|
|||
<Checkbox
|
||||
id="packing"
|
||||
checked={store.packing}
|
||||
disabled={isMac}
|
||||
onCheckedChange={(v) => store.setPacking(!!v)}
|
||||
/>
|
||||
<label
|
||||
htmlFor="packing"
|
||||
className="text-xs cursor-pointer text-muted-foreground"
|
||||
className={`text-xs text-muted-foreground ${
|
||||
isMac
|
||||
? "cursor-not-allowed opacity-60"
|
||||
: "cursor-pointer"
|
||||
}`}
|
||||
>
|
||||
{t("studio.params.enablePacking")}
|
||||
</label>
|
||||
{isMac && (
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<button
|
||||
type="button"
|
||||
className="text-foreground/70 hover:text-foreground"
|
||||
>
|
||||
<HugeiconsIcon
|
||||
icon={InformationCircleIcon}
|
||||
className="size-3"
|
||||
/>
|
||||
</button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent>
|
||||
Packing is not supported on Apple Silicon (MLX).
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
{!store.isEmbeddingModel && !isCpt && !isRawText && (
|
||||
|
|
|
|||
|
|
@ -19,7 +19,8 @@ import {
|
|||
PopoverTrigger,
|
||||
} from "@/components/ui/popover";
|
||||
import { Progress } from "@/components/ui/progress";
|
||||
import { OPTIMIZER_OPTIONS } from "@/config/training";
|
||||
import { usePlatformStore } from "@/config/env";
|
||||
import { MLX_OPTIMIZER_OPTIONS, OPTIMIZER_OPTIONS } from "@/config/training";
|
||||
import { setTrainingCompareHandoff } from "@/features/chat";
|
||||
import {
|
||||
useTrainingActions,
|
||||
|
|
@ -92,6 +93,7 @@ export function ProgressSection({
|
|||
}: ProgressSectionProps): ReactElement {
|
||||
const t = useT();
|
||||
const navigate = useNavigate();
|
||||
const platformDeviceType = usePlatformStore((s) => s.deviceType);
|
||||
const trainingMethodLabel = getTrainingMethodLabel(data.trainingMethod);
|
||||
|
||||
const config = useTrainingConfigStore(
|
||||
|
|
@ -187,9 +189,18 @@ export function ProgressSection({
|
|||
const cfgLoraDropout = cfg?.loraDropout;
|
||||
const cfgLoraVariant = cfg?.loraVariant;
|
||||
|
||||
// Mirror the training form: on Mac the CUDA/bitsandbytes optimizer names run
|
||||
// as plain AdamW (the MLX backend normalizes them), so label them AdamW here
|
||||
// too rather than by the requested, unnormalized name.
|
||||
const effectiveOptimizer =
|
||||
platformDeviceType === "mac" &&
|
||||
OPTIMIZER_OPTIONS.some((o) => o.value === cfgOptimizerType)
|
||||
? "adamw"
|
||||
: cfgOptimizerType;
|
||||
const optimizerLabel =
|
||||
OPTIMIZER_OPTIONS.find((o) => o.value === cfgOptimizerType)?.label ??
|
||||
cfgOptimizerType;
|
||||
[...OPTIMIZER_OPTIONS, ...MLX_OPTIMIZER_OPTIONS].find(
|
||||
(o) => o.value === effectiveOptimizer,
|
||||
)?.label ?? effectiveOptimizer;
|
||||
|
||||
const configItems: ConfigGroup[] = [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -1051,6 +1051,8 @@ export const en = {
|
|||
optimizer: "Optimizer",
|
||||
optimizerTooltip:
|
||||
"Optimization algorithm. 8-bit variants reduce memory usage. Fused is recommended for vision models.",
|
||||
optimizerTooltipMlx:
|
||||
"Optimization algorithm. AdamW is the default. Lion uses less memory but usually needs a lower learning rate.",
|
||||
lrScheduler: "LR scheduler",
|
||||
lrSchedulerTooltip:
|
||||
"How the learning rate changes over training. Linear decays steadily; cosine decays in a curve.",
|
||||
|
|
|
|||
|
|
@ -38,7 +38,11 @@ def fail(m):
|
|||
raise AssertionError(f"[font-scale] FAIL: {m}")
|
||||
|
||||
|
||||
def near(a, b, tol = 0.35):
|
||||
def near(
|
||||
a,
|
||||
b,
|
||||
tol = 0.35,
|
||||
):
|
||||
return a is not None and b is not None and abs(a - b) <= tol
|
||||
|
||||
|
||||
|
|
@ -86,9 +90,7 @@ def open_appearance(page):
|
|||
page.wait_for_timeout(700)
|
||||
if page.get_by_role("dialog").count() == 0:
|
||||
fail("settings dialog did not open")
|
||||
page.get_by_role("dialog").get_by_role("button").filter(
|
||||
has_text = "Appearance"
|
||||
).first.click()
|
||||
page.get_by_role("dialog").get_by_role("button").filter(has_text = "Appearance").first.click()
|
||||
page.wait_for_timeout(600)
|
||||
|
||||
|
||||
|
|
@ -153,9 +155,7 @@ def main():
|
|||
page.wait_for_timeout(400)
|
||||
|
||||
step("overflowing select scrolls its Radix viewport")
|
||||
page.get_by_role("dialog").get_by_role("button").filter(
|
||||
has_text = "Voice"
|
||||
).first.click()
|
||||
page.get_by_role("dialog").get_by_role("button").filter(has_text = "Voice").first.click()
|
||||
page.wait_for_timeout(600)
|
||||
page.set_viewport_size({"width": 1440, "height": 480})
|
||||
page.locator("[aria-label='Dictation language']").click()
|
||||
|
|
@ -194,9 +194,7 @@ def main():
|
|||
page.wait_for_timeout(400)
|
||||
|
||||
step("default restores exactly")
|
||||
page.get_by_role("dialog").get_by_role("button").filter(
|
||||
has_text = "Appearance"
|
||||
).first.click()
|
||||
page.get_by_role("dialog").get_by_role("button").filter(has_text = "Appearance").first.click()
|
||||
page.wait_for_timeout(500)
|
||||
set_input(page, "UI font size", DEFAULT)
|
||||
final = measure(page)
|
||||
|
|
|
|||
|
|
@ -15,9 +15,7 @@ from pathlib import Path
|
|||
REPO = Path(__file__).resolve().parents[2]
|
||||
SRC = REPO / "studio/frontend/src"
|
||||
INDEX_CSS = (SRC / "index.css").read_text(encoding = "utf-8")
|
||||
STORE = (SRC / "features/settings/stores/appearance-custom-store.ts").read_text(
|
||||
encoding = "utf-8"
|
||||
)
|
||||
STORE = (SRC / "features/settings/stores/appearance-custom-store.ts").read_text(encoding = "utf-8")
|
||||
SELECT = (SRC / "components/ui/select.tsx").read_text(encoding = "utf-8")
|
||||
|
||||
# Raw numeric fontSize props are only allowed where a scaled stylesheet rule
|
||||
|
|
@ -70,9 +68,7 @@ def test_ui_token_families_exist():
|
|||
|
||||
|
||||
def test_explicit_code_font_size_is_never_multiplied():
|
||||
match = re.search(
|
||||
r"html\[data-code-font-size\][^{]*\{([^}]*)\}", INDEX_CSS
|
||||
)
|
||||
match = re.search(r"html\[data-code-font-size\][^{]*\{([^}]*)\}", INDEX_CSS)
|
||||
assert match is not None
|
||||
body = match.group(1)
|
||||
assert "var(--custom-code-font-size)" in body
|
||||
|
|
@ -83,7 +79,9 @@ def test_radix_select_viewport_owns_the_scroll_state():
|
|||
viewport = SELECT[SELECT.index("SelectPrimitive.Viewport") :]
|
||||
assert "overflow-y-auto" in viewport.split("</SelectPrimitive.Viewport>")[0]
|
||||
# The rounded surface itself must not scroll (WebKit squares its corners).
|
||||
content_cls = re.search(r"SelectPrimitive\.Content[\s\S]*?className=\{cn\(\s*\"([^\"]+)\"", SELECT)
|
||||
content_cls = re.search(
|
||||
r"SelectPrimitive\.Content[\s\S]*?className=\{cn\(\s*\"([^\"]+)\"", SELECT
|
||||
)
|
||||
assert content_cls is not None
|
||||
assert "overflow-hidden" in content_cls.group(1)
|
||||
assert "overflow-y-auto" not in content_cls.group(1)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue