pass grad norm in mlx worker

studio/backend/core/inference/mlx_inference.py

@@ -94,11 +94,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/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:
@@ -732,6 +780,7 @@ 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 = float(config.get("max_grad_norm", 0.0) or 0.0),
             logging_steps = 1,
             max_seq_length = max_seq_length,
             seed = config.get("random_seed", 3407),
@@ -820,7 +869,10 @@ 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 +883,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 +898,7 @@ 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 +910,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/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/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.runtime 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}")

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,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os, platform, importlib.util
+import os, importlib.util
 
 os.environ["UNSLOTH_IS_PRESENT"] = "1"
 
+
+def _is_mlx_available():
+    try:
+        from unsloth_zoo.mlx.runtime 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 +36,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.runtime 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,9 @@ 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 +209,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: