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 ef5cafb175..763183a505 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: @@ -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 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 35fbaba8f0..c910c24e2c 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.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}") 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 a42d2b9fab..221b8b42ab 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.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 diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index dd5a9cbf0e..86e76e08f7 100644 --- a/unsloth/kernels/utils.py +++ b/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: