MLX training support for Studio on Apple Silicon (#5340)

* mlx fixes

* Fix studio integration, local dataset files, chat templates without the torch gpu imports

* pass grad norm in mlx worker

* fix(studio): pass MLX grad clipping settings

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* mlx: update grad value

* fix(mlx): address ci and clipping review

* fix backward compatibility and CI tests

* unsloth local is mlx function

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* dont reference runtime

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio mlx: hardcode value clipping, drop max_grad_value from frontend

Simplifies the MLX grad-clipping plumbing now that we are standardising on
elementwise value clipping at [-5, 5] for the compiled MLX path and norm
clipping disabled. The MLX worker no longer reads max_grad_norm /
max_grad_value from the request; both are pinned in one place. Frontend
stops sending the field at all, and the TypeScript request type drops it
to match. Non-MLX (CUDA/AMD/Intel) is untouched and continues to pick up
HF TrainingArguments' default max_grad_norm = 1.0.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>

studio/backend/core/inference/mlx_inference.py

@@ -116,11 +116,11 @@ class MLXInferenceBackend:
         )
 
         try:
-            from unsloth_zoo.mlx_loader import FastMLXModel
+            from unsloth_zoo.mlx.loader import FastMLXModel
         except ImportError as e:
             raise ImportError(
                 "Unsloth: MLX inference requires unsloth-zoo with the MLX modules "
-                "(unsloth_zoo.mlx_loader). Reinstall via install.sh on Apple Silicon."
+                "(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon."
             ) from e
 
         model, tokenizer_or_processor = FastMLXModel.from_pretrained(

studio/backend/core/training/training.py

@@ -214,6 +214,8 @@ class TrainingBackend:
             "max_steps": kwargs.get("max_steps", 0),
             "save_steps": kwargs.get("save_steps", 0),
             "weight_decay": kwargs.get("weight_decay", 0.001),
+            "max_grad_norm": kwargs.get("max_grad_norm", 0.0),
+            "max_grad_value": kwargs.get("max_grad_value"),
             "random_seed": kwargs.get("random_seed", 3407),
             "packing": kwargs.get("packing", False),
             "optim": kwargs.get("optim", "adamw_8bit"),

studio/backend/core/training/worker.py

@@ -417,6 +417,55 @@ def _normalize_mlx_studio_scheduler(value):
     return raw
 
 
+def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]:
+    """Resolve Studio local dataset uploads without importing the GPU trainer."""
+    from utils.paths import resolve_dataset_path
+
+    all_files: list[str] = []
+    for dataset_file in file_paths or []:
+        file_path = (
+            dataset_file
+            if os.path.isabs(dataset_file)
+            else str(resolve_dataset_path(dataset_file))
+        )
+        file_path_obj = Path(file_path)
+
+        if file_path_obj.is_dir():
+            parquet_dir = (
+                file_path_obj / "parquet-files"
+                if (file_path_obj / "parquet-files").exists()
+                else file_path_obj
+            )
+            parquet_files = sorted(parquet_dir.glob("*.parquet"))
+            if parquet_files:
+                all_files.extend(str(p) for p in parquet_files)
+                continue
+
+            candidates: list[Path] = []
+            for ext in (".json", ".jsonl", ".csv", ".parquet"):
+                candidates.extend(sorted(file_path_obj.glob(f"*{ext}")))
+            if candidates:
+                all_files.extend(str(c) for c in candidates)
+                continue
+
+            raise ValueError(f"No supported data files in directory: {file_path_obj}")
+
+        all_files.append(str(file_path_obj))
+
+    return all_files
+
+
+def _mlx_local_dataset_loader_for_files(files: list[str]) -> str:
+    first_ext = Path(files[0]).suffix.lower()
+    if first_ext in (".json", ".jsonl"):
+        return "json"
+    if first_ext == ".csv":
+        return "csv"
+    if first_ext == ".parquet":
+        return "parquet"
+    raise ValueError(f"Unsupported dataset format: {files[0]}")
+
+
 def _run_mlx_training(event_queue, stop_queue, config):
     """Self-contained MLX training path for Apple Silicon.
 
@@ -442,8 +491,8 @@ def _run_mlx_training(event_queue, stop_queue, config):
     import mlx.core as mx
 
     try:
-        from unsloth_zoo.mlx_loader import FastMLXModel
-        from unsloth_zoo.mlx_trainer import (
+        from unsloth_zoo.mlx.loader import FastMLXModel
+        from unsloth_zoo.mlx.trainer import (
             MLXTrainer,
             MLXTrainingConfig,
             train_on_responses_only,
@@ -451,7 +500,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
     except ImportError as e:
         raise ImportError(
             "Unsloth: MLX training requires unsloth-zoo with the MLX modules "
-            "(unsloth_zoo.mlx_loader / unsloth_zoo.mlx_trainer). Reinstall via "
+            "(unsloth_zoo.mlx.loader / unsloth_zoo.mlx.trainer). Reinstall via "
             "install.sh on Apple Silicon."
         ) from e
     from datasets import load_dataset
@@ -572,7 +621,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
         return ds
 
     def _load_local(file_paths):
-        from core.training.trainer import UnslothTrainer
         from datasets import load_from_disk
 
         if len(file_paths) == 1:
@@ -581,10 +629,10 @@ def _run_mlx_training(event_queue, stop_queue, config):
                 (p / "dataset_info.json").exists() or (p / "state.json").exists()
             ):
                 return load_from_disk(str(p))
-        all_files = UnslothTrainer._resolve_local_files(file_paths)
+        all_files = _resolve_mlx_local_dataset_files(file_paths)
         if not all_files:
             raise ValueError("No local dataset files found")
-        loader = UnslothTrainer._loader_for_files(all_files)
+        loader = _mlx_local_dataset_loader_for_files(all_files)
         return load_dataset(loader, data_files = all_files, split = "train")
 
     if hf_dataset:
@@ -718,6 +766,10 @@ def _run_mlx_training(event_queue, stop_queue, config):
     else:
         eval_steps_val = int(eval_steps_val)
 
+    # MLX: value-clip grads to [-5, 5]; norm clipping disabled for compile-friendliness.
+    max_grad_norm = 0.0
+    max_grad_value = 5.0  # TODO: expose MLX grad-clip in Studio UI for power users
+
     trainer = MLXTrainer(
         model = model,
         tokenizer = tokenizer,
@@ -732,6 +784,8 @@ def _run_mlx_training(event_queue, stop_queue, config):
             lr_scheduler_type = lr_scheduler_type,
             optim = optim_name,
             weight_decay = float(config.get("weight_decay", 0.001) or 0.001),
+            max_grad_norm = max_grad_norm,
+            max_grad_value = max_grad_value,
             logging_steps = 1,
             max_seq_length = max_seq_length,
             seed = config.get("random_seed", 3407),
@@ -820,7 +874,17 @@ def _run_mlx_training(event_queue, stop_queue, config):
     # ── 9. Real-time progress callback ──
     _send("status", status_message = f"Training {model_name}...")
 
-    def _on_step(step, total, loss, lr, tok_s, peak_gb, elapsed, num_tokens):
+    def _on_step(
+        step,
+        total,
+        loss,
+        lr,
+        tok_s,
+        peak_gb,
+        elapsed,
+        num_tokens,
+        grad_norm = None,
+    ):
         eta = (elapsed / step * (total - step)) if step > 0 else 0
         _send(
             "progress",
@@ -831,7 +895,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
             total_steps = total,
             elapsed_seconds = elapsed,
             eta_seconds = max(0, eta),
-            grad_norm = None,
+            grad_norm = grad_norm,
             num_tokens = num_tokens,
             eval_loss = None,
             status_message = None,
@@ -846,6 +910,11 @@ def _run_mlx_training(event_queue, stop_queue, config):
                         "train/tokens_per_sec": tok_s,
                         "train/peak_gb": peak_gb,
                         "train/num_tokens": num_tokens,
+                        **(
+                            {"train/grad_norm": grad_norm}
+                            if grad_norm is not None
+                            else {}
+                        ),
                     },
                     step = step,
                 )
@@ -857,6 +926,8 @@ def _run_mlx_training(event_queue, stop_queue, config):
                 tb_writer.add_scalar("train/learning_rate", lr, step)
                 tb_writer.add_scalar("train/tokens_per_sec", tok_s, step)
                 tb_writer.add_scalar("train/peak_gb", peak_gb, step)
+                if grad_norm is not None:
+                    tb_writer.add_scalar("train/grad_norm", grad_norm, step)
             except Exception:
                 pass
 

studio/backend/models/training.py

@@ -262,6 +262,19 @@ class TrainingStartRequest(BaseModel):
     max_steps: Optional[int] = Field(None, description = "Maximum training steps")
     save_steps: int = Field(100, description = "Steps between checkpoints")
     weight_decay: float = Field(0.001, description = "Weight decay")
+    max_grad_norm: float = Field(
+        0.0,
+        ge = 0,
+        description = "Global gradient norm clipping threshold. Set 0 to disable.",
+    )
+    max_grad_value: Optional[float] = Field(
+        None,
+        ge = 0,
+        description = (
+            "Elementwise gradient value clipping threshold. Set 0 to disable. "
+            "If omitted, MLX defaults to 1 unless max_grad_norm is set."
+        ),
+    )
     random_seed: int = Field(42, description = "Random seed")
     packing: bool = Field(False, description = "Enable sequence packing")
     optim: str = Field("adamw_8bit", description = "Optimizer")

studio/backend/routes/training.py

@@ -215,6 +215,8 @@ async def start_training(
             "max_steps": request.max_steps,
             "save_steps": request.save_steps,
             "weight_decay": request.weight_decay,
+            "max_grad_norm": request.max_grad_norm,
+            "max_grad_value": request.max_grad_value,
             "random_seed": request.random_seed,
             "packing": request.packing,
             "optim": request.optim,

studio/backend/tests/test_mlx_inference_backend.py

@@ -56,11 +56,14 @@ def _install_fake_fast_mlx(monkeypatch, calls):
             return _DummyModel(), _DummyTokenizer()
 
     unsloth_zoo_pkg = types.ModuleType("unsloth_zoo")
-    mlx_loader = types.ModuleType("unsloth_zoo.mlx_loader")
+    mlx_pkg = types.ModuleType("unsloth_zoo.mlx")
+    mlx_loader = types.ModuleType("unsloth_zoo.mlx.loader")
     mlx_loader.FastMLXModel = _FastMLXModel
-    unsloth_zoo_pkg.mlx_loader = mlx_loader
+    unsloth_zoo_pkg.mlx = mlx_pkg
+    mlx_pkg.loader = mlx_loader
     monkeypatch.setitem(sys.modules, "unsloth_zoo", unsloth_zoo_pkg)
-    monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx_loader", mlx_loader)
+    monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx", mlx_pkg)
+    monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.loader", mlx_loader)
 
 
 def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch):

studio/backend/tests/test_training_raw_support.py

@@ -70,6 +70,50 @@ class TestTrainingRawSupport(unittest.TestCase):
         self.assertTrue(config["load_in_4bit"])
         self.assertEqual(config["embedding_learning_rate"], 1e-5)
 
+    def test_training_backend_forwards_grad_clipping_controls(self):
+        backend = TrainingBackend()
+
+        class DummyProcess:
+            pid = 12345
+
+            def start(self):
+                return None
+
+        class DummyThread:
+            def start(self):
+                return None
+
+        dummy_queue = object()
+
+        with (
+            patch(
+                "core.training.training.prepare_gpu_selection",
+                return_value = ([0], {"selection_mode": "auto"}),
+            ),
+            patch(
+                "core.training.training._CTX.Queue",
+                side_effect = [dummy_queue, dummy_queue],
+            ),
+            patch(
+                "core.training.training._CTX.Process", return_value = DummyProcess()
+            ) as mock_process,
+            patch(
+                "core.training.training.threading.Thread",
+                return_value = DummyThread(),
+            ),
+        ):
+            backend.start_training(
+                job_id = "test-grad-clip",
+                model_name = "unsloth/test",
+                training_type = "LoRA/QLoRA",
+                max_grad_norm = 0.7,
+                max_grad_value = 0.0,
+            )
+
+        config = mock_process.call_args.kwargs["kwargs"]["config"]
+        self.assertEqual(config["max_grad_norm"], 0.7)
+        self.assertEqual(config["max_grad_value"], 0.0)
+
     def test_training_route_forwards_embedding_learning_rate(self):
         training_route = _load_route_module(
             "training_route_module_raw_support",

studio/backend/utils/datasets/chat_templates.py

@@ -28,6 +28,23 @@ DEFAULT_ALPACA_TEMPLATE = """Below is an instruction that describes a task, pair
 {}"""
 
 
+def _is_mlx_runtime() -> bool:
+    try:
+        from unsloth_zoo.mlx import is_mlx_available
+    except ImportError:
+        return False
+    return is_mlx_available()
+
+
+def _chat_template_kwargs() -> dict:
+    if not _is_mlx_runtime():
+        return {}
+    return {
+        "patch_saving": False,
+        "use_zoo_tokenizer_patch": True,
+    }
+
+
 def get_tokenizer_chat_template(tokenizer, model_name):
     """
     Gets appropriate chat template for tokenizer based on model.
@@ -60,6 +77,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
             tokenizer = get_chat_template(
                 tokenizer,
                 chat_template = matched_template,
+                **_chat_template_kwargs(),
             )
         except Exception as e:
             logger.info(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}")
@@ -79,6 +97,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
                 tokenizer = get_chat_template(
                     tokenizer,
                     chat_template = "chatml",
+                    **_chat_template_kwargs(),
                 )
             except Exception as e:
                 logger.info(f"⚠️ Failed to apply default ChatML template: {e}")
@@ -255,7 +274,11 @@ def apply_chat_template_to_dataset(
         if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template):
             try:
                 from unsloth.chat_templates import get_chat_template
-                tokenizer = get_chat_template(tokenizer, chat_template = "alpaca")
+                tokenizer = get_chat_template(
+                    tokenizer,
+                    chat_template = "alpaca",
+                    **_chat_template_kwargs(),
+                )
                 logger.info(f"📝 Set alpaca chat template on tokenizer for model saving")
             except Exception as e:
                 logger.info(f"⚠️ Could not set alpaca template on tokenizer: {e}")

studio/frontend/src/features/training/api/mappers.ts

@@ -83,6 +83,7 @@ export function buildTrainingStartPayload(
     save_steps: config.saveSteps,
     eval_steps: config.evalSteps,
     weight_decay: config.weightDecay,
+    max_grad_norm: 0.0,
     random_seed: config.randomSeed,
     packing: isEmbedding ? false : config.packing,
     optim: config.optimizerType,

studio/frontend/src/features/training/types/api.ts

@@ -31,6 +31,7 @@ export interface TrainingStartRequest {
   save_steps: number;
   eval_steps: number;
   weight_decay: number;
+  max_grad_norm: number;
   random_seed: number;
   packing: boolean;
   optim: string;

tests/studio/run_real_mlx_smoke.py

@@ -211,8 +211,8 @@ def cmd_train(args) -> int:
     workdir.mkdir(parents = True, exist_ok = True)
 
     import mlx.core as mx
-    from unsloth_zoo.mlx_loader import FastMLXModel
-    from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig
+    from unsloth_zoo.mlx.loader import FastMLXModel
+    from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig
 
     hf_token = os.environ.get("HF_TOKEN") or None
 
@@ -440,7 +440,7 @@ def cmd_reload(args) -> int:
         return _reload_gguf(save_dir, metrics)
 
     import mlx.core as mx
-    from unsloth_zoo.mlx_loader import FastMLXModel
+    from unsloth_zoo.mlx.loader import FastMLXModel
     from mlx_lm import generate
 
     hf_token = os.environ.get("HF_TOKEN") or None

tests/studio/test_is_mlx_dispatch_gate.py

@@ -7,8 +7,9 @@ Two gates drive every dispatch decision in Studio's MLX path:
 
   1. ``unsloth._IS_MLX`` at the top of ``unsloth/__init__.py`` -- evaluated
      once at import time and read by Studio worker code to choose between
-     the GPU and MLX trainer / inference / export paths. Defined as
-     ``Darwin AND arm64 AND find_spec("mlx") is not None``.
+     the GPU and MLX trainer / inference / export paths. It delegates to
+     the shared zoo MLX runtime gate, with a local import barrier while the
+     paired unsloth-zoo runtime rollout is in flight.
 
   2. ``utils.hardware.detect_hardware()`` -- runtime probe in the Studio
      backend. Priority order: CUDA -> XPU -> MLX -> CPU. The MLX branch is
@@ -18,8 +19,8 @@ Two gates drive every dispatch decision in Studio's MLX path:
 These gates are the canaries for "MLX support accidentally hijacks
 CUDA/AMD/Intel users". The tests here:
 
-  * verify the source-level structure of the ``_IS_MLX`` expression so an
-    accidental rewrite (e.g. dropping the ``arm64`` check) is caught,
+  * verify the source-level structure of the ``_IS_MLX`` helper so an
+    accidental rewrite importing zoo before the local MLX precheck is caught,
   * exercise the runtime gate logic under a spoofed Darwin+arm64 platform
     with a fake ``mlx`` module in ``sys.modules`` to confirm both gates
     flip True together,
@@ -64,20 +65,36 @@ def test_is_mlx_gate_uses_three_required_predicates():
             target = node.value
             break
     assert target is not None, "_IS_MLX assignment not found in unsloth/__init__.py"
-    assert isinstance(target, ast.BoolOp) and isinstance(
-        target.op, ast.And
-    ), "_IS_MLX must be a BoolOp(And) of platform + mlx checks"
-
+    assert isinstance(target, ast.Call), "_IS_MLX must call the shared MLX helper"
     expr_src = ast.unparse(target)
     assert (
-        "platform.system()" in expr_src and "Darwin" in expr_src
-    ), "_IS_MLX must check platform.system() == 'Darwin'"
+        expr_src == "_is_mlx_available()"
+    ), "_IS_MLX must delegate to the shared MLX runtime gate"
+
+    helper = None
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name == "_is_mlx_available":
+            helper = node
+            break
+    assert helper is not None, "_is_mlx_available helper not found"
+
+    helper_src = ast.unparse(helper)
     assert (
-        "platform.machine()" in expr_src and "arm64" in expr_src
-    ), "_IS_MLX must check platform.machine() == 'arm64'"
+        "platform.system()" in helper_src
+        and "'Darwin'" in helper_src
+        and "platform.machine()" in helper_src
+        and "'arm64'" in helper_src
+        and "find_spec" in helper_src
+        and "'mlx'" in helper_src
+        and "from unsloth_zoo.mlx import is_mlx_available" in helper_src
+    ), "_IS_MLX helper must precheck local MLX predicates before importing zoo"
     assert (
-        "find_spec" in expr_src and "'mlx'" in expr_src
-    ), "_IS_MLX must check importlib.util.find_spec('mlx')"
+        "from unsloth_zoo.mlx import is_mlx_available" in helper_src
+        and "return is_mlx_available()" in helper_src
+    ), "_IS_MLX helper must delegate final detection to the shared zoo MLX runtime gate"
+    assert helper_src.index("UNSLOTH_FORCE_GPU_PATH") < helper_src.index(
+        "from unsloth_zoo.mlx import is_mlx_available"
+    ), "_IS_MLX helper must run the local MLX precheck before importing zoo"
 
 
 # ---------------------------------------------------------------------------
@@ -87,13 +104,14 @@ def test_is_mlx_gate_uses_three_required_predicates():
 # ---------------------------------------------------------------------------
 
 
-def _evaluate_is_mlx_gate(platform_module, importlib_util):
-    """Re-evaluate the _IS_MLX expression using injected dependencies.
+def _evaluate_is_mlx_precheck(platform_module, importlib_util, os_module):
+    """Re-evaluate the local _is_mlx_available precheck using injected dependencies.
 
-    Mirrors the assignment in unsloth/__init__.py exactly.
+    Mirrors only the cheap import barrier before unsloth imports unsloth_zoo.
     """
     return (
-        platform_module.system() == "Darwin"
+        os_module.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1"
+        and platform_module.system() == "Darwin"
         and platform_module.machine() == "arm64"
         and importlib_util.find_spec("mlx") is not None
     )
@@ -112,7 +130,9 @@ def test_is_mlx_gate_true_on_apple_silicon_with_mlx_present(monkeypatch):
     monkeypatch.setattr(platform, "system", lambda: "Darwin")
     monkeypatch.setattr(platform, "machine", lambda: "arm64")
 
-    assert _evaluate_is_mlx_gate(platform, importlib.util) is True
+    import os
+
+    assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is True
 
 
 def test_is_mlx_gate_false_when_mlx_missing(monkeypatch):
@@ -133,7 +153,9 @@ def test_is_mlx_gate_false_when_mlx_missing(monkeypatch):
 
     monkeypatch.setattr(importlib.util, "find_spec", _no_mlx)
 
-    assert _evaluate_is_mlx_gate(platform, importlib.util) is False
+    import os
+
+    assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is False
 
 
 def test_is_mlx_gate_false_on_non_apple_silicon():
@@ -147,7 +169,9 @@ def test_is_mlx_gate_false_on_non_apple_silicon():
 
         pytest.skip("Test host is Apple Silicon; CUDA-side canary doesn't apply.")
 
-    assert _evaluate_is_mlx_gate(platform, importlib.util) is False
+    import os
+
+    assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is False
 
 
 # ---------------------------------------------------------------------------

tests/studio/test_mlx_training_worker_behaviors.py

@@ -46,8 +46,8 @@ def test_wandb_init_strips_secret_keys():
 
 def test_local_dataset_loader_uses_load_dataset_path():
     src = WORKER.read_text()
-    assert "_resolve_local_files" in src
-    assert "_loader_for_files" in src
+    assert "_resolve_mlx_local_dataset_files" in src
+    assert "_mlx_local_dataset_loader_for_files" in src
     assert "data_files = all_files" in src or "data_files=all_files" in src
 
 
@@ -84,7 +84,7 @@ def test_poll_stop_returns_on_broken_pipe():
 
 def test_unsloth_zoo_mlx_imports_have_friendly_error():
     src = WORKER.read_text()
-    assert "from unsloth_zoo.mlx_loader import FastMLXModel" in src
-    assert "from unsloth_zoo.mlx_trainer import" in src
+    assert "from unsloth_zoo.mlx.loader import FastMLXModel" in src
+    assert "from unsloth_zoo.mlx.trainer import" in src
     assert "raise ImportError" in src
     assert "install.sh" in src

unsloth/__init__.py

@@ -12,16 +12,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os, platform, importlib.util
+import os, importlib.util, platform
 
 os.environ["UNSLOTH_IS_PRESENT"] = "1"
 
+
+def _is_mlx_available():
+    # Transitional import barrier: while the paired unsloth-zoo MLX runtime
+    # rollout is in flight, keep non-Apple-Silicon imports from touching
+    # unsloth_zoo here. After both PRs are released together and
+    # unsloth_zoo.mlx is guaranteed to be import-safe on GPU hosts,
+    # this helper can collapse back to the centralized zoo runtime call below.
+    if (
+        os.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") == "1"
+        or platform.system() != "Darwin"
+        or platform.machine() != "arm64"
+        or importlib.util.find_spec("mlx") is None
+    ):
+        return False
+    try:
+        from unsloth_zoo.mlx import is_mlx_available
+    except ImportError:
+        return False
+    return is_mlx_available()
+
+
 # Detect Apple Silicon + MLX before any torch/numpy imports
-_IS_MLX = (
-    platform.system() == "Darwin"
-    and platform.machine() == "arm64"
-    and importlib.util.find_spec("mlx") is not None
-)
+_IS_MLX = _is_mlx_available()
 
 if _IS_MLX:
     try:
@@ -31,18 +48,18 @@ if _IS_MLX:
             "Unsloth: MLX support requires `unsloth-zoo` with MLX modules. "
             "Reinstall with `pip install unsloth-zoo` or rerun install.sh."
         ) from _e
-    # The mlx_trainer / mlx_loader submodules ship with unsloth-zoo's MLX
+    # The mlx.trainer / mlx.loader submodules ship with unsloth-zoo's MLX
     # support. An older installed unsloth-zoo (e.g. from PyPI before the
     # MLX release lands) will satisfy `import unsloth_zoo` but be missing
     # these submodules. Surface the same friendly install hint instead of
     # a raw ImportError on the submodule path.
     try:
-        from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig
-        from unsloth_zoo.mlx_loader import FastMLXModel
+        from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig
+        from unsloth_zoo.mlx.loader import FastMLXModel
     except ImportError as _e:
         raise ImportError(
             "Unsloth: MLX support requires an unsloth-zoo build that includes "
-            "`unsloth_zoo.mlx_trainer` and `unsloth_zoo.mlx_loader`. Upgrade with "
+            "`unsloth_zoo.mlx.trainer` and `unsloth_zoo.mlx.loader`. Upgrade with "
             "`pip install -U unsloth-zoo` or rerun install.sh."
         ) from _e
 

unsloth/chat_templates.py

@@ -30,11 +30,9 @@ __all__ = [
 from transformers import StoppingCriteria, StoppingCriteriaList
 from torch import LongTensor, FloatTensor
 from transformers.models.llama.modeling_llama import logger
-from .save import patch_saving_functions
 import os
 import shutil
 from .tokenizer_utils import *
-from .models._utils import patch_tokenizer
 import re
 from .ollama_template_mappers import OLLAMA_TEMPLATES
 from unsloth_zoo.dataset_utils import (
@@ -1844,6 +1842,8 @@ def get_chat_template(
     mapping = {"role" : "role", "content" : "content", "user" : "user", "assistant" : "assistant"},
     map_eos_token = True,
     system_message = None,
+    patch_saving = True,
+    use_zoo_tokenizer_patch = False,
 ):
     assert(type(map_eos_token) is bool)
     old_tokenizer = tokenizer
@@ -2026,6 +2026,12 @@ def get_chat_template(
         .replace("'user'",      "'" + mapping["user"]      + "'")\
         .replace("'assistant'", "'" + mapping["assistant"] + "'")
 
+    if use_zoo_tokenizer_patch:
+        # Studio MLX avoids the model-utils tokenizer wrapper because that
+        # import path pulls in Torch/GPU-specific modules before MLX training.
+        from unsloth_zoo.tokenizer_utils import patch_tokenizer
+    else:
+        from .models._utils import patch_tokenizer
     _, tokenizer = patch_tokenizer(model = None, tokenizer = tokenizer)
     tokenizer.padding_side = old_padding_side
 
@@ -2059,7 +2065,9 @@ def get_chat_template(
     # stopping_criteria = create_stopping_criteria(tokenizer, stop_word)
 
     # Patch saving functions
-    tokenizer = patch_saving_functions(tokenizer)
+    if patch_saving:
+        from .save import patch_saving_functions
+        tokenizer = patch_saving_functions(tokenizer)
 
     # Add Ollama
     tokenizer._ollama_modelfile = ollama_modelfile

unsloth/device_type.py

@@ -20,21 +20,40 @@ __all__ = [
     "DEVICE_COUNT",
     "ALLOW_PREQUANTIZED_MODELS",
     "ALLOW_BITSANDBYTES",
+    "is_mlx_available",
 ]
 
-import torch
 import functools
 import inspect
+import os
 from unsloth_zoo.utils import Version
 
 
+def is_mlx_available():
+    try:
+        from unsloth_zoo.mlx import is_mlx_available as _is_mlx_available
+    except ImportError:
+        return False
+    return _is_mlx_available()
+
+
+_IS_MLX = is_mlx_available()
+
+if not _IS_MLX:
+    import torch
+
+
 @functools.cache
 def is_hip():
+    if _IS_MLX:
+        return False
     return bool(getattr(getattr(torch, "version", None), "hip", None))
 
 
 @functools.cache
 def get_device_type():
+    if _IS_MLX:
+        return "mlx"
     if hasattr(torch, "cuda") and torch.cuda.is_available():
         if is_hip():
             return "hip"
@@ -64,6 +83,8 @@ DEVICE_TYPE: str = get_device_type()
 DEVICE_TYPE_TORCH = DEVICE_TYPE
 if DEVICE_TYPE_TORCH == "hip":
     DEVICE_TYPE_TORCH = "cuda"
+elif DEVICE_TYPE_TORCH == "mlx":
+    DEVICE_TYPE_TORCH = "mps"
 
 
 @functools.cache

unsloth/kernels/utils.py

@@ -160,6 +160,10 @@ else:
 # INTEL GPU Specific Logic
 if DEVICE_TYPE == "xpu":
     _gpu_getCurrentRawStream = torch._C._xpu_getCurrentRawStream
+elif DEVICE_TYPE == "mlx":
+
+    def _gpu_getCurrentRawStream(_index = 0):
+        return 0
 # NVIDIA GPU Default Logic
 elif hasattr(torch._C, "_cuda_getCurrentRawStream"):
     _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream
@@ -206,6 +210,11 @@ if DEVICE_TYPE == "xpu":
         XPU_STREAMS = ()
         WEIGHT_BUFFERS = []
         ABSMAX_BUFFERS = []
+elif DEVICE_TYPE == "mlx":
+    CUDA_STREAMS = ()
+    XPU_STREAMS = ()
+    WEIGHT_BUFFERS = []
+    ABSMAX_BUFFERS = []
 else:
     # NVIDIA GPU Default Logic
     if DEVICE_COUNT > 0: