diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index ee7c9bbd51..e7bce2d33e 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/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( diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 549d733252..556545e680 100644 --- a/studio/backend/core/training/training.py +++ b/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"), diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index ef5cafb175..8cd1d1f9c8 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_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 diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 31f1d575d7..34ae9cfe72 100644 --- a/studio/backend/models/training.py +++ b/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") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 19202f3883..6875a13206 100644 --- a/studio/backend/routes/training.py +++ b/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, 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/tests/test_training_raw_support.py b/studio/backend/tests/test_training_raw_support.py index 876ee34686..c4aaaf3298 100644 --- a/studio/backend/tests/test_training_raw_support.py +++ b/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", diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index 35fbaba8f0..cfdd811853 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/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}") diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 5e68ccd72c..5d81f0df9c 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/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, diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index fb8a2f899e..0cb881e634 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/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; diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 168e9ad329..f0c90dd9c6 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/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 diff --git a/tests/studio/test_is_mlx_dispatch_gate.py b/tests/studio/test_is_mlx_dispatch_gate.py index fc07a497e7..043f2796aa 100644 --- a/tests/studio/test_is_mlx_dispatch_gate.py +++ b/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 # --------------------------------------------------------------------------- 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..75fb169e83 100644 --- a/unsloth/__init__.py +++ b/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 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 a42d2b9fab..9bad9be0e4 100644 --- a/unsloth/device_type.py +++ b/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 diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index dd5a9cbf0e..f77baea281 100644 --- a/unsloth/kernels/utils.py +++ b/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: