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 <danielhanchen@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Long Yixing
GitHub
parent
de60a3a994
commit
38dacb8a1f
6 changed files with 1162 additions and 262 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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-<step>/`` partials after a cancel.
|
||||
|
||||
Completed ``checkpoint-<int>/`` 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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
243
unsloth-cli.py
243
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__":
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue