mlx fixes

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:

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('<bos>')

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"

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: