From 797ae4cf40b023d42920ede68133c0cb14fe1757 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Thu, 7 May 2026 12:43:40 -0500 Subject: [PATCH] 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: