feat(studio): MLX training tab on Apple Silicon (LoRA / full FT, VLM, export)#5265
Conversation
Rewrite __init__.py: detect MLX on macOS arm64 before any torch imports Extract original GPU init to _gpu_init.py (unchanged) MLX path imports FastMLXModel from unsloth_zoo, skips all GPU code GPU path unchanged: from ._gpu_init import *
- Rewrite __init__.py: detect MLX on macOS arm64 before any torch imports - Extract original GPU init to _gpu_init.py (unchanged) - MLX path imports FastMLXModel from unsloth_zoo, skips all GPU code - GPU path unchanged: from ._gpu_init import *
…streaming for VLM
…streaming for VLM
…wide Studio UI was showing ~95 GB during MLX training because get_gpu_utilization read "In use system memory" from IORegistry's AGXAccelerator — system-wide GPU memory across all processes (training + backend + browser + Display). Now the trainer's mx.get_peak_memory value is forwarded through the progress event and surfaced via /api/train/hardware while training is active. Falls back to the system-wide reading when training is not running.
…wide Studio UI was showing ~95 GB during MLX training because get_gpu_utilization read "In use system memory" from IORegistry's AGXAccelerator — system-wide GPU memory across all processes (training + backend + browser + Display). Now the trainer's mx.get_peak_memory() value is forwarded through the progress event and surfaced via /api/train/hardware while training is active. Falls back to the system-wide reading when training is not running.
M1 and M2 chips emulate bf16 in software on the GPU, causing 40-70% slower prefill compared to native fp16. M3+ have native bf16 (macOS Sonoma+ MPSGraph). Replaces the always-True stub with chip-aware detection via mx.device_info.
M1 and M2 chips emulate bf16 in software on the GPU, causing 40-70% slower prefill compared to native fp16. M3+ have native bf16 (macOS Sonoma+ MPSGraph). Replaces the always-True stub with chip-aware detection via mx.device_info().
danielhanchen
added
the
auto-addresses-issue
|
This PR appears to address open issue(s). The duplicate detector matched the following open issues with HIGH confidence:
If this PR fixes any of them, consider adding |
danielhanchen
added
the
auto-has-duplicate
|
Possible duplicate of a trusted maintainer's PR. This PR looks like it solves the same underlying problem as unslothai/unsloth#4258 by @danielhanchen (trusted maintainer).
Canonical PR summary: This PR adds initial Apple Silicon MLX support through a new FastMLXModel API, MLX LoRA training helpers, MPS device detection, and a macOS mlx/mlx-lm dependency extra. It does this by gating CUDA/Triton-specific exports on MPS, adding MLX model loading/inference/training utilities, and wiring MLX tuner-based LoRA optimization. The auto-review is still running against this PR — reviewers will factor in the canonical above. If this PR is genuinely different, call out the delta in the review discussion so the maintainer can decide which to merge. |
Studio export Restore Tuple[bool, str, Optional[str]] contract on export_merged_model, export_base_model, export_gguf, and export_lora_adapter, populating output_path on successful local saves so routes/worker/CLI/frontend details.output_path is non-empty again. Lift the GPU save_method assignment out of the local-save branch so Hub-only merged exports (save_directory='', push_to_hub=True) no longer hit UnboundLocalError on the push branch. For MLX merged and base hub-only export, stage to a tempfile.TemporaryDirectory before push_to_hub_merged instead of passing save_directory=''. Source _IS_MLX from unsloth instead of recomputing the platform check (single source of truth, also enforces mlx-package availability). Studio MLX training/inference Pass token=hf_token into FastMLXModel.from_pretrained for gated/private models, matching the inference path. Strip hf_token and wandb_token from wandb.init(config=...) so secrets do not leak into the W&B run config. Replace load_from_disk(local_datasets[0]) with the existing UnslothTrainer._resolve_local_files / _loader_for_files helpers so uploaded JSON/JSONL/CSV/Parquet files train through the normal datasets loader (load_from_disk still used for HF save_to_disk directories). Make the dataset slice helper inclusive at the end and treat 0 as a real index instead of "unset", matching the GPU and embedding paths. Add a status_message -> message alias inside _send so the existing parent pump (training.py) renders MLX status updates instead of blanks. Forward min_p through generate_chat_response into _generate_text / _generate_vlm and into make_sampler / vlm_kwargs so the sampling control is no longer a no-op on MLX. Wrap unsloth_zoo.mlx_loader / mlx_trainer imports with a clearer ImportError pointing users at install.sh for Apple Silicon. Exit the MLX stop-polling thread on EOFError/OSError instead of busy-looping when the queue/pipe is permanently closed (one-line why-safe rationale inline). Studio frontend ParamsSection subscribes to platform deviceType via the Zustand hook so the gradient checkpointing dropdown re-renders after the async device fetch completes. Studio hardware get_gpu_utilization MLX branch now reads _read_apple_gpu_stats once and derives VRAM totals from psutil, removing the second ioreg subprocess per utilization poll. Unsloth core Restore the os.geteuid == 0 guard around the CUDA ldconfig recovery that was lost when GPU initialization moved into _gpu_init.py, plus the non-root manual-fix warning branch. Non-root CUDA users no longer shell out to ldconfig at import time. Load dataprep/raw_text via importlib so the MLX import path no longer pulls torch in through dataprep/__init__.py -> synthetic.py. FastVisionModel.from_pretrained overrides the inherited delegator only to inject text_only=False; this is an extension, not a duplication, and is needed so VLM checkpoint loads keep the vision tower. Wrap the MLX-branch unsloth_zoo import with a clearer ImportError.
…g guard tests/python/test_gpu_init_ldconfig_guard.py asserts the geteuid root check still wraps the ldconfig recovery and the non-root branch warns bnb users; AST + source-text inspection so the test runs without torch. tests/studio/test_export_output_path_contract.py covers the Tuple[bool, str, Optional[str]] return contract on every export method, the output_path assignment after successful local save, the Hub-only GPU save_method binding fix, the MLX hub-only TemporaryDirectory staging, and the single-source `_IS_MLX` import from unsloth. tests/studio/test_mlx_training_worker_behaviors.py covers token forwarding to FastMLXModel.from_pretrained, wandb config secret stripping, file-aware local dataset loading, status_message -> message aliasing, inclusive slice semantics, EOFError/OSError stop thread exit, and the friendly mlx_loader / mlx_trainer ImportError.
danielhanchen
added
auto-approved
|
Auto-review verdict: Approved Adds an Apple Silicon MLX training/inference path to Unsloth Studio (FastMLXModel + MLXTrainer routing for Mac ARM64), with frontend, hardware, and export plumbing. After the review fixes, the PR correctly enables Mac users to train, infer, and export models without falling back to CPU/chat-only mode while preserving the existing GPU/CUDA path. Reason: All 20 accepted findings (10 P1, 7 P2, 3 P3) addressed; no remaining real bugs; merge-back pushed cleanly |
for more information, see https://pre-commit.ci
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7d4e896b9c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| return ds | ||
|
|
||
| def _load_local(file_paths): | ||
| from core.training.trainer import UnslothTrainer |
There was a problem hiding this comment.
The new MLX path calls from core.training.trainer import UnslothTrainer inside _load_local, but core/training/trainer.py imports torch at module import time. This reintroduces a hard torch dependency in the supposed torch-free Apple/MLX flow, so local-file training (local_datasets / local_eval_datasets) can fail on MLX environments where torch is intentionally absent (for example no-torch installs) before dataset parsing even starts.
Useful? React with 👍 / 👎.
…r pre-pin Three memory-hardening fixes for Studio's MLX path: 1. Inference applies the same Metal caps as the trainer. load_model previously only called set_wired_limit(100% of recommended) with no upper memory_limit, leaving large VLM checkpoints unbounded during the loader allocation. Add _configure_memory_limits() that sets memory_limit to 85% of recommended and wired_limit to min(recommended, memory_limit) — matching MLXTrainer's defaults so behavior is the same whether the user trains or just runs inference. 2. unload_model releases pinned memory back to the OS — but only when the cache is empty. Without this, pinned wired bytes stayed allocated to MLX after the model was gone, starving other apps. The release is guarded on `not self.models` so unloading one of several cached models doesn't un-pin weights still in use. 3. Worker pre-cap is conservative instead of aggressive. The previous pre-pin set_wired_limit(100% of recommended) competed with MLXTrainer's later more conservative cap. Replace with the same 85%-memory / min(rec, memory) pair that the trainer applies later (idempotent re-apply). Bounds the model load + LoRA setup window without over-pinning.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 48ec7db71a
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if cmd_type == "generate": | ||
| cancel_event.clear() | ||
| _handle_generate(backend, cmd, resp_queue, cancel_event) | ||
| elif cmd_type == "load": | ||
| if backend.active_model_name: | ||
| backend.unload_model(backend.active_model_name) | ||
| _handle_load(backend, cmd, resp_queue) | ||
| elif cmd_type == "unload": | ||
| _handle_unload(backend, cmd, resp_queue) |
There was a problem hiding this comment.
The MLX command loop only routes generate, load, unload, etc., but omits generate_audio and generate_audio_input. The orchestrator still emits those command types from generate_audio_response / _generate_audio_input_inner, so on Apple-Silicon runs with an audio-capable active model the worker never replies and callers block until timeout (e.g., 120s in TTS). Add explicit handling (or immediate error responses) for audio command types in this MLX branch.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: df0cc6e0cb
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| class FastLanguageModel: | ||
| @staticmethod | ||
| def from_pretrained(*args, **kwargs): | ||
| return FastMLXModel.from_pretrained(*args, **kwargs) | ||
|
|
||
| @staticmethod | ||
| def get_peft_model(*args, **kwargs): | ||
| return FastMLXModel.get_peft_model(*args, **kwargs) | ||
|
|
||
| @staticmethod | ||
| def for_inference(*args, **kwargs): | ||
| return args[0] if args else None |
There was a problem hiding this comment.
On the MLX code path, FastLanguageModel no longer defines for_training, while GPU imports still expose that method via the existing public API. Any Apple-Silicon script that calls FastLanguageModel.for_training(model, ...) (a common pattern in Unsloth training flows) will now raise AttributeError before training starts. Please add a no-op shim (like the new FastVisionModel.for_training) to keep the API surface compatible across backends.
Useful? React with 👍 / 👎.
Two gates drive every MLX-vs-CUDA dispatch decision in Studio:
1. unsloth._IS_MLX in unsloth/__init__.py — evaluated once at import
time, read by Studio worker code to choose the GPU vs MLX trainer
and inference paths. Defined as
Darwin AND arm64 AND find_spec("mlx") is not None.
2. utils.hardware.detect_hardware() — runtime probe with priority
CUDA > XPU > MLX > CPU. The MLX branch is reached only when both
CUDA and XPU are unavailable and the host is Apple Silicon and
mlx is importable.
Neither gate had a direct test. Adds tests/studio/test_is_mlx_dispatch_gate.py
with six tests:
test_is_mlx_gate_uses_three_required_predicates
AST-walks unsloth/__init__.py and asserts the _IS_MLX assignment
is a BoolOp(And) of platform.system()=="Darwin",
platform.machine()=="arm64", and find_spec("mlx") is not None.
Catches accidental rewrites that drop a predicate.
test_is_mlx_gate_true_on_apple_silicon_with_mlx_present
Spoofs platform to Darwin/arm64, injects a fake mlx module so
find_spec returns a real ModuleSpec, re-evaluates the gate
expression. Verifies it flips True under the exact conditions
Studio expects.
test_is_mlx_gate_false_when_mlx_missing
Spoofs Apple Silicon but with mlx absent. Verifies the gate stays
False (so a Mac without mlx installed does not pretend to have
MLX support).
test_is_mlx_gate_false_on_non_apple_silicon
Canary on the actual Linux+CUDA / AMD / Intel test host: the gate
must remain False regardless of whether mlx happens to be
importable. Protects existing GPU users from accidental MLX
hijack when MLX support evolves.
test_detect_hardware_picks_mlx_when_only_apple_silicon_available
Forces torch.cuda and torch.xpu off, spoofs Apple Silicon, injects
fake mlx and mlx.core. detect_hardware() must return DeviceType.MLX.
test_detect_hardware_picks_cuda_on_real_host
Canary: on a real CUDA host detect_hardware() must return
DeviceType.CUDA. Protects against the MLX branch shadowing CUDA
dispatch on NVIDIA / AMD ROCm hosts.
Uses the same monkeypatch.setitem(sys.modules, ...) fake-mlx pattern as
the existing test_mlx_inference_backend.py — no new test infrastructure,
no real mlx install required.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 64683731ef
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # MLX LoRA adapter toggling not yet supported — generate normally | ||
| yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs) |
There was a problem hiding this comment.
generate_with_adapter_control currently ignores the use_adapter argument and always calls normal generation, so compare-mode requests like use_adapter=False or selecting a specific adapter produce LoRA-on output instead of the requested base/adapter state. This breaks API parity with the GPU backend and returns incorrect comparison results whenever clients use adapter toggling (e.g., A/B evaluation flows routed through orchestrator.generate_with_adapter_control). At minimum this path should either enforce the requested adapter state or return an explicit unsupported error instead of silently ignoring it.
Useful? React with 👍 / 👎.
Four Studio MLX test files shipped without an SPDX-License-Identifier: studio/backend/tests/test_mlx_training_worker_config.py tests/studio/test_mlx_training_worker_behaviors.py tests/studio/test_export_output_path_contract.py tests/studio/test_is_mlx_dispatch_gate.py They sit in or alongside studio/backend/, which is governed by studio/LICENSE.AGPL-3.0, and exercise AGPL Studio code. Add the same "# SPDX-License-Identifier: AGPL-3.0-only" header that's already on test_mlx_inference_backend.py so the license declaration matches the code under test rather than defaulting to the repo-root Apache-2.0.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cad3597838
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| huggingface = [ | ||
| "unsloth[huggingfacenotorch]", | ||
| "unsloth_zoo>=2026.5.1", | ||
| "unsloth_zoo>=2026.4.8", |
There was a problem hiding this comment.
This dependency now allows unsloth_zoo==2026.4.8, but the new Apple-Silicon import path immediately imports unsloth_zoo.mlx_trainer and unsloth_zoo.mlx_loader (unsloth/__init__.py, and the MLX training worker does the same). A fresh install can therefore resolve an older zoo build that satisfies this spec but lacks those MLX entry points, causing import unsloth or MLX training/inference startup to fail before users can recover. Keep the lower bound at the first zoo release that ships these modules (and mirror it in colab-new).
Useful? React with 👍 / 👎.
The _IS_MLX block at the top of unsloth/__init__.py already catches the missing-package case with a friendly install hint, but the follow-up "from unsloth_zoo.mlx_trainer import ..." and "from unsloth_zoo.mlx_loader import ..." lines run unguarded. An Apple Silicon user who has unsloth-zoo installed but on an older version (e.g. the current PyPI release, before the MLX modules ship) sees a raw ImportError on the submodule rather than the hint that points at install.sh. Wrap the two submodule imports in the same try/except shape so the friendly install message fires whether the package is missing entirely or just predates the MLX submodules. No-op once both packages release together; smooths the transitional window where unsloth/main has merged but unsloth-zoo on PyPI has not.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 494059746e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| info = format_and_template_dataset( | ||
| dataset, | ||
| model_name = model_name, | ||
| tokenizer = tokenizer, | ||
| is_vlm = False, | ||
| format_type = format_type, | ||
| dataset_name = hf_dataset or "local", | ||
| ) |
There was a problem hiding this comment.
When Apple-Silicon users train on an uploaded/custom dataset that relies on custom_format_mapping from the Studio UI, this MLX formatting call drops that config, so the formatter cannot map non-standard columns and will either fail detection or train on the wrong fields. The existing GPU path forwards the same mapping through trainer.load_and_format_dataset(...), so the MLX branch should pass custom_format_mapping=config.get("custom_format_mapping") for both train and eval/VLM formatting calls.
Useful? React with 👍 / 👎.
| output_dir = config.get("output_dir", "") | ||
| if not output_dir: | ||
| output_dir = f"{model_name.replace('/', '_')}_{int(time.time())}" |
There was a problem hiding this comment.
When a user resumes a previous Studio training run on Apple Silicon, resume_from_checkpoint is still sent in the config, but this MLX branch ignores it and falls back to a fresh timestamped output directory before calling trainer.train() without any resume argument. That makes the Resume action silently restart training from step 0 instead of continuing the checkpoint, unlike the GPU branches that derive output_dir from the checkpoint and pass resume_from_checkpoint into training.
Useful? React with 👍 / 👎.
4 participants
Summary
Routes Studio's training pipeline through
MLXTraineron Apple Silicon, replacing the torch /SFTTrainerpath that doesn't run on Mac. Same UI, same one-click flow, same export. Studio now trains models on M1-M5 Macs with the memory wins fromunsloth-zoo's MLX integration.Depends on
unslothai/unsloth-zoo#XXX(Apple Silicon training PR).What's included
Backend (
studio/backend/)utils/hardware/hardware.py: detect MLX on Apple Silicon, setCHAT_ONLY = False.core/training/worker.py: MLX fast-path that bypasses torch / SFTTrainer entirely. BuildsFastMLXModel+MLXTrainer, hooks progress / loss / memory events into theexisting event_queue. Supports LoRA + full FT, gradient checkpointing, CCE,
train_on_responses_only, and thefinetune_language / attention / mlp / visionflags withauto-imply guardrails (e.g. picking attention without picking a scope auto-implies language). Elementwise grad clip is on by default.
core/training/training.py: skip CUDA / GPU validation on MLX.core/inference/mlx_inference.py: text + VLM streaming inference. Wirestop_k,top_p,repetition_penalty,temperaturethrough to mlx-lm / mlx-vlm samplers (fixessilent VLM temperature drop).
core/export/export.py: maps Studio's format dropdown tosave_method(LoRA-only / merged 16-bit / merged 4-bit / GGUF variants); passessave_directorytopush_to_hub_merged; private/public toggle.Frontend (
studio/frontend/)src/config/env.ts: drop the Mac fallback that hard-codedchat_only = true; respect backendchat_only.What's preserved
worker.pywhen the device is MLX.