From 797ae4cf40b023d42920ede68133c0cb14fe1757 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Thu, 7 May 2026 12:43:40 -0500 Subject: [PATCH 1/3] mlx fixes --- studio/backend/core/training/worker.py | 54 +++++++++++++++- .../backend/utils/datasets/chat_templates.py | 64 +++++++++++++++++-- unsloth/device_type.py | 18 +++++- unsloth/kernels/utils.py | 15 +++-- 4 files changed, 135 insertions(+), 16 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index ef5cafb175..9d9eac7743 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -417,6 +417,55 @@ def _normalize_mlx_studio_scheduler(value): return raw +def _resolve_local_dataset_files(file_paths: list) -> list[str]: + """Resolve local dataset paths 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 _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. @@ -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_local_dataset_files(file_paths) if not all_files: raise ValueError("No local dataset files found") - loader = UnslothTrainer._loader_for_files(all_files) + loader = _local_dataset_loader_for_files(all_files) return load_dataset(loader, data_files = all_files, split = "train") if hf_dataset: diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index 35fbaba8f0..c70ed2380f 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -8,6 +8,10 @@ This module contains functions for applying chat templates to datasets and generating dataset info summaries. """ +import importlib.util +import os +import platform + from .format_detection import detect_dataset_format, detect_multimodal_dataset, detect_custom_format_heuristic from .model_mappings import MODEL_TO_TEMPLATE_MAPPER from loggers import get_logger @@ -28,6 +32,33 @@ DEFAULT_ALPACA_TEMPLATE = """Below is an instruction that describes a task, pair {}""" +def _is_mlx_runtime() -> bool: + return ( + os.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1" + and platform.system() == "Darwin" + and platform.machine() == "arm64" + and importlib.util.find_spec("mlx") is not None + ) + + +def _fallback_apply_chat_template(convo) -> str: + parts = [] + for message in convo or []: + role = message.get("role") or message.get("from") or "user" + content = message.get("content") + if content is None: + content = message.get("value", "") + if isinstance(content, list): + content = "\n".join( + str(part.get("text", part)) + if isinstance(part, dict) else str(part) + for part in content + if not (isinstance(part, dict) and part.get("type") == "image") + ) + parts.append(f"<|im_start|>{role}\n{content}<|im_end|>") + return "\n".join(parts) + + def get_tokenizer_chat_template(tokenizer, model_name): """ Gets appropriate chat template for tokenizer based on model. @@ -40,6 +71,9 @@ def get_tokenizer_chat_template(tokenizer, model_name): Returns: tokenizer: Tokenizer with appropriate chat template applied """ + if _is_mlx_runtime(): + return tokenizer + try: from unsloth.chat_templates import get_chat_template except ImportError: @@ -252,7 +286,10 @@ def apply_chat_template_to_dataset( # Set alpaca chat template on tokenizer for saving (if not already set) # This ensures the template is saved with the model for inference - if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template): + if ( + not _is_mlx_runtime() + and 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") @@ -330,17 +367,32 @@ def apply_chat_template_to_dataset( if model_name: tokenizer = get_tokenizer_chat_template(tokenizer, model_name) + is_mlx_runtime = _is_mlx_runtime() + def _format_chatml(examples): convos = examples[chat_column] texts = [] for convo in convos: try: - text = tokenizer.apply_chat_template( - convo, - tokenize = False, - add_generation_prompt = False - ) + if is_mlx_runtime: + if ( + hasattr(tokenizer, "apply_chat_template") + and getattr(tokenizer, "chat_template", None) + ): + text = tokenizer.apply_chat_template( + convo, + tokenize = False, + add_generation_prompt = False + ) + else: + text = _fallback_apply_chat_template(convo) + else: + text = tokenizer.apply_chat_template( + convo, + tokenize = False, + add_generation_prompt = False + ) if remove_bos_prefix: text = text.removeprefix('') diff --git a/unsloth/device_type.py b/unsloth/device_type.py index a42d2b9fab..f8fb557df7 100644 --- a/unsloth/device_type.py +++ b/unsloth/device_type.py @@ -22,19 +22,35 @@ __all__ = [ "ALLOW_BITSANDBYTES", ] -import torch import functools import inspect +import importlib.util +import os +import platform from unsloth_zoo.utils import Version +_IS_MLX = ( + os.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1" + and platform.system() == "Darwin" + and platform.machine() == "arm64" + and importlib.util.find_spec("mlx") is not None +) + +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" diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index dd5a9cbf0e..40d7b8ec1e 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -160,14 +160,12 @@ else: # INTEL GPU Specific Logic if DEVICE_TYPE == "xpu": _gpu_getCurrentRawStream = torch._C._xpu_getCurrentRawStream -# NVIDIA GPU Default Logic -elif hasattr(torch._C, "_cuda_getCurrentRawStream"): - _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream -else: - # CPU-only torch wheel (no compiled CUDA backend). _get_tensor_stream - # is only invoked during real GPU work, so a no-op binding is safe. +elif DEVICE_TYPE == "mlx": def _gpu_getCurrentRawStream(_index = 0): return 0 +# NVIDIA GPU Default Logic +else: + _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream c_void_p = ctypes.c_void_p @@ -206,6 +204,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: From 051b2024be54628d4ad22b76ac4da3ff9619f473 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Thu, 7 May 2026 20:51:27 -0500 Subject: [PATCH 2/3] Fix studio integration, local dataset files, chat templates without the torch gpu imports --- .../backend/core/inference/mlx_inference.py | 4 +- studio/backend/core/training/worker.py | 16 ++-- .../tests/test_mlx_inference_backend.py | 9 ++- .../backend/utils/datasets/chat_templates.py | 79 ++++++------------- .../test_mlx_training_worker_behaviors.py | 8 +- unsloth/__init__.py | 25 +++--- unsloth/chat_templates.py | 14 +++- unsloth/device_type.py | 21 +++-- unsloth/kernels/utils.py | 7 +- 9 files changed, 90 insertions(+), 93 deletions(-) diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 1d2b03ecb9..6f6d90a27f 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/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( diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 9d9eac7743..84ac7f6409 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -417,8 +417,8 @@ def _normalize_mlx_studio_scheduler(value): return raw -def _resolve_local_dataset_files(file_paths: list) -> list[str]: - """Resolve local dataset paths without importing the GPU trainer.""" +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] = [] @@ -455,7 +455,7 @@ def _resolve_local_dataset_files(file_paths: list) -> list[str]: return all_files -def _local_dataset_loader_for_files(files: list[str]) -> str: +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" @@ -491,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, @@ -500,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 @@ -629,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 = _resolve_local_dataset_files(file_paths) + all_files = _resolve_mlx_local_dataset_files(file_paths) if not all_files: raise ValueError("No local dataset files found") - loader = _local_dataset_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: diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 868e537372..ce447bdd1f 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/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): diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index c70ed2380f..c910c24e2c 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -8,10 +8,6 @@ This module contains functions for applying chat templates to datasets and generating dataset info summaries. """ -import importlib.util -import os -import platform - from .format_detection import detect_dataset_format, detect_multimodal_dataset, detect_custom_format_heuristic from .model_mappings import MODEL_TO_TEMPLATE_MAPPER from loggers import get_logger @@ -33,30 +29,20 @@ DEFAULT_ALPACA_TEMPLATE = """Below is an instruction that describes a task, pair def _is_mlx_runtime() -> bool: - return ( - os.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1" - and platform.system() == "Darwin" - and platform.machine() == "arm64" - and importlib.util.find_spec("mlx") is not None - ) + try: + from unsloth_zoo.mlx.runtime import is_mlx_available + except ImportError: + return False + return is_mlx_available() -def _fallback_apply_chat_template(convo) -> str: - parts = [] - for message in convo or []: - role = message.get("role") or message.get("from") or "user" - content = message.get("content") - if content is None: - content = message.get("value", "") - if isinstance(content, list): - content = "\n".join( - str(part.get("text", part)) - if isinstance(part, dict) else str(part) - for part in content - if not (isinstance(part, dict) and part.get("type") == "image") - ) - parts.append(f"<|im_start|>{role}\n{content}<|im_end|>") - return "\n".join(parts) +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): @@ -71,9 +57,6 @@ def get_tokenizer_chat_template(tokenizer, model_name): Returns: tokenizer: Tokenizer with appropriate chat template applied """ - if _is_mlx_runtime(): - return tokenizer - try: from unsloth.chat_templates import get_chat_template except ImportError: @@ -94,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}") @@ -113,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}") @@ -286,13 +271,14 @@ def apply_chat_template_to_dataset( # Set alpaca chat template on tokenizer for saving (if not already set) # This ensures the template is saved with the model for inference - if ( - not _is_mlx_runtime() - and not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template) - ): + 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}") @@ -367,32 +353,17 @@ def apply_chat_template_to_dataset( if model_name: tokenizer = get_tokenizer_chat_template(tokenizer, model_name) - is_mlx_runtime = _is_mlx_runtime() - def _format_chatml(examples): convos = examples[chat_column] texts = [] for convo in convos: try: - if is_mlx_runtime: - if ( - hasattr(tokenizer, "apply_chat_template") - and getattr(tokenizer, "chat_template", None) - ): - text = tokenizer.apply_chat_template( - convo, - tokenize = False, - add_generation_prompt = False - ) - else: - text = _fallback_apply_chat_template(convo) - else: - text = tokenizer.apply_chat_template( - convo, - tokenize = False, - add_generation_prompt = False - ) + text = tokenizer.apply_chat_template( + convo, + tokenize = False, + add_generation_prompt = False + ) if remove_bos_prefix: text = text.removeprefix('') diff --git a/tests/studio/test_mlx_training_worker_behaviors.py b/tests/studio/test_mlx_training_worker_behaviors.py index 6c067ea00b..78b229d6e9 100644 --- a/tests/studio/test_mlx_training_worker_behaviors.py +++ b/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 diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 9b620a5c76..c1e2214f9f 100644 --- a/unsloth/__init__.py +++ b/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 diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 8376dc7e39..1c94e10f9f 100644 --- a/unsloth/chat_templates.py +++ b/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 diff --git a/unsloth/device_type.py b/unsloth/device_type.py index f8fb557df7..221b8b42ab 100644 --- a/unsloth/device_type.py +++ b/unsloth/device_type.py @@ -20,21 +20,24 @@ __all__ = [ "DEVICE_COUNT", "ALLOW_PREQUANTIZED_MODELS", "ALLOW_BITSANDBYTES", + "is_mlx_available", ] import functools import inspect -import importlib.util import os -import platform from unsloth_zoo.utils import Version -_IS_MLX = ( - os.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1" - and platform.system() == "Darwin" - and platform.machine() == "arm64" - and importlib.util.find_spec("mlx") is not None -) + +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 @@ -80,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 diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 40d7b8ec1e..86e76e08f7 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -164,8 +164,13 @@ elif DEVICE_TYPE == "mlx": def _gpu_getCurrentRawStream(_index = 0): return 0 # NVIDIA GPU Default Logic -else: +elif hasattr(torch._C, "_cuda_getCurrentRawStream"): _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream +else: + # CPU-only torch wheel (no compiled CUDA backend). _get_tensor_stream + # is only invoked during real GPU work, so a no-op binding is safe. + def _gpu_getCurrentRawStream(_index = 0): + return 0 c_void_p = ctypes.c_void_p From da2b371916fa02d3f98ae05f33db47fe0cbf58c7 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Fri, 8 May 2026 00:04:55 -0500 Subject: [PATCH 3/3] pass grad norm in mlx worker --- studio/backend/core/training/worker.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 84ac7f6409..763183a505 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -780,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), @@ -868,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", @@ -879,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, @@ -894,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, ) @@ -905,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