From 38dacb8a1f49906d2fa23e4dace198a9c56b5220 Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Wed, 8 Jul 2026 18:25:26 +0800 Subject: [PATCH] Add MLX backend support for CLI unsloth train (#6709) * feat(studio): route CLI trainer to MLX backend * fix(studio): harden MLX trainer routing * fix(studio): harden MLX trainer adapter routing * test(studio): assert MLX CLI activation order * fix(studio): address MLX CLI review feedback * feat(cli): support MLX in legacy script * fix(cli): adapt MLX tokenizer for raw text * fix(cli): omit unsupported MLX eval batch arg * fix(cli): feed raw text to MLX trainer * Fix CLI MLX routing and Python 3.9 annotations Route the MLX backend through create_mlx_trainer_adapter so the torch-free Apple Silicon path never imports trainer.py (torch/unsloth/trl). Replace from __future__ import annotations with typing.Optional/Union so the CLI annotations stay Python 3.9 compatible without the unused-import lint hit. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Strip return_tensors from MLX raw-text tokenizer proxy On a torch-free MLX install, RawTextDataLoader calls the tokenizer with return_tensors='pt'; the callable proxy forwarded that to the HF tokenizer, which tried to build torch tensors and failed before training. Drop return_tensors so the MLX path returns plain token ids. * Tighten CLI MLX-backend comments --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/training/trainer.py | 31 +- studio/backend/core/training/training.py | 624 +++++++++++++++--- studio/backend/core/training/worker.py | 248 +++++-- .../backend/tests/test_training_preflight.py | 232 +++++++ unsloth-cli.py | 243 ++++--- unsloth_cli/commands/train.py | 46 +- 6 files changed, 1162 insertions(+), 262 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 20b2305a5a..958f8f4197 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -62,7 +62,6 @@ from loggers import get_logger import time from pathlib import Path from typing import Any, Dict, List, Optional, Callable -from dataclasses import dataclass import pandas as pd from datasets import Dataset from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset @@ -86,6 +85,11 @@ from utils.native_path_leases import child_env_without_native_path_secret from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) +from .training import ( + TrainingProgress, + create_mlx_trainer_adapter, + should_use_mlx_training_backend, +) logger = get_logger(__name__) @@ -104,31 +108,16 @@ def _build_report_targets(training_args) -> list[str] | str: return report_to or "none" -@dataclass -class TrainingProgress: - """Training progress tracking""" - - epoch: float = 0 - step: int = 0 - total_steps: int = 0 - loss: Optional[float] = None - learning_rate: Optional[float] = None - is_training: bool = False - is_completed: bool = False - error: Optional[str] = None - status_message: str = "Ready to train" # Current stage - elapsed_seconds: Optional[float] = None - eta_seconds: Optional[float] = None - grad_norm: Optional[float] = None - num_tokens: Optional[int] = None - eval_loss: Optional[float] = None - - class UnslothTrainer: """ Unsloth Training Backend """ + def __new__(cls, *args, **kwargs): + if cls is UnslothTrainer and should_use_mlx_training_backend(): + return create_mlx_trainer_adapter(*args, **kwargs) + return super().__new__(cls) + def __init__(self): self.model = None self.tokenizer = None diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index f4233fcf04..2ddda19951 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -14,17 +14,19 @@ import json as _json import math import multiprocessing as mp import os +import platform import queue import re import shutil import threading import time +import traceback import structlog from datetime import datetime, timezone from loggers import get_logger -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Optional, Tuple, Any, TYPE_CHECKING +from typing import Optional, Tuple, Any, Callable, Union, TYPE_CHECKING if TYPE_CHECKING: import matplotlib.pyplot as plt @@ -98,6 +100,107 @@ def _coerce_optional_nonneg_float(name: str, value): return coerced +def is_apple_silicon_training_platform() -> bool: + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def is_mlx_training_device(device: Any) -> bool: + return ( + str(device).lower() == "mlx" + or str(device).lower().endswith(".mlx") + or getattr(device, "name", "").lower() == "mlx" + ) + + +def should_use_mlx_training_backend(*, device: Optional[Any] = None) -> bool: + if device is not None: + return is_mlx_training_device(device) + return is_apple_silicon_training_platform() + + +def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: + """Build the normalized worker config shared by Studio and the CLI adapter.""" + config = { + "model_name": values["model_name"], + "project_name": values.get("project_name"), + "training_type": values.get("training_type", "LoRA/QLoRA"), + "hf_token": values.get("hf_token", ""), + "load_in_4bit": values.get("load_in_4bit", True), + "max_seq_length": values.get("max_seq_length", 2048), + "vision_image_size": values.get("vision_image_size"), + "hf_dataset": values.get("hf_dataset", ""), + "local_datasets": values.get("local_datasets"), + "local_eval_datasets": values.get("local_eval_datasets"), + "format_type": values.get("format_type", ""), + "subset": values.get("subset"), + "train_split": values.get("train_split", "train"), + "eval_split": values.get("eval_split"), + "eval_steps": values.get("eval_steps", 0.00), + "dataset_streaming": values.get("dataset_streaming", False), + "dataset_slice_start": values.get("dataset_slice_start"), + "dataset_slice_end": values.get("dataset_slice_end"), + "custom_format_mapping": values.get("custom_format_mapping"), + "is_dataset_image": values.get("is_dataset_image", False), + "is_dataset_audio": values.get("is_dataset_audio", False), + "is_embedding": values.get("is_embedding", False), + "num_epochs": values.get("num_epochs", 3), + "learning_rate": values.get("learning_rate", "2e-4"), + "embedding_learning_rate": values.get("embedding_learning_rate"), + "batch_size": values.get("batch_size", 2), + "gradient_accumulation_steps": values.get("gradient_accumulation_steps", 4), + "warmup_steps": values.get("warmup_steps"), + "warmup_ratio": values.get("warmup_ratio"), + "max_steps": values.get("max_steps", 0), + "save_steps": values.get("save_steps", 0), + "weight_decay": values.get("weight_decay", 0.001), + "max_grad_norm": values.get("max_grad_norm", 0.0), + "max_grad_value": _coerce_optional_nonneg_float( + "max_grad_value", values.get("max_grad_value") + ), + "max_grad_leaf_norm": _coerce_optional_nonneg_float( + "max_grad_leaf_norm", values.get("max_grad_leaf_norm") + ), + "cast_norm_output_to_input_dtype": _coerce_optional_bool( + values.get("cast_norm_output_to_input_dtype"), True + ), + "random_seed": _coerce_seed(values.get("random_seed")), + "packing": values.get("packing", False), + "optim": values.get("optim", "adamw_8bit"), + "lr_scheduler_type": values.get("lr_scheduler_type", "linear"), + "use_lora": values.get("use_lora", True), + "lora_r": values.get("lora_r", 16), + "lora_alpha": values.get("lora_alpha", 16), + "lora_dropout": values.get("lora_dropout", 0.0), + "target_modules": values.get("target_modules"), + "gradient_checkpointing": values.get("gradient_checkpointing", "unsloth"), + "use_rslora": values.get("use_rslora", False), + "use_loftq": values.get("use_loftq", False), + "train_on_completions": values.get("train_on_completions", False), + "finetune_vision_layers": values.get("finetune_vision_layers", True), + "finetune_language_layers": values.get("finetune_language_layers", True), + "finetune_attention_modules": values.get("finetune_attention_modules", True), + "finetune_mlp_modules": values.get("finetune_mlp_modules", True), + "enable_wandb": values.get("enable_wandb", False), + "wandb_token": values.get("wandb_token"), + "wandb_project": values.get("wandb_project", "unsloth-training"), + "enable_tensorboard": values.get("enable_tensorboard", False), + "tensorboard_dir": values.get("tensorboard_dir", "runs"), + "resume_from_checkpoint": values.get("resume_from_checkpoint"), + "trust_remote_code": values.get("trust_remote_code", False), + "approved_remote_code_fingerprint": values.get("approved_remote_code_fingerprint"), + "subject": values.get("subject"), + "gpu_ids": values.get("gpu_ids"), + "s3_config": values.get("s3_config"), + "disable_xet": values.get("disable_xet", False), + } + for key in ("output_dir", "allow_external_output_dir"): + if key in values: + config[key] = values.get(key) + if config["training_type"] == "Full Finetuning": + config["load_in_4bit"] = False + return config + + _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$") @@ -133,7 +236,7 @@ def _s3_dataset_name(s3_dataset: Any) -> Optional[str]: return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}" -def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None: +def _cleanup_cancelled_checkpoints(output_dir: Union[str, os.PathLike]) -> None: """Remove only HF Trainer ``tmp-checkpoint-/`` partials after a cancel. Completed ``checkpoint-/`` dirs survive. Symlinked output_dir / children @@ -183,7 +286,7 @@ PLOT_HEIGHT = 3.5 @dataclass class TrainingProgress: - """Mirror of trainer.TrainingProgress so the parent never imports heavy ML modules.""" + """Shared training progress payload for Studio and backend-aware trainers.""" epoch: float = 0 step: int = 0 @@ -200,6 +303,423 @@ class TrainingProgress: num_tokens: Optional[int] = None eval_loss: Optional[float] = None peak_memory_gb: Optional[float] = None + output_dir: Optional[str] = None + + +class _MLXTrainerAdapter: + """Adapts the legacy UnslothTrainer API to the shared Studio MLX worker path.""" + + def __init__(self): + self.model = None + self.tokenizer = None + self.trainer = None + self.training_thread = None + self.training_progress = TrainingProgress() + self.progress_callbacks: list[Callable[[TrainingProgress], None]] = [] + self.is_training = False + self.should_stop = False + self.save_on_stop = True + self.load_in_4bit = True + self.output_dir = None + + self.is_cpt = False + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model_name = None + self.max_seq_length = None + + self._model_config: dict[str, Any] = {} + self._peft_config: dict[str, Any] = {} + self._dataset_config: dict[str, Any] = {} + self._event_queue: Optional[queue.Queue] = None + self._stop_queue: Optional[queue.Queue] = None + self._pump_thread: Optional[threading.Thread] = None + self._lock = threading.Lock() + + def _activate_transformers_for_model(self, model_name: str, hf_token: Optional[str]) -> None: + try: + from utils.transformers_version import activate_transformers_for_subprocess + activate_transformers_for_subprocess(model_name, hf_token) + except Exception as exc: + logger.warning("MLX trainer adapter Transformers activation failed", error = str(exc)) + + def add_progress_callback(self, callback: Callable[[TrainingProgress], None]): + self.progress_callbacks.append(callback) + + def _update_progress(self, **kwargs): + with self._lock: + for key, value in kwargs.items(): + if hasattr(self.training_progress, key): + setattr(self.training_progress, key, value) + progress = self.training_progress + for callback in self.progress_callbacks: + try: + callback(progress) + except Exception: + pass + + def load_model( + self, + model_name: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + hf_token: Optional[str] = None, + is_dataset_image: bool = False, + is_dataset_audio: bool = False, + trust_remote_code: bool = False, + full_finetuning: bool = False, + gpu_ids: Optional[list[int]] = None, + ) -> bool: + self.model_name = model_name + self.max_seq_length = max_seq_length + self.load_in_4bit = load_in_4bit + self._audio_type = None + self._activate_transformers_for_model(model_name, hf_token) + try: + from utils.models import detect_audio_type, is_vision_model + + self._audio_type = detect_audio_type(model_name, hf_token) + if self._audio_type == "audio_vlm": + self.is_audio = False + self.is_audio_vlm = bool(is_dataset_audio) + self._audio_type = None + else: + self.is_audio = self._audio_type is not None + self.is_audio_vlm = False + vision = is_vision_model(model_name, hf_token = hf_token) if not self.is_audio else False + self.is_vlm = not self.is_audio_vlm and vision and bool(is_dataset_image) + except Exception as exc: + logger.warning("MLX trainer adapter model type detection failed", error = str(exc)) + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model = object() + self.tokenizer = object() + self._model_config = { + "model_name": model_name, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + "hf_token": hf_token or "", + "is_dataset_image": bool(is_dataset_image), + "is_dataset_audio": bool(is_dataset_audio), + "trust_remote_code": bool(trust_remote_code), + "gpu_ids": gpu_ids, + } + self._update_progress( + is_training = False, + is_completed = False, + error = None, + step = 0, + loss = 0.0, + epoch = 0, + status_message = f"Queued MLX model load: {model_name}", + ) + return True + + def prepare_model_for_training( + self, + use_lora: bool = True, + finetune_vision_layers: bool = True, + finetune_language_layers: bool = True, + finetune_attention_modules: bool = True, + finetune_mlp_modules: bool = True, + target_modules: Optional[Union[list, str]] = None, + lora_r: int = 16, + lora_alpha: int = 16, + lora_dropout: float = 0.0, + use_gradient_checkpointing: Union[str, bool] = "unsloth", + use_rslora: bool = False, + use_loftq: bool = False, + ) -> bool: + self._peft_config = { + "use_lora": bool(use_lora), + "lora_r": lora_r, + "lora_alpha": lora_alpha, + "lora_dropout": lora_dropout, + "target_modules": target_modules, + "gradient_checkpointing": use_gradient_checkpointing, + "use_rslora": bool(use_rslora), + "use_loftq": bool(use_loftq), + "finetune_vision_layers": bool(finetune_vision_layers), + "finetune_language_layers": bool(finetune_language_layers), + "finetune_attention_modules": bool(finetune_attention_modules), + "finetune_mlp_modules": bool(finetune_mlp_modules), + } + self._update_progress(status_message = "Queued MLX training setup") + return True + + def load_and_format_dataset( + self, + dataset_source: Optional[str], + format_type: str = "auto", + local_datasets: Optional[list[str]] = None, + local_eval_datasets: Optional[list[str]] = None, + custom_format_mapping: Optional[dict[str, Any]] = None, + subset: Optional[str] = None, + train_split: str = "train", + eval_split: Optional[str] = None, + dataset_streaming: bool = False, + eval_steps: float = 0.00, + dataset_slice_start: Optional[int] = None, + dataset_slice_end: Optional[int] = None, + is_cpt: bool = False, + s3_config: dict = None, + ) -> Optional[tuple]: + self._dataset_config = { + "hf_dataset": dataset_source or "", + "local_datasets": local_datasets, + "local_eval_datasets": local_eval_datasets, + "format_type": format_type or "", + "custom_format_mapping": custom_format_mapping, + "subset": subset, + "train_split": train_split or "train", + "eval_split": eval_split, + "dataset_streaming": bool(dataset_streaming), + "eval_steps": eval_steps or 0.0, + "dataset_slice_start": dataset_slice_start, + "dataset_slice_end": dataset_slice_end, + "s3_config": s3_config, + } + self.is_cpt = bool(is_cpt) + self._update_progress(status_message = "Queued MLX dataset load") + return ({"dataset": [], "final_format": "deferred_mlx_cli", "success": True}, None) + + def start_training( + self, + dataset = None, + eval_dataset = None, + **training_args, + ) -> bool: + if self.is_training and self.training_thread and self.training_thread.is_alive(): + return False + if self._pump_thread and self._pump_thread.is_alive(): + self._pump_thread.join(timeout = 2.0) + if self._pump_thread.is_alive(): + self._update_progress(error = "Previous training event pump is still finalizing") + return False + if not self._model_config: + self._update_progress(error = "Model not loaded") + return False + if not self._dataset_config: + self._update_progress(error = "Dataset not loaded") + return False + if self.is_cpt: + self._update_progress( + error = "Continued Pretraining is not supported for MLX training yet.", + is_training = False, + is_completed = False, + ) + return False + + config = self._build_worker_config(training_args) + event_queue = queue.Queue() + stop_queue = queue.Queue() + self._event_queue = event_queue + self._stop_queue = stop_queue + self.should_stop = False + self.is_training = True + self.training_progress = TrainingProgress( + is_training = True, + status_message = "Initializing MLX training...", + ) + + self.training_thread = threading.Thread( + target = self._run_training_thread, + args = (config, event_queue, stop_queue), + daemon = True, + ) + self._pump_thread = threading.Thread( + target = self._pump_events, + args = (event_queue, self.training_thread), + daemon = True, + ) + self.training_thread.start() + self._pump_thread.start() + return True + + def _build_worker_config(self, training_args: dict[str, Any]) -> dict[str, Any]: + peft = { + "use_lora": True, + "lora_r": 16, + "lora_alpha": 16, + "lora_dropout": 0.0, + "target_modules": None, + "gradient_checkpointing": "unsloth", + "use_rslora": False, + "use_loftq": False, + "finetune_vision_layers": True, + "finetune_language_layers": True, + "finetune_attention_modules": True, + "finetune_mlp_modules": True, + **self._peft_config, + } + output_dir = training_args.get("output_dir") + if output_dir: + output_dir = os.path.abspath(os.path.expanduser(str(output_dir))) + values = { + **self._model_config, + **self._dataset_config, + **training_args, + "training_type": ( + "Continued Pretraining" + if self.is_cpt + else "LoRA/QLoRA" + if peft["use_lora"] + else "Full Finetuning" + ), + **peft, + "output_dir": output_dir, + "allow_external_output_dir": bool(output_dir), + } + config = _build_training_worker_config(values) + config["resolved_gpu_ids"] = None + config["gpu_selection"] = None + return config + + def _run_training_thread( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + try: + self._run_mlx_worker(config, event_queue, stop_queue) + except Exception as exc: + if event_queue is not None: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def _run_mlx_worker( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + from .worker import run_mlx_training_process + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + ) + + def _pump_events(self, event_queue: queue.Queue, training_thread: threading.Thread): + while True: + event = None + try: + event = event_queue.get(timeout = 0.25) + except queue.Empty: + pass + if event is not None: + self._handle_event(event) + continue + if not training_thread.is_alive(): + self._drain_events(event_queue) + with self._lock: + if self.training_progress.is_training: + self.training_progress.is_training = False + if self.should_stop: + self.training_progress.status_message = "Training stopped." + elif ( + not self.training_progress.error + and not self.training_progress.is_completed + ): + self.training_progress.error = "Training process exited unexpectedly" + self.is_training = False + self._event_queue = None + self._stop_queue = None + return + + def _drain_events(self, event_queue: Optional[queue.Queue] = None): + event_queue = event_queue or self._event_queue + if event_queue is None: + return + while True: + try: + self._handle_event(event_queue.get_nowait()) + except queue.Empty: + return + + def _handle_event(self, event: dict[str, Any]): + etype = event.get("type") + if etype == "status": + self._update_progress( + status_message = event.get("status_message") or event.get("message") or "" + ) + return + if etype == "progress": + self._update_progress( + step = event.get("step", self.training_progress.step), + epoch = event.get("epoch", self.training_progress.epoch), + loss = event.get("loss", self.training_progress.loss), + learning_rate = event.get("learning_rate", self.training_progress.learning_rate), + total_steps = event.get("total_steps", self.training_progress.total_steps), + elapsed_seconds = event.get( + "elapsed_seconds", + self.training_progress.elapsed_seconds, + ), + eta_seconds = event.get("eta_seconds", self.training_progress.eta_seconds), + grad_norm = event.get("grad_norm", self.training_progress.grad_norm), + num_tokens = event.get("num_tokens", self.training_progress.num_tokens), + eval_loss = event.get("eval_loss", self.training_progress.eval_loss), + peak_memory_gb = event.get("peak_memory_gb", self.training_progress.peak_memory_gb), + ) + return + if etype == "complete": + status_message = event.get("status_message") or "Training completed" + output_dir = event.get("output_dir") + was_cancelled = self.should_stop or status_message.strip().lower() in { + "training cancelled", + "training stopped", + } + self.output_dir = output_dir + self._update_progress( + is_training = False, + is_completed = not was_cancelled, + error = None, + status_message = status_message, + output_dir = output_dir, + ) + self.is_training = False + return + if etype == "error": + self._update_progress( + is_training = False, + is_completed = False, + error = event.get("error") or event.get("message") or "Training failed", + ) + self.is_training = False + return + + def stop_training(self, save: bool = True): + self.should_stop = True + self.save_on_stop = bool(save) + if self._stop_queue is not None: + self._stop_queue.put({"type": "stop", "save": save}) + status_message = ( + "Stopping training and saving checkpoint..." if save else "Cancelling training..." + ) + self._update_progress(status_message = status_message) + return True + + def get_training_progress(self) -> TrainingProgress: + pump_thread = self._pump_thread + training_thread = self.training_thread + if ( + pump_thread is not None + and pump_thread.is_alive() + and (training_thread is None or not training_thread.is_alive()) + and threading.current_thread() is not pump_thread + ): + pump_thread.join(timeout = 5.0) + if pump_thread is None or not pump_thread.is_alive(): + self._drain_events() + with self._lock: + return replace(self.training_progress) + + +def create_mlx_trainer_adapter(*args, **kwargs): + return _MLXTrainerAdapter(*args, **kwargs) class TrainingBackend: @@ -296,86 +816,7 @@ class TrainingBackend: # treat this fresh setup as a recoverable death. self._pump_running = False - # Build config dict for the subprocess - config = { - "model_name": kwargs["model_name"], - "project_name": kwargs.get("project_name"), - "training_type": kwargs.get("training_type", "LoRA/QLoRA"), - "hf_token": kwargs.get("hf_token", ""), - "load_in_4bit": kwargs.get("load_in_4bit", True), - "max_seq_length": kwargs.get("max_seq_length", 2048), - "vision_image_size": kwargs.get("vision_image_size"), - "hf_dataset": kwargs.get("hf_dataset", ""), - "local_datasets": kwargs.get("local_datasets"), - "local_eval_datasets": kwargs.get("local_eval_datasets"), - "format_type": kwargs.get("format_type", ""), - "subset": kwargs.get("subset"), - "train_split": kwargs.get("train_split", "train"), - "eval_split": kwargs.get("eval_split"), - "eval_steps": kwargs.get("eval_steps", 0.00), - "dataset_streaming": kwargs.get("dataset_streaming", False), - "dataset_slice_start": kwargs.get("dataset_slice_start"), - "dataset_slice_end": kwargs.get("dataset_slice_end"), - "custom_format_mapping": kwargs.get("custom_format_mapping"), - "is_dataset_image": kwargs.get("is_dataset_image", False), - "is_dataset_audio": kwargs.get("is_dataset_audio", False), - "is_embedding": kwargs.get("is_embedding", False), - "num_epochs": kwargs.get("num_epochs", 3), - "learning_rate": kwargs.get("learning_rate", "2e-4"), - "embedding_learning_rate": kwargs.get("embedding_learning_rate"), - "batch_size": kwargs.get("batch_size", 2), - "gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4), - "warmup_steps": kwargs.get("warmup_steps"), - "warmup_ratio": kwargs.get("warmup_ratio"), - "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": _coerce_optional_nonneg_float( - "max_grad_value", kwargs.get("max_grad_value") - ), - "max_grad_leaf_norm": _coerce_optional_nonneg_float( - "max_grad_leaf_norm", kwargs.get("max_grad_leaf_norm") - ), - "cast_norm_output_to_input_dtype": _coerce_optional_bool( - kwargs.get("cast_norm_output_to_input_dtype"), True - ), - # MLX/CUDA/embedding workers need an int (transformers.set_seed(None) raises). - "random_seed": _coerce_seed(kwargs.get("random_seed")), - "packing": kwargs.get("packing", False), - "optim": kwargs.get("optim", "adamw_8bit"), - "lr_scheduler_type": kwargs.get("lr_scheduler_type", "linear"), - "use_lora": kwargs.get("use_lora", True), - "lora_r": kwargs.get("lora_r", 16), - "lora_alpha": kwargs.get("lora_alpha", 16), - "lora_dropout": kwargs.get("lora_dropout", 0.0), - "target_modules": kwargs.get("target_modules"), - "gradient_checkpointing": kwargs.get("gradient_checkpointing", "unsloth"), - "use_rslora": kwargs.get("use_rslora", False), - "use_loftq": kwargs.get("use_loftq", False), - "train_on_completions": kwargs.get("train_on_completions", False), - "finetune_vision_layers": kwargs.get("finetune_vision_layers", True), - "finetune_language_layers": kwargs.get("finetune_language_layers", True), - "finetune_attention_modules": kwargs.get("finetune_attention_modules", True), - "finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True), - "enable_wandb": kwargs.get("enable_wandb", False), - "wandb_token": kwargs.get("wandb_token"), - "wandb_project": kwargs.get("wandb_project", "unsloth-training"), - "enable_tensorboard": kwargs.get("enable_tensorboard", False), - "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), - "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), - "trust_remote_code": kwargs.get("trust_remote_code", False), - "approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"), - "subject": kwargs.get("subject"), - "gpu_ids": kwargs.get("gpu_ids"), - "s3_config": kwargs.get("s3_config"), - # Flipped to True only by the HTTP-fallback respawn after a stall. - "disable_xet": kwargs.get("disable_xet", False), - } - - # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. - if config["training_type"] == "Full Finetuning": - config["load_in_4bit"] = False + config = _build_training_worker_config(kwargs) # Split GPU validation from placement around the VRAM hook: # * Explicit gpu_ids are validated here (raises -> the route returns 400 @@ -401,7 +842,7 @@ class TrainingBackend: ) defer_auto_selection = False - if _hw.DEVICE == _hw.DeviceType.MLX: + if should_use_mlx_training_backend(device = _hw.DEVICE): config["resolved_gpu_ids"] = None config["gpu_selection"] = None elif gpu_ids: @@ -1022,17 +1463,22 @@ class TrainingBackend: self._progress.is_training = True elif etype == "complete": - self._progress.is_training = False - self._progress.is_completed = True - self._output_dir = event.get("output_dir") msg = event.get("status_message", "Training completed") + stopped = self._should_stop or msg.strip().lower() in { + "training cancelled", + "training stopped", + } + self._progress.is_training = False + self._progress.is_completed = not stopped + self._output_dir = event.get("output_dir") + self._progress.output_dir = self._output_dir self._progress.status_message = msg if not self._db_run_created and self.current_job_id and self._db_config: db_action = "create_and_finalize" else: db_action = "finalize" db_action_kwargs = { - "status": "stopped" if self._should_stop else "completed", + "status": "stopped" if stopped else "completed", "output_dir": self._output_dir, } diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 17dc1299ca..0ff4d517ed 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1309,14 +1309,18 @@ def _normalize_mlx_studio_scheduler(value): def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]: - """Resolve Studio local dataset uploads without importing the GPU trainer.""" + """Resolve CLI paths and 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)) - ) + dataset_path = Path(os.path.expanduser(str(dataset_file))) + if dataset_path.is_absolute(): + file_path = str(dataset_path) + elif dataset_path.exists(): + file_path = str(dataset_path.resolve()) + else: + file_path = str(resolve_dataset_path(str(dataset_file))) file_path_obj = Path(file_path) if file_path_obj.is_dir(): @@ -1355,6 +1359,58 @@ def _mlx_local_dataset_loader_for_files(files: list[str]) -> str: raise ValueError(f"Unsupported dataset format: {files[0]}") +_MLX_WORKER_COMPLETE = "_mlx_worker_complete" + + +def _start_mlx_stop_poller(stop_queue): + import queue as _queue + import threading + + stop_save = [True] + stop_requested = [False] + trainer_ref = [None] + + def is_stop_requested(): + return stop_requested[0] + + def poll_stop(): + while True: + try: + msg = stop_queue.get(timeout = 0.25) + if msg and msg.get("type") == _MLX_WORKER_COMPLETE: + return + if msg and msg.get("type") == "stop": + stop_save[0] = msg.get("save", True) + stop_requested[0] = True + trainer = trainer_ref[0] + if trainer is not None: + trainer.stop_requested = True + return + except _queue.Empty: + continue + except (EOFError, OSError): + return + + stop_thread = threading.Thread(target = poll_stop, daemon = True) + stop_thread.start() + return stop_save, stop_requested, trainer_ref, is_stop_requested, stop_thread + + +def _resolve_mlx_output_dir(config, model_name): + from utils.paths import resolve_output_dir, default_run_dir_name + + output_dir = config.get("output_dir", "") + if not output_dir: + output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + return str(resolve_output_dir(output_dir)) + if config.get("allow_external_output_dir"): + output_path = Path(output_dir).expanduser() + if not output_path.is_absolute(): + output_path = Path.cwd() / output_path + return str(output_path.resolve()) + return str(resolve_output_dir(output_dir)) + + def _run_mlx_training(event_queue, stop_queue, config): """Self-contained MLX training path for Apple Silicon. @@ -1363,8 +1419,6 @@ def _run_mlx_training(event_queue, stop_queue, config): """ import time import math - import threading - import queue as _queue from pathlib import Path def _send(event_type, **kwargs): @@ -1374,31 +1428,9 @@ def _run_mlx_training(event_queue, stop_queue, config): kwargs["message"] = sm event_queue.put({"type": event_type, "ts": time.time(), **kwargs}) - _stop_save = [True] - _stop_requested = [False] - _trainer_ref = [None] - - def _is_stop_requested(): - return _stop_requested[0] - - def _poll_stop(): - while True: - try: - msg = stop_queue.get(timeout = 1.0) - if msg and msg.get("type") == "stop": - _stop_save[0] = msg.get("save", True) - _stop_requested[0] = True - trainer = _trainer_ref[0] - if trainer is not None: - trainer.stop_requested = True - return - except _queue.Empty: - continue - except (EOFError, OSError): - return - - stop_thread = threading.Thread(target = _poll_stop, daemon = True) - stop_thread.start() + _stop_save, _stop_requested, _trainer_ref, _is_stop_requested, _stop_thread = ( + _start_mlx_stop_poller(stop_queue) + ) _send("status", status_message = "Loading MLX libraries...") @@ -1804,21 +1836,14 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 5. Build output dir ── # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it - from utils.paths import resolve_output_dir, ensure_dir + from utils.paths import ensure_dir - output_dir = config.get("output_dir", "") - if not output_dir: - output_dir = build_default_output_dir_name( - model_name, - config.get("project_name"), - ) - output_dir = str(resolve_output_dir(output_dir)) + output_dir = _resolve_mlx_output_dir(config, model_name) ensure_dir(Path(output_dir)) # ── 6. Create trainer ── eval_steps_val = config.get("eval_steps", 0) or 0 if isinstance(eval_steps_val, float) and 0 < eval_steps_val < 1: - # Studio sometimes sends fraction-of-total-steps eval_steps_val = max(1, int(eval_steps_val * max_steps)) else: eval_steps_val = int(eval_steps_val) @@ -2043,12 +2068,27 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 11. Run training ── gc.collect() mx.synchronize() - trainer.train(resume_from_checkpoint = resume_from_checkpoint) + _save_model = trainer.save_model + + def _skip_internal_final_save(*args, **kwargs): + raise ValueError("worker owns final save") + + trainer.save_model = _skip_internal_final_save + try: + trainer.train(resume_from_checkpoint = resume_from_checkpoint) + finally: + trainer.save_model = _save_model # ── 12. Save and finalize ── - if trainer.stop_requested and not _stop_save[0]: - # User clicked "Cancel" (save=False) — skip saving - _send("complete", output_dir = None, status_message = "Training cancelled") + if trainer.stop_requested: + if not _stop_save[0]: + # Cancel (save=False): skip saving. + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving stopped model...") + mx.synchronize() + trainer.save_model(output_dir) + _send("complete", output_dir = output_dir, status_message = "Training stopped") else: _send("status", status_message = "Saving model...") mx.synchronize() @@ -2067,6 +2107,79 @@ def _run_mlx_training(event_queue, stop_queue, config): pass +def _is_current_process_apple_silicon() -> bool: + import platform + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def run_mlx_training_process( + *, + event_queue: Any, + stop_queue: Any, + config: dict, + transformers_activated: bool = False, +) -> None: + """MLX worker entrypoint shared by Studio subprocesses and the CLI adapter.""" + model_name = config["model_name"] + + backend_path = str(Path(__file__).resolve().parent.parent.parent) + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.hf_xet_fallback import child_should_disable_xet + + if child_should_disable_xet(config): + os.environ["HF_HUB_DISABLE_XET"] = "1" + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + + if not transformers_activated: + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + + from utils.hardware import hardware as _hw + + _hw.detect_hardware() + if _hw.DEVICE != _hw.DeviceType.MLX: + event_queue.put( + { + "type": "error", + "error": "MLX training requires Apple Silicon with the MLX backend available.", + "stack": "", + "ts": time.time(), + } + ) + return + + if config.get("is_dataset_audio"): + event_queue.put( + { + "type": "error", + "error": "Audio dataset training is not yet supported on Apple Silicon.", + "stack": "", + "ts": time.time(), + } + ) + return + + try: + try: + _run_mlx_training(event_queue, stop_queue, config) + finally: + try: + stop_queue.put({"type": _MLX_WORKER_COMPLETE}) + except (EOFError, OSError, ValueError): + pass + except Exception as exc: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None: """Subprocess entrypoint. Fresh Python — no stale module state. @@ -2141,36 +2254,26 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) + from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend + + mlx_backend_requested = is_apple_silicon_training_platform() + + mlx_transformers_activated = False + if mlx_backend_requested and _is_current_process_apple_silicon(): + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + mlx_transformers_activated = True + from utils.hardware import hardware as _hw _hw.detect_hardware() - if _hw.DEVICE == _hw.DeviceType.MLX: - if config.get("is_dataset_audio"): - event_queue.put( - { - "type": "error", - "error": "Audio dataset training is not yet supported on Apple Silicon.", - "stack": "", - "ts": time.time(), - } - ) - return - # Activate correct transformers version (Gemma-4 needs a 5.x sidecar, etc.) - # Must happen before any transformers/mlx-lm imports in _run_mlx_training. - # Non-fatal: fall through with whatever version is installed, but log - # the failure instead of swallowing it (issue #6103). - _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) - try: - _run_mlx_training(event_queue, stop_queue, config) - except Exception as exc: - event_queue.put( - { - "type": "error", - "error": str(exc), - "stack": traceback.format_exc(limit = 20), - "ts": time.time(), - } - ) + if mlx_backend_requested or should_use_mlx_training_backend(device = _hw.DEVICE): + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + transformers_activated = mlx_transformers_activated, + ) return # ── 1. Activate correct transformers version BEFORE any ML imports ── @@ -2693,7 +2796,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) - from core.training.trainer import UnslothTrainer, TrainingProgress + from core.training.training import TrainingProgress + from core.training.trainer import UnslothTrainer from utils.paths import ( ensure_dir, resolve_output_dir, diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 54048a65dd..47c6669f8f 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -6,9 +6,15 @@ empty-chat-template crash) before train(). The real methods are bound onto a lig fake self so the production logic runs against controlled batches.""" import importlib +import json +import os +import queue +import subprocess import sys +import threading import types import unittest +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -184,5 +190,231 @@ class TestChatTemplateRendersEmpty(unittest.TestCase): self.assertFalse(s._chat_template_renders_empty()) +def _clear_trainer_module(package: str): + sys.modules.pop(f"{package}.trainer", None) + pkg = sys.modules.get(package) + if pkg is not None and hasattr(pkg, "trainer"): + delattr(pkg, "trainer") + + +def _set_training_platform(monkeypatch, package: str, backend: str): + training_mod = importlib.import_module(f"{package}.training") + from utils.hardware import hardware as hw + + monkeypatch.setattr(hw, "DEVICE", None) + monkeypatch.setattr( + training_mod.platform, + "system", + lambda: "Darwin" if backend == "mlx" else "Linux", + ) + monkeypatch.setattr( + training_mod.platform, + "machine", + lambda: "arm64" if backend == "mlx" else "x86_64", + ) + + +def _load_trainer_module( + monkeypatch, + backend: str, + package: str = "core.training", +): + _set_training_platform(monkeypatch, package, backend) + _clear_trainer_module(package) + if package in sys.modules: + importlib.reload(sys.modules[package]) + trainer_mod = importlib.import_module(f"{package}.trainer") + training_mod = importlib.import_module(f"{package}.training") + monkeypatch.setattr( + training_mod._MLXTrainerAdapter, + "_activate_transformers_for_model", + lambda self, model_name, hf_token: None, + ) + return trainer_mod + + +class _ExitedProc: + def join(self, timeout = None): + return None + + def is_alive(self): + return False + + +class _TerminableProc: + def __init__(self): + self.terminated = False + self._done = threading.Event() + + def join(self, timeout = None): + self._done.wait(timeout = timeout or 5) + + def is_alive(self): + return not self.terminated + + def terminate(self): + self.terminated = True + self._done.set() + + +def test_unsloth_trainer_dispatches_for_mlx_and_torch(monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + + mlx_trainer = trainer_mod.UnslothTrainer() + + assert type(mlx_trainer).__module__ == "core.training.training" + assert mlx_trainer.get_training_progress().status_message == "Ready to train" + + trainer_mod = _load_trainer_module(monkeypatch, "torch") + + assert trainer_mod.UnslothTrainer().__class__ is trainer_mod.UnslothTrainer + + +def test_cli_mlx_trainer_activates_before_importing_trainer(): + repo_root = Path(__file__).resolve().parents[3] + script = """ +import json +import sys +import unsloth_cli.commands.train as train_cmd +from studio.backend.core.training import training as training_mod +from utils.hardware import hardware as hw + +training_mod.platform.system = lambda: "Darwin" +training_mod.platform.machine = lambda: "arm64" +hw.DEVICE = None +events = [] + +def fake_activate(model_name, hf_token): + events.append({ + "model_name": model_name, + "trainer_loaded": "studio.backend.core.training.trainer" in sys.modules, + }) + +train_cmd._activate_mlx_transformers = fake_activate +trainer = train_cmd._create_cli_trainer("mlx-community/Qwen3-0.6B-4bit", None) +print(json.dumps({ + "trainer_module": type(trainer).__module__, + "events": events, +})) +""" + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + [str(repo_root), str(repo_root / "studio" / "backend"), env.get("PYTHONPATH", "")] + ) + result = subprocess.run( + [sys.executable, "-c", script], + cwd = repo_root, + env = env, + text = True, + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + check = True, + ) + payload = json.loads(result.stdout) + + assert payload["trainer_module"] == "studio.backend.core.training.training" + assert payload["events"] == [ + {"model_name": "mlx-community/Qwen3-0.6B-4bit", "trainer_loaded": False} + ] + + +def test_mlx_adapter_builds_config_and_reports_completion(tmp_path, monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + captured = {} + + def fake_run_worker(config, event_queue, stop_queue): + captured["config"] = config + event_queue.put({"type": "progress", "step": 1, "total_steps": 1, "loss": 0.25}) + event_queue.put( + {"type": "complete", "status_message": "done", "output_dir": config["output_dir"]} + ) + + trainer = trainer_mod.UnslothTrainer() + monkeypatch.setattr(trainer, "_run_mlx_worker", fake_run_worker) + + assert trainer.load_model("mlx-community/Qwen3-0.6B-4bit", max_seq_length = 1024) + assert trainer.prepare_model_for_training(use_lora = False) + dataset, eval_dataset = trainer.load_and_format_dataset("org/dataset") + output_dir = tmp_path / "mlx-out" + + assert trainer.start_training( + dataset = dataset, + eval_dataset = eval_dataset, + output_dir = output_dir, + project_name = "Sales Assistant", + max_steps = 1, + learning_rate = 3e-4, + ) + trainer.training_thread.join(timeout = 5) + + progress = trainer.get_training_progress() + config = captured["config"] + assert progress.is_completed + assert progress.output_dir == str(output_dir.resolve()) + progress.status_message = "mutated" + assert trainer.get_training_progress().status_message == "done" + assert config["model_name"] == "mlx-community/Qwen3-0.6B-4bit" + assert config["project_name"] == "Sales Assistant" + assert config["hf_dataset"] == "org/dataset" + assert config["training_type"] == "Full Finetuning" + assert config["load_in_4bit"] is False + assert config["max_seq_length"] == 1024 + assert config["learning_rate"] == 3e-4 + assert config["output_dir"] == str(output_dir.resolve()) + assert config["allow_external_output_dir"] is True + + +def test_mlx_worker_helpers_cover_cli_paths(tmp_path, monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training.worker import ( + _resolve_mlx_local_dataset_files, + _resolve_mlx_output_dir, + ) + + dataset = tmp_path / "train.jsonl" + dataset.write_text('{"text":"hello"}\n', encoding = "utf-8") + monkeypatch.chdir(tmp_path) + + assert _resolve_mlx_local_dataset_files(["train.jsonl"]) == [str(dataset)] + assert _resolve_mlx_output_dir( + {"output_dir": "cli-out", "allow_external_output_dir": True}, + "mlx-community/Qwen3-0.6B-4bit", + ) == str((tmp_path / "cli-out").resolve()) + + +def test_run_mlx_training_process_applies_side_effects_before_hardware_detection(monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training import worker + from utils.hardware import hardware as hw + + order = [] + + def fake_activate(model_name, hf_token): + order.append(("activate", model_name, hf_token)) + + def fake_detect_hardware(): + order.append("detect") + hw.DEVICE = hw.DeviceType.CPU + return hw.DEVICE + + monkeypatch.delenv("HF_HUB_DISABLE_XET", raising = False) + monkeypatch.delenv("HF_HUB_ENABLE_HF_TRANSFER", raising = False) + monkeypatch.setattr(worker, "_activate_transformers_version_or_warn", fake_activate) + monkeypatch.setattr(hw, "detect_hardware", fake_detect_hardware) + + event_queue = queue.Queue() + worker.run_mlx_training_process( + event_queue = event_queue, + stop_queue = queue.Queue(), + config = {"model_name": "mlx-community/Gemma-4-12B", "disable_xet": True}, + ) + + event = event_queue.get_nowait() + assert order == [("activate", "mlx-community/Gemma-4-12B", None), "detect"] + assert os.environ["HF_HUB_DISABLE_XET"] == "1" + assert os.environ["HF_HUB_ENABLE_HF_TRANSFER"] == "0" + assert "MLX training requires Apple Silicon" in event["error"] + + if __name__ == "__main__": unittest.main() diff --git a/unsloth-cli.py b/unsloth-cli.py index 756efef0d0..be48893d8d 100644 --- a/unsloth-cli.py +++ b/unsloth-cli.py @@ -22,24 +22,179 @@ import argparse import os +def _is_mlx_backend(unsloth_module): + return bool(getattr(unsloth_module, "_IS_MLX", False)) + + +def _normalize_dtype(dtype, is_mlx): + if is_mlx and isinstance(dtype, str) and dtype.strip().lower() in {"", "none", "auto"}: + return None + return dtype + + +def _prepare_device_map(is_mlx): + if is_mlx: + return None, False + + from unsloth.models.loader_utils import prepare_device_map + return prepare_device_map() + + +class _CallableTokenizerProxy: + def __init__(self, tokenizer): + self._tokenizer = tokenizer + + def __getattr__(self, name): + return getattr(self._tokenizer, name) + + def __call__(self, text, *args, **kwargs): + # MLX/torch-free: never request torch tensors; keep plain python ids. + kwargs.pop("return_tensors", None) + wrapped = getattr(self._tokenizer, "_tokenizer", None) + if callable(wrapped): + return wrapped(text, *args, **kwargs) + + add_special_tokens = kwargs.get("add_special_tokens", False) + input_ids = self._tokenizer.encode(text, add_special_tokens = add_special_tokens) + return {"input_ids": input_ids} + + +def _tokenizer_for_raw_text_loader(tokenizer, is_mlx): + if not is_mlx or callable(tokenizer): + return tokenizer + return _CallableTokenizerProxy(tokenizer) + + +def _raw_text_loader_for_backend( + RawTextDataLoader, + tokenizer, + is_mlx, + chunk_size = 2048, + stride = 512, +): + return RawTextDataLoader( + _tokenizer_for_raw_text_loader(tokenizer, is_mlx), + chunk_size, + stride, + return_tokenized = not is_mlx, + ) + + +def _train_with_legacy_save_control(trainer, is_mlx): + if not is_mlx: + return trainer.train() + + original_save_model = getattr(trainer, "save_model", None) + if original_save_model is None: + return trainer.train() + + def skip_internal_final_save(*args, **kwargs): + raise ValueError("legacy unsloth-cli.py owns final save") + + trainer.save_model = skip_internal_final_save + try: + return trainer.train() + finally: + trainer.save_model = original_save_model + + +def _iter_quantization_methods(quantization): + if isinstance(quantization, list): + return quantization + return [quantization] + + +def _save_or_push_model(model, tokenizer, args, is_mlx): + if not args.save_model: + print("Warning: The model is not saved!") + return + + # Enter the GGUF branch when saving or pushing GGUF, so --push_gguf works + # without --save_gguf (the local save is guarded separately below). + if args.save_gguf or args.push_gguf: + if not args.save_gguf: + print("Warning: --save_gguf not set, pushing GGUF to hub without saving locally.") + for quantization_method in _iter_quantization_methods(args.quantization): + if args.save_gguf: + print(f"Saving model with quantization method: {quantization_method}") + model.save_pretrained_gguf( + args.save_path, + tokenizer, + quantization_method = quantization_method, + ) + if args.push_model or args.push_gguf: + model.push_to_hub_gguf( + args.hub_path, + tokenizer, + quantization_method = quantization_method, + token = args.hub_token, + ) + return + + if is_mlx: + model.save_pretrained_merged( + args.save_path, + tokenizer, + save_method = args.save_method, + push_to_hub = args.push_model, + repo_id = args.hub_path if args.push_model else None, + token = args.hub_token, + ) + return + + model.save_pretrained_merged(args.save_path, tokenizer, save_method = args.save_method) + if args.push_model: + model.push_to_hub_merged(args.hub_path, tokenizer, args.save_method, token = args.hub_token) + + +def _build_sft_config(SFTConfig, args, is_mlx, bf16_supported): + config_kwargs = dict( + per_device_train_batch_size = args.per_device_train_batch_size, + gradient_accumulation_steps = args.gradient_accumulation_steps, + warmup_steps = args.warmup_steps, + max_steps = args.max_steps, + learning_rate = args.learning_rate, + fp16 = not bf16_supported, + bf16 = bf16_supported, + logging_steps = args.logging_steps, + optim = args.optim, + weight_decay = args.weight_decay, + lr_scheduler_type = args.lr_scheduler_type, + seed = args.seed, + output_dir = args.output_dir, + report_to = args.report_to, + max_length = args.max_seq_length, + dataset_num_proc = 2, + packing = args.packing, + ) + if is_mlx: + if args.per_device_eval_batch_size != 4: + print("Warning: --per_device_eval_batch_size is ignored on MLX without eval data.") + else: + config_kwargs["per_device_eval_batch_size"] = args.per_device_eval_batch_size + return SFTConfig(**config_kwargs) + + def run(args): + import unsloth from unsloth import FastLanguageModel from datasets import load_dataset from transformers.utils import strtobool from trl import SFTTrainer, SFTConfig from unsloth import is_bfloat16_supported - from unsloth.models.loader_utils import prepare_device_map import logging from unsloth import RawTextDataLoader logging.getLogger("hf-to-gguf").setLevel(logging.WARNING) + is_mlx = _is_mlx_backend(unsloth) + # Load model and tokenizer - device_map, distributed = prepare_device_map() + device_map, distributed = _prepare_device_map(is_mlx) model, tokenizer = FastLanguageModel.from_pretrained( model_name = args.model_name, max_seq_length = args.max_seq_length, - dtype = args.dtype, + dtype = _normalize_dtype(args.dtype, is_mlx), load_in_4bit = args.load_in_4bit, device_map = device_map, ) @@ -92,11 +247,13 @@ def run(args): def load_dataset_smart(args): from transformers.utils import strtobool if args.raw_text_file: - loader = RawTextDataLoader(tokenizer, args.chunk_size, args.stride) + loader = _raw_text_loader_for_backend( + RawTextDataLoader, tokenizer, is_mlx, args.chunk_size, args.stride + ) dataset = loader.load_from_file(args.raw_text_file) elif args.dataset.endswith((".txt", ".md", ".json", ".jsonl")): # Auto-detect local raw text files - loader = RawTextDataLoader(tokenizer) + loader = _raw_text_loader_for_backend(RawTextDataLoader, tokenizer, is_mlx) dataset = loader.load_from_file(args.dataset) else: use_modelscope = strtobool(os.environ.get("UNSLOTH_USE_MODELSCOPE", "False")) @@ -115,27 +272,9 @@ def run(args): print("Data is formatted and ready!") # Configure training arguments - training_args = SFTConfig( - per_device_train_batch_size = args.per_device_train_batch_size, - per_device_eval_batch_size = args.per_device_eval_batch_size, - gradient_accumulation_steps = args.gradient_accumulation_steps, - warmup_steps = args.warmup_steps, - max_steps = args.max_steps, - learning_rate = args.learning_rate, - fp16 = not is_bfloat16_supported(), - bf16 = is_bfloat16_supported(), - logging_steps = args.logging_steps, - optim = args.optim, - weight_decay = args.weight_decay, - lr_scheduler_type = args.lr_scheduler_type, - seed = args.seed, - output_dir = args.output_dir, - report_to = args.report_to, - max_length = args.max_seq_length, - dataset_num_proc = 2, - ddp_find_unused_parameters = False if distributed else None, - packing = args.packing, - ) + training_args = _build_sft_config(SFTConfig, args, is_mlx, is_bfloat16_supported()) + if distributed: + training_args.ddp_find_unused_parameters = False # Initialize trainer trainer = SFTTrainer( @@ -145,57 +284,9 @@ def run(args): args = training_args, ) - trainer.train() + _train_with_legacy_save_control(trainer, is_mlx) - # Save model - if args.save_model: - # If args.quantization is a list, save once per quantization method - # Enter the GGUF branch when saving *or* pushing GGUF, so --push_gguf - # works even when --save_gguf is omitted (the local save is guarded - # separately below). - if args.save_gguf or args.push_gguf: - # Push-only GGUF (no --save_gguf) skips the local save; warn so it is not silent. - if not args.save_gguf: - print("Warning: --save_gguf not set, pushing GGUF to hub without saving locally.") - if isinstance(args.quantization, list): - for quantization_method in args.quantization: - if args.save_gguf: - print(f"Saving model with quantization method: {quantization_method}") - model.save_pretrained_gguf( - args.save_path, - tokenizer, - quantization_method = quantization_method, - ) - if args.push_model or args.push_gguf: - model.push_to_hub_gguf( - args.hub_path, - tokenizer, - quantization_method = quantization_method, - token = args.hub_token, - ) - else: - if args.save_gguf: - print(f"Saving model with quantization method: {args.quantization}") - model.save_pretrained_gguf( - args.save_path, - tokenizer, - quantization_method = args.quantization, - ) - if args.push_model or args.push_gguf: - model.push_to_hub_gguf( - args.hub_path, - tokenizer, - quantization_method = args.quantization, - token = args.hub_token, - ) - else: - model.save_pretrained_merged(args.save_path, tokenizer, args.save_method) - if args.push_model: - model.push_to_hub_merged( - args.hub_path, tokenizer, args.save_method, token = args.hub_token - ) - else: - print("Warning: The model is not saved!") + _save_or_push_model(model, tokenizer, args, is_mlx) if __name__ == "__main__": diff --git a/unsloth_cli/commands/train.py b/unsloth_cli/commands/train.py index 9d47a574d5..c52c2344b7 100644 --- a/unsloth_cli/commands/train.py +++ b/unsloth_cli/commands/train.py @@ -7,10 +7,42 @@ from typing import Optional import typer +from unsloth_cli._inference import ensure_studio_backend_path from unsloth_cli.config import Config, load_config from unsloth_cli.options import add_options_from_config +def _should_use_mlx_backend_for_cli() -> bool: + ensure_studio_backend_path() + from studio.backend.core.training.training import should_use_mlx_training_backend + return should_use_mlx_training_backend() + + +def _activate_mlx_transformers(model_name: str, hf_token: Optional[str]) -> None: + # Activate before any transformers import: adapter model-type detection imports utils.models. + ensure_studio_backend_path() + from utils.transformers_version import activate_transformers_for_subprocess + try: + activate_transformers_for_subprocess(model_name, hf_token) + except Exception as exc: + typer.echo(f"Warning: failed to activate Transformers sidecar: {exc}", err = True) + + +def _create_cli_trainer(model_name: str, hf_token: Optional[str]): + if _should_use_mlx_backend_for_cli(): + _activate_mlx_transformers(model_name, hf_token) + # MLX is torch-free: use the lightweight adapter, not trainer.py (imports torch/unsloth/trl at load). + ensure_studio_backend_path() + from studio.backend.core.training.training import create_mlx_trainer_adapter + + return create_mlx_trainer_adapter() + + ensure_studio_backend_path() + from studio.backend.core.training.trainer import UnslothTrainer + + return UnslothTrainer() + + @add_options_from_config(Config) def train( config: Optional[Path] = typer.Option( @@ -39,6 +71,7 @@ def train( typer.echo(f"Error: {e}", err = True) raise typer.Exit(code = 2) + config_overrides = config_overrides or {} cfg.apply_overrides(**config_overrides) # CLI/env tokens take precedence; guard against unresolved typer.Option @@ -83,9 +116,7 @@ def train( ) raise typer.Exit(code = 2) - from studio.backend.core.training.trainer import UnslothTrainer - - trainer = UnslothTrainer() + trainer = _create_cli_trainer(cfg.model, hf_token) # Load model (trainer.is_vlm is set after this) if not trainer.load_model( @@ -124,13 +155,20 @@ def train( try: while trainer.training_thread and trainer.training_thread.is_alive(): + progress = trainer.get_training_progress() + if getattr(progress, "error", None): + break time.sleep(1) except KeyboardInterrupt: typer.echo("Stopping training (Ctrl+C detected)...") trainer.stop_training() finally: if trainer.training_thread: - trainer.training_thread.join() + progress = trainer.get_training_progress() + if getattr(progress, "error", None): + trainer.training_thread.join(timeout = 5) + else: + trainer.training_thread.join() final = trainer.get_training_progress() if getattr(final, "error", None):