# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os, importlib.util, platform os.environ["UNSLOTH_IS_PRESENT"] = "1" # ── Windows console UTF-8 safety ───────────────────────────────────────────── # Legacy Windows consoles (cp1252) can't encode Unsloth's emoji/box-drawing # glyphs and crash with UnicodeEncodeError. Force stdout/stderr to UTF-8 only on # Windows and only when not already UTF-8; no-op elsewhere. errors="replace" # guarantees we never crash on an unencodable glyph. if platform.system() == "Windows": import sys as _sys for _name in ("stdout", "stderr"): _s = getattr(_sys, _name, None) try: _enc = (getattr(_s, "encoding", None) or "").lower() if _s is not None and hasattr(_s, "reconfigure") and "utf" not in _enc: _s.reconfigure(encoding = "utf-8", errors = "replace") except Exception: pass class _UnslothDeviceStats: """Portable device metadata used by backend memory-reporting helpers.""" def __init__( self, name, total_memory = 0, ): """Store a display name and total memory in bytes.""" self.name = name self.total_memory = int(total_memory or 0) self.major = 0 self.minor = 0 self.multi_processor_count = 0 def _bytes_to_gb(value): """Convert byte counts to GiB rounded""" return round(float(value or 0) / 1024 / 1024 / 1024, 3) def _is_mlx_available(): # Transitional import barrier: keep non-Apple-Silicon imports from touching # unsloth_zoo until unsloth_zoo.mlx is import-safe on GPU hosts. Then this # can collapse back to the centralized zoo runtime call below. if ( os.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") == "1" or platform.system() != "Darwin" or platform.machine() != "arm64" or importlib.util.find_spec("mlx") is None ): return False try: from unsloth_zoo.mlx import is_mlx_available except ImportError: return False return is_mlx_available() # Detect Apple Silicon + MLX before any torch/numpy imports _IS_MLX = _is_mlx_available() if _IS_MLX: try: import unsloth_zoo except ImportError as _e: raise ImportError( "Unsloth: MLX support requires `unsloth-zoo` with MLX modules. " "Reinstall with `pip install unsloth-zoo` or rerun install.sh." ) from _e # An older unsloth-zoo satisfies `import unsloth_zoo` but lacks the # mlx.trainer / mlx.loader submodules. Surface a friendly install hint # instead of a raw ImportError on the submodule path. try: from unsloth_zoo.mlx.trainer import ( MLXTrainer, MLXTrainingConfig, _is_vlm_model, _normalize_mlx_optimizer_name, ) from unsloth_zoo.mlx.loader import FastMLXModel except ImportError as _e: raise ImportError( "Unsloth: MLX support requires an unsloth-zoo build that includes " "`unsloth_zoo.mlx.trainer` and `unsloth_zoo.mlx.loader`. Upgrade with " "`pip install -U unsloth-zoo` or rerun install.sh." ) from _e import dataclasses as _dataclasses import inspect as _inspect import importlib.machinery as _machinery import sys as _sys import types as _types import warnings as _warnings __version__ = unsloth_zoo.__version__ DEVICE_TYPE = "mlx" _MLX_TRAINER_ACCEPTS_VAR_KWARGS = False _MLX_TRAINER_SUPPORTED_KWARGS = frozenset() try: _MLX_TRAINER_INIT_PARAMETERS = _inspect.signature(MLXTrainer.__init__).parameters _MLX_TRAINER_ACCEPTS_VAR_KWARGS = any( param.kind is _inspect.Parameter.VAR_KEYWORD for param in _MLX_TRAINER_INIT_PARAMETERS.values() ) _MLX_TRAINER_SUPPORTED_KWARGS = frozenset( name for name, param in _MLX_TRAINER_INIT_PARAMETERS.items() if name != "self" and param.kind in ( _inspect.Parameter.POSITIONAL_OR_KEYWORD, _inspect.Parameter.KEYWORD_ONLY, ) ) except (TypeError, ValueError): pass def _mlx_trainer_supports_kwarg(name): """Return whether the installed zoo MLXTrainer accepts a kwarg.""" return _MLX_TRAINER_ACCEPTS_VAR_KWARGS or name in _MLX_TRAINER_SUPPORTED_KWARGS def _is_mlx_cuda_device_target(device): """Return True when a torch .to/.cuda target asks for CUDA on MLX.""" if device is None: return False return str(device).lower().startswith("cuda") def _patch_mlx_batch_encoding_to_cuda(): """Treat tokenizer_output.to("cuda") as a no-op on the MLX backend.""" try: from transformers.tokenization_utils_base import BatchEncoding except Exception: return original_to = getattr(BatchEncoding, "to", None) if original_to is None or getattr(original_to, "_unsloth_mlx_cuda_noop", False): return def batch_encoding_to( self, device = None, *args, **kwargs, ): target = kwargs.get("device", device) if _is_mlx_cuda_device_target(target): return self # device given by keyword: don't also pass the positional None, or the # original raises "multiple values for 'device'" (e.g. .to(device="cpu")). if "device" in kwargs: return original_to(self, *args, **kwargs) return original_to(self, device, *args, **kwargs) batch_encoding_to._unsloth_mlx_cuda_noop = True batch_encoding_to._unsloth_original_to = original_to BatchEncoding.to = batch_encoding_to _patch_mlx_batch_encoding_to_cuda() # Load raw_text helpers without executing dataprep/__init__.py, which # imports synthetic.py -> torch and would defeat the torch-free MLX path. from pathlib import Path as _Path _raw_text_path = _Path(__file__).resolve().parent / "dataprep" / "raw_text.py" _raw_text_spec = importlib.util.spec_from_file_location("unsloth._mlx_raw_text", _raw_text_path) if _raw_text_spec is None or _raw_text_spec.loader is None: raise ImportError("Unsloth: could not load MLX raw_text dataprep helpers.") _raw_text = importlib.util.module_from_spec(_raw_text_spec) _raw_text_spec.loader.exec_module(_raw_text) RawTextDataLoader = _raw_text.RawTextDataLoader TextPreprocessor = _raw_text.TextPreprocessor del _raw_text, _raw_text_spec, _raw_text_path, _Path class FastLanguageModel: @staticmethod def from_pretrained(*args, **kwargs): return FastMLXModel.from_pretrained(*args, **kwargs) @staticmethod def get_peft_model(*args, **kwargs): return FastMLXModel.get_peft_model(*args, **kwargs) @staticmethod def for_inference(*args, **kwargs): return args[0] if args else None class FastVisionModel(FastLanguageModel): @staticmethod def from_pretrained(*args, **kwargs): kwargs.setdefault("text_only", False) return FastMLXModel.from_pretrained(*args, **kwargs) @staticmethod def for_training(*args, **kwargs): return args[0] if args else None FastTextModel = FastLanguageModel FastModel = FastLanguageModel class FastSentenceTransformer: @staticmethod def from_pretrained(*args, **kwargs): raise NotImplementedError( "Unsloth: FastSentenceTransformer is not yet supported on MLX." ) @staticmethod def get_peft_model(*args, **kwargs): raise NotImplementedError( "Unsloth: FastSentenceTransformer is not yet supported on MLX." ) def is_bfloat16_supported(): try: import mlx.core as mx name = mx.device_info().get("device_name", "") or "" return not name.startswith(("Apple M1", "Apple M2")) except Exception: return True is_bf16_supported = is_bfloat16_supported def get_gpu_memory_stats(): """Return MLX device stats, peak memory, and total memory in GiB.""" import mlx.core as mx info = mx.device_info() total = info.get("memory_size") or info.get("max_recommended_working_set_size") or 0 get_peak_memory = getattr(mx, "get_peak_memory", None) if get_peak_memory is None and hasattr(mx, "metal"): get_peak_memory = getattr(mx.metal, "get_peak_memory", None) peak = get_peak_memory() if callable(get_peak_memory) else 0 stats = _UnslothDeviceStats(info.get("device_name", "Apple GPU"), total) max_memory = _bytes_to_gb(total) or 1.0 return stats, _bytes_to_gb(peak), max_memory def clear_gpu_memory(): """Clear MLX's cached GPU memory for compatibility cleanup helpers.""" import mlx.core as mx clear_cache = getattr(mx, "clear_cache", None) if clear_cache is None and hasattr(mx, "metal"): clear_cache = getattr(mx.metal, "clear_cache", None) if callable(clear_cache): clear_cache() def _patch_mlx_torch_cuda_compat_api(): """Expose CUDA-shaped torch helpers for compatibility callers on MLX.""" try: import torch except Exception: return cuda = getattr(torch, "cuda", None) if cuda is not None and not getattr(cuda, "_unsloth_mlx_cuda_compat_api", False): def get_device_properties(device = None): """Return MLX device stats through torch.cuda's compatibility API.""" return get_gpu_memory_stats()[0] def get_device_name(device = None): """Return the MLX device name through torch.cuda's compatibility API.""" return get_device_properties(device).name def max_memory_reserved(device = None): """Return MLX peak memory in bytes for torch.cuda compatibility API.""" return int(get_gpu_memory_stats()[1] * 1024 * 1024 * 1024) def empty_cache(): """Clear MLX cache through torch.cuda.empty_cache().""" clear_gpu_memory() def _mlx_active_memory_bytes(): """Current active MLX memory in bytes (not the peak high-water mark).""" import mlx.core as mx get_active = getattr(mx, "get_active_memory", None) if get_active is None and hasattr(mx, "metal"): get_active = getattr(mx.metal, "get_active_memory", None) return int(get_active()) if callable(get_active) else 0 def memory_current(device = None): """Return CURRENT MLX memory in bytes. torch.cuda.memory_reserved / memory_allocated report live usage, not the peak (that is max_*).""" return _mlx_active_memory_bytes() def mem_get_info(device = None): """Return (free, total) bytes for torch.cuda compatibility API. Free uses CURRENT active memory, not the peak high-water mark, so a capacity check stays accurate after a transient spike.""" total = int(get_gpu_memory_stats()[2] * 1024 * 1024 * 1024) return (max(total - _mlx_active_memory_bytes(), 0), total) def reset_peak_memory_stats(device = None): """Reset MLX's peak-memory counter so a later max_memory_reserved / max_memory_allocated scopes to the run, not earlier model-load peaks.""" import mlx.core as mx reset = getattr(mx, "reset_peak_memory", None) if reset is None and hasattr(mx, "metal"): reset = getattr(mx.metal, "reset_peak_memory", None) if callable(reset): reset() def synchronize(device = None): """Wait for queued MLX work when torch.cuda.synchronize() is called.""" import mlx.core as mx sync = getattr(mx, "synchronize", None) if callable(sync): sync() cuda.get_device_properties = get_device_properties cuda.get_device_name = get_device_name cuda.max_memory_reserved = max_memory_reserved cuda.max_memory_allocated = max_memory_reserved cuda.memory_reserved = memory_current cuda.memory_allocated = memory_current cuda.empty_cache = empty_cache cuda.mem_get_info = mem_get_info cuda.reset_peak_memory_stats = reset_peak_memory_stats cuda.synchronize = synchronize cuda.current_device = lambda: 0 cuda.device_count = lambda: 1 cuda.set_device = lambda device = None: None cuda.get_device_capability = lambda device = None: (0, 0) cuda.is_bf16_supported = lambda *args, **kwargs: is_bfloat16_supported() cuda._unsloth_mlx_cuda_compat_api = True tensor_to = getattr(torch.Tensor, "to", None) if tensor_to is not None and not getattr(tensor_to, "_unsloth_mlx_cuda_noop", False): def _coerce_mlx_dtype_to_torch(value): """Map MLX dtype objects to their torch dtype equivalents.""" try: import mlx.core as mx except Exception: return value dtype_map = { mx.bool_: torch.bool, mx.int8: torch.int8, mx.int16: torch.int16, mx.int32: torch.int32, mx.int64: torch.int64, mx.uint8: torch.uint8, mx.float16: torch.float16, mx.float32: torch.float32, mx.bfloat16: torch.bfloat16, } mapped = dtype_map.get(value, None) if mapped is not None: return mapped dtype_name = str(value).rsplit(".", 1)[-1] name_map = { "bool_": torch.bool, "int8": torch.int8, "int16": torch.int16, "int32": torch.int32, "int64": torch.int64, "uint8": torch.uint8, "float16": torch.float16, "float32": torch.float32, "bfloat16": torch.bfloat16, } return name_map.get(dtype_name, value) def mlx_tensor_to(self, *args, **kwargs): """Ignore CUDA device targets while preserving dtype conversions.""" args = list(args) kwargs = dict(kwargs) removed_cuda_device = False if args and _is_mlx_cuda_device_target(args[0]): args.pop(0) removed_cuda_device = True if _is_mlx_cuda_device_target(kwargs.get("device", None)): kwargs.pop("device", None) removed_cuda_device = True if removed_cuda_device and not args: cuda_only_kwargs = ("non_blocking", "copy", "memory_format") if all(key in cuda_only_kwargs for key in kwargs): return self if removed_cuda_device and not args and not kwargs: return self if args: args[0] = _coerce_mlx_dtype_to_torch(args[0]) if "dtype" in kwargs: kwargs["dtype"] = _coerce_mlx_dtype_to_torch(kwargs["dtype"]) return tensor_to(self, *args, **kwargs) mlx_tensor_to._unsloth_mlx_cuda_noop = True mlx_tensor_to._unsloth_original_to = tensor_to torch.Tensor.to = mlx_tensor_to tensor_cuda = getattr(torch.Tensor, "cuda", None) if tensor_cuda is not None and not getattr(tensor_cuda, "_unsloth_mlx_cuda_noop", False): def mlx_tensor_cuda(self, *args, **kwargs): """Treat tensor.cuda() as a no-op on MLX.""" return self mlx_tensor_cuda._unsloth_mlx_cuda_noop = True mlx_tensor_cuda._unsloth_original_cuda = tensor_cuda torch.Tensor.cuda = mlx_tensor_cuda _patch_mlx_torch_cuda_compat_api() _MLX_TRAINING_CONFIG_FIELDS = {_field.name for _field in _dataclasses.fields(MLXTrainingConfig)} _MLX_TRAINING_ARGUMENT_ALIASES = { "max_length": "max_seq_length", } _MLX_COMPAT_EXTRA_ARGUMENTS = frozenset( ( "bf16", "dataloader_num_workers", "dataloader_pin_memory", "dataset_kwargs", "ddp_find_unused_parameters", "disable_tqdm", "eval_strategy", "evaluation_strategy", "fp16", "full_determinism", "gradient_checkpointing_kwargs", "hub_model_id", "hub_token", "log_level", "logging_strategy", "neftune_noise_alpha", "optim_args", "padding_free", "push_to_hub", "remove_unused_columns", "save_on_each_node", "save_safetensors", "save_strategy", "torch_compile", ) ) _MLX_IMPLEMENTED_EXTRA_ARGUMENTS = frozenset( ( "image_size", "preserve_dataset_order", "warmup_ratio", ) ) _MLX_ALLOWED_EXTRA_ARGUMENTS = _MLX_COMPAT_EXTRA_ARGUMENTS | _MLX_IMPLEMENTED_EXTRA_ARGUMENTS _MLX_UNSUPPORTED_TASK_ARGUMENTS = frozenset( ( "assistant_only_loss", "completion_only_loss", ) ) def _is_mlx_no_save_strategy(value): if hasattr(value, "value"): value = value.value strategy = str(value or "").strip().lower() strategy = strategy.rsplit(".", 1)[-1] return strategy in ("no", "none", "false") _MLX_ADAMW_OPTIMIZER_ALIASES = frozenset( ( "adamw_8bit", "paged_adamw_8bit", "adamw_bnb_8bit", "paged_adamw_32bit", "adamw_torch", "adamw_torch_fused", "paged_adamw", "adamw_32bit", "adamw_hf", "adamw_anyprecision", "adamw_apex_fused", ) ) def _normalize_mlx_training_value(key, value): if key == "eval_steps" and value is None: return 0 if key == "num_train_epochs" and value is not None and not isinstance(value, bool): try: epochs = float(value) except (TypeError, ValueError): pass else: if epochs.is_integer(): return int(epochs) if key == "lr_scheduler_type" and hasattr(value, "value"): return value.value if key != "optim": return value try: return _normalize_mlx_optimizer_name(value) except ValueError: # Older unsloth-zoo lacks CUDA/TRL optimizer aliases; map common # adamw_* names so notebook defaults (optim="adamw_8bit") still work. opt = str(getattr(value, "value", value) or "adamw").strip().lower() opt = opt.rsplit(".", 1)[-1].replace("-", "_") if opt in _MLX_ADAMW_OPTIMIZER_ALIASES: return "adamw" raise def _mlx_training_argument_values(args): values = {} for field in _dataclasses.fields(MLXTrainingConfig): if hasattr(args, field.name): values[field.name] = _normalize_mlx_training_value( field.name, getattr(args, field.name), ) for alias, target in _MLX_TRAINING_ARGUMENT_ALIASES.items(): if target not in values and hasattr(args, alias): values[target if target in _MLX_ALLOWED_EXTRA_ARGUMENTS else alias] = getattr( args, alias ) for name in _MLX_ALLOWED_EXTRA_ARGUMENTS: if hasattr(args, name): values[name] = getattr(args, name) for name in _MLX_UNSUPPORTED_TASK_ARGUMENTS: if hasattr(args, name): value = getattr(args, name) if ( name == "completion_only_loss" and value is not None and name in _MLX_TRAINING_CONFIG_FIELDS ): values[name] = value elif value is not None and value is not False: values[name] = value if _is_mlx_no_save_strategy(values.get("save_strategy", None)): values["save_steps"] = 0 return values def _split_mlx_trainer_kwargs(kwargs): trainer_kwargs = {} config_kwargs = {} ignored_kwargs = {} for key, value in kwargs.items(): if key in _MLX_TRAINER_KWARGS: trainer_kwargs[key] = value continue target = _MLX_TRAINING_ARGUMENT_ALIASES.get(key, key) if target in _MLX_TRAINING_CONFIG_FIELDS or key in _MLX_ALLOWED_EXTRA_ARGUMENTS: config_kwargs[key] = value else: ignored_kwargs[key] = value return trainer_kwargs, config_kwargs, ignored_kwargs def _is_mlx_training_args_like(value): if isinstance(value, (MLXTrainingConfig, dict, str, os.PathLike)): return True return any( hasattr(value, name) for name in ( "output_dir", "per_device_train_batch_size", "gradient_accumulation_steps", "max_steps", "learning_rate", ) ) def _should_use_trl_positional_schema(args): if len(args) < 2: return False if _is_mlx_training_args_like(args[1]): return True # TRL callers often pass explicit defaults: # SFTTrainer(model, None, None, train_dataset, ...) return len(args) >= 3 and args[1] is None and (args[2] is None or callable(args[2])) def _assign_mlx_positional_kwarg(kwargs, name, value): if name in kwargs: raise TypeError( f"UnslothTrainer.__init__() got multiple values for argument " f"{name!r}" ) kwargs[name] = value def _normalize_mlx_trainer_init_args(args, kwargs): kwargs = dict(kwargs) if len(args) == 0: return kwargs use_trl_schema = _should_use_trl_positional_schema(args) positional_names = ( _TRL_SFT_TRAINER_POSITIONAL_KWARGS if use_trl_schema else _MLX_TRAINER_POSITIONAL_KWARGS ) if len(args) > len(positional_names): raise TypeError( f"UnslothTrainer.__init__() takes at most " f"{len(positional_names)} positional arguments on MLX " f"({len(args)} given)" ) for name, value in zip(positional_names, args): _assign_mlx_positional_kwarg(kwargs, name, value) return kwargs def _is_meaningful_mlx_extra_value(value): if value is None or value is False: return False if isinstance(value, (str, bytes)) and len(value) == 0: return False if isinstance(value, (dict, list, tuple, set, frozenset)) and len(value) == 0: return False return True def _warn_ignored_mlx_training_args(extra_kwargs): names = sorted( key for key, value in extra_kwargs.items() if (key in _MLX_COMPAT_EXTRA_ARGUMENTS and _is_meaningful_mlx_extra_value(value)) ) if not names: return _warnings.warn( "Unsloth MLX: accepting but not applying unsupported " "TrainingArguments kwargs: " f"{', '.join(names)}. These options are not implemented by " "MLXTrainer yet.", RuntimeWarning, stacklevel = 3, ) def _is_meaningful_mlx_trainer_kwarg(key, value): if key == "optimizers" and value == (None, None): return False return _is_meaningful_mlx_extra_value(value) def _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs): names = sorted( key for key, value in ignored_kwargs.items() if _is_meaningful_mlx_trainer_kwarg(key, value) ) if not names: return raise NotImplementedError( "Unsloth MLX: unsupported SFTTrainer kwargs cannot be ignored safely: " f"{', '.join(names)}. Remove these kwargs or use a supported MLX " "trainer configuration." ) def _raise_unknown_mlx_training_args(extra_kwargs): names = sorted(key for key in extra_kwargs if key not in _MLX_ALLOWED_EXTRA_ARGUMENTS) if not names: return raise NotImplementedError( "Unsloth MLX: unsupported TrainingArguments/SFTConfig kwargs: " f"{', '.join(names)}. Remove these kwargs or use fields implemented " "by MLXTrainingConfig." ) def _positive_mlx_context_length(value): if value is None or isinstance(value, bool): return None try: length = int(value) except (TypeError, ValueError, OverflowError): return None if length <= 0: return None return length def _positive_mlx_training_number(value): if value is None or isinstance(value, bool): return None try: number = float(value) except (TypeError, ValueError, OverflowError): return None if number <= 0: return None return number def _set_mlx_cuda_style_context_length(args, length): args.max_seq_length = length args.max_length = length args._unsloth_mlx_max_length_value = length return args class UnslothTrainingArguments(MLXTrainingConfig): """MLX-compatible public training arguments for Unsloth notebooks.""" def __init__(self, *args, **kwargs): if len(args) == 1 and isinstance(args[0], dict): kwargs = {**args[0], **kwargs} elif len(args) == 1 and isinstance(args[0], (str, os.PathLike)): kwargs = {"output_dir": os.fspath(args[0]), **kwargs} elif args: raise TypeError( "UnslothTrainingArguments on MLX accepts keyword arguments, " "a dict, or a single positional output_dir." ) max_length_value = kwargs.get("max_length", None) # Only the canonical max_seq_length marks context length explicit; TRL # max_length stays a compatibility alias and defers to the model's # context length when one is available. max_seq_length_explicit = ( _positive_mlx_context_length(kwargs.get("max_seq_length", None)) is not None ) if "max_length" in kwargs and "max_seq_length" not in kwargs: kwargs["max_seq_length"] = kwargs["max_length"] elif ( "max_length" in kwargs and _positive_mlx_context_length(kwargs.get("max_seq_length", None)) is not None ): max_length_value = kwargs["max_seq_length"] if "num_train_epochs" in kwargs and "max_steps" not in kwargs: kwargs["max_steps"] = -1 dataset_order_explicit = "dataset_order" in kwargs or bool( kwargs.get("preserve_dataset_order", False) ) append_eos_explicit = "append_eos" in kwargs grad_clip_explicit = any( name in kwargs for name in ("max_grad_norm", "max_grad_value", "max_grad_leaf_norm") ) warmup_ratio = kwargs.get("warmup_ratio", None) warmup_steps_supplied = "warmup_steps" in kwargs warmup_steps_value = kwargs.get("warmup_steps", None) warmup_steps_explicit = False if warmup_steps_supplied: try: warmup_steps_explicit = int(warmup_steps_value) > 0 except (TypeError, ValueError): warmup_steps_explicit = True filtered_kwargs = {} extra_kwargs = {} for key, value in kwargs.items(): target = _MLX_TRAINING_ARGUMENT_ALIASES.get(key, key) if key != target and target in kwargs: continue value = _normalize_mlx_training_value(target, value) if target in _MLX_UNSUPPORTED_TASK_ARGUMENTS: if ( target == "completion_only_loss" and value is not None and target in _MLX_TRAINING_CONFIG_FIELDS ): filtered_kwargs[target] = value elif _is_meaningful_mlx_extra_value(value): extra_kwargs[key] = value continue if target in _MLX_TRAINING_CONFIG_FIELDS: filtered_kwargs[target] = value else: extra_kwargs[target if target in _MLX_ALLOWED_EXTRA_ARGUMENTS else key] = value _raise_unknown_mlx_training_args(extra_kwargs) if _is_mlx_no_save_strategy(extra_kwargs.get("save_strategy", None)): filtered_kwargs["save_steps"] = 0 if warmup_ratio is not None and not warmup_steps_explicit: import math as _math max_steps = filtered_kwargs.get( "max_steps", getattr(MLXTrainingConfig, "max_steps", 60), ) try: if int(max_steps) > 0: filtered_kwargs["warmup_steps"] = max( 0, _math.ceil(int(max_steps) * float(warmup_ratio)), ) except (TypeError, ValueError): pass super().__init__(**filtered_kwargs) self._unsloth_mlx_dataset_order_explicit = dataset_order_explicit self._unsloth_mlx_append_eos_explicit = append_eos_explicit self._unsloth_mlx_max_seq_length_explicit = max_seq_length_explicit self._unsloth_mlx_max_length_value = max_length_value if "max_length" in kwargs: self.max_length = max_length_value self._unsloth_mlx_grad_clip_explicit = grad_clip_explicit self._unsloth_mlx_warmup_steps_explicit = warmup_steps_explicit self._unsloth_mlx_extra_args = extra_kwargs for key, value in extra_kwargs.items(): setattr(self, key, value) _warn_ignored_mlx_training_args(extra_kwargs) def _resolve_mlx_cuda_style_max_seq_length(args, model = None): model_max_seq_length = _positive_mlx_context_length( getattr(model, "max_seq_length", None), ) args_max_seq_length = _positive_mlx_context_length( getattr(args, "max_seq_length", None), ) args_max_seq_length_explicit = getattr( args, "_unsloth_mlx_max_seq_length_explicit", None, ) if args_max_seq_length_explicit is None: default_max_seq_length = getattr(MLXTrainingConfig, "max_seq_length", 2048) args_max_seq_length_explicit = ( args_max_seq_length is not None and args_max_seq_length != default_max_seq_length ) if not args_max_seq_length_explicit: args_max_seq_length = None if args_max_seq_length is None and model_max_seq_length is not None: args_max_seq_length = model_max_seq_length elif ( args_max_seq_length is not None and model_max_seq_length is not None and args_max_seq_length > model_max_seq_length ): print( "Unsloth: You set `max_seq_length` as " f"{args_max_seq_length} but the maximum the model supports is " f"{model_max_seq_length}. We shall reduce it." ) args_max_seq_length = model_max_seq_length if args_max_seq_length is not None: _set_mlx_cuda_style_context_length(args, args_max_seq_length) return args model_max_length = model_max_seq_length if model_max_length is None: model_max_length = _positive_mlx_context_length( getattr(model, "max_length", None), ) if model_max_length is not None: _set_mlx_cuda_style_context_length(args, model_max_length) return args args_max_length = _positive_mlx_context_length( getattr(args, "max_length", None), ) if args_max_length is None: args_max_length = _positive_mlx_context_length( getattr(args, "_unsloth_mlx_max_length_value", None), ) if args_max_length is not None: _set_mlx_cuda_style_context_length(args, args_max_length) if model is not None: setattr(model, "max_seq_length", args_max_length) return args _set_mlx_cuda_style_context_length(args, 1024) return args def _apply_unsloth_trainer_mlx_defaults( args, model = None, max_seq_length_explicit = False, ): if ( not getattr(args, "streaming", False) and not getattr(args, "preserve_dataset_order", False) and not getattr(args, "_unsloth_mlx_dataset_order_explicit", False) ): default_order = getattr(MLXTrainingConfig, "dataset_order", "default") if getattr(args, "dataset_order", default_order) in (None, default_order): args.dataset_order = "torch_randperm" if isinstance(args, UnslothTrainingArguments) and not getattr( args, "_unsloth_mlx_append_eos_explicit", False ): args.append_eos = False if isinstance(args, UnslothTrainingArguments) and not getattr( args, "_unsloth_mlx_grad_clip_explicit", False ): max_grad_norm = _positive_mlx_training_number( getattr(args, "max_grad_norm", None), ) max_grad_value = _positive_mlx_training_number( getattr(args, "max_grad_value", None), ) max_grad_leaf_norm = _positive_mlx_training_number( getattr(args, "max_grad_leaf_norm", None), ) if max_grad_norm is None and max_grad_value is None and max_grad_leaf_norm is None: args.max_grad_norm = 1.0 if not max_seq_length_explicit: _resolve_mlx_cuda_style_max_seq_length(args, model = model) return args def _coerce_mlx_training_args(args, overrides = None): overrides = overrides or {} if isinstance(args, MLXTrainingConfig) and not overrides: return args dataset_order_explicit = None append_eos_explicit = None max_seq_length_explicit = None max_length_value = None grad_clip_explicit = None if args is None: values = {} elif isinstance(args, dict): values = dict(args) elif isinstance(args, (str, os.PathLike)): values = {"output_dir": os.fspath(args)} else: dataset_order_explicit = getattr( args, "_unsloth_mlx_dataset_order_explicit", False, ) append_eos_explicit = getattr( args, "_unsloth_mlx_append_eos_explicit", None, ) max_seq_length_explicit = getattr( args, "_unsloth_mlx_max_seq_length_explicit", None, ) if max_seq_length_explicit is None: args_max_seq_length = _positive_mlx_context_length( getattr(args, "max_seq_length", None), ) default_max_seq_length = getattr(MLXTrainingConfig, "max_seq_length", 2048) max_seq_length_explicit = ( args_max_seq_length is not None and args_max_seq_length != default_max_seq_length ) max_length_value = getattr( args, "_unsloth_mlx_max_length_value", getattr(args, "max_length", None), ) grad_clip_explicit = getattr( args, "_unsloth_mlx_grad_clip_explicit", None, ) values = _mlx_training_argument_values(args) if hasattr(args, "max_length"): values["max_length"] = getattr(args, "max_length") values.update(overrides) coerced = UnslothTrainingArguments(**values) if ( dataset_order_explicit is not None and "dataset_order" not in overrides and "preserve_dataset_order" not in overrides ): coerced._unsloth_mlx_dataset_order_explicit = dataset_order_explicit if append_eos_explicit is not None and "append_eos" not in overrides: coerced._unsloth_mlx_append_eos_explicit = append_eos_explicit if ( max_seq_length_explicit is not None and "max_seq_length" not in overrides and "max_length" not in overrides ): coerced._unsloth_mlx_max_seq_length_explicit = max_seq_length_explicit if max_length_value is not None and "max_length" not in overrides: coerced._unsloth_mlx_max_length_value = max_length_value coerced.max_length = max_length_value if ( grad_clip_explicit is not None and "max_grad_norm" not in overrides and "max_grad_value" not in overrides and "max_grad_leaf_norm" not in overrides ): coerced._unsloth_mlx_grad_clip_explicit = grad_clip_explicit return coerced _MLX_TRAINER_POSITIONAL_KWARGS = ( "model", "tokenizer", "train_dataset", "eval_dataset", "dataset_text_field", "max_seq_length", "packing", "data_collator", "args", "formatting_func", "processor", "callbacks", ) _TRL_SFT_TRAINER_POSITIONAL_KWARGS = ( "model", "args", "data_collator", "train_dataset", "eval_dataset", "processing_class", "compute_loss_func", "compute_metrics", "callbacks", "optimizers", "optimizer_cls_and_kwargs", "preprocess_logits_for_metrics", "peft_config", "formatting_func", ) _MLX_TRAINER_KWARGS = frozenset(_MLX_TRAINER_POSITIONAL_KWARGS) def _filter_supported_mlx_trainer_kwargs(trainer_kwargs): """Drop inert/empty kwargs unsupported by this zoo MLXTrainer.""" unsupported = { key: value for key, value in trainer_kwargs.items() if not _mlx_trainer_supports_kwarg(key) } names = sorted( key for key, value in unsupported.items() if _is_meaningful_mlx_extra_value(value) ) if names: subject = ", ".join(names) verb = "requires" if len(names) == 1 else "require" raise NotImplementedError( "Unsloth MLX: " f"{subject} {verb} an unsloth-zoo build with " "matching MLXTrainer support. Upgrade unsloth-zoo together " "with unsloth." ) for key in unsupported: trainer_kwargs.pop(key, None) return trainer_kwargs def _is_mlx_native_text_collator(collator): """HF pad/copy collators are redundant on MLX; match by class name.""" for klass in type(collator).__mro__: name = klass.__name__ if name in ( "DataCollatorForSeq2Seq", "DataCollatorWithPadding", "DefaultDataCollator", ): return True if name == "DataCollatorForLanguageModeling": # Plain causal padding is fine; MLM masking changes semantics. return not bool(getattr(collator, "mlm", False)) return False _MLX_VISION_COLLATOR_FORWARDED_KWARGS = frozenset( ("completion_only_loss", "formatting_func", "max_seq_length") ) _MLX_VISION_COLLATOR_IMAGE_KWARGS = frozenset(("image_size", "resize")) _MLX_VISION_COLLATOR_POSITIONAL_KWARGS = ( "max_seq_length", "formatting_func", "resize", "ignore_index", "train_on_responses_only", "instruction_part", "response_part", "force_match", "num_proc", "completion_only_loss", "pad_to_multiple_of", "resize_dimension", "snap_to_patch_size", "last_response_only", ) _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS = { "ignore_index": -100, "train_on_responses_only": False, "instruction_part": None, "response_part": None, "force_match": True, "num_proc": None, "pad_to_multiple_of": None, "resize_dimension": 0, "snap_to_patch_size": False, "last_response_only": False, } def _is_default_mlx_vision_collator_value(key, value): """Return whether an unsupported collator value is the CUDA default.""" if key not in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS: return False default = _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS[key] if default is None: return value is None if isinstance(default, bool): return value is default return value == default and type(value) is type(default) def _has_mlx_training_arg_value(args, key): """Return whether training args already carry an explicit config value.""" if args is None or isinstance(args, (str, os.PathLike)): return False if isinstance(args, dict): return key in args return getattr(args, key, None) is not None def _raise_unsupported_mlx_vision_collator_kwargs(collator_kwargs): """Reject VLM collator kwargs that cannot be ignored safely on MLX.""" unsupported = sorted( key for key, value in collator_kwargs.items() if ( key not in _MLX_VISION_COLLATOR_FORWARDED_KWARGS and key not in _MLX_VISION_COLLATOR_IMAGE_KWARGS and ( ( key in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS and not _is_default_mlx_vision_collator_value(key, value) ) or ( key not in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS and _is_meaningful_mlx_extra_value(value) ) ) ) ) if unsupported: raise NotImplementedError( "Unsloth MLX: unsupported UnslothVisionDataCollator kwargs " f"cannot be ignored safely: {', '.join(unsupported)}." ) class UnslothTrainer(MLXTrainer): """Backend-aware public trainer that routes supported SFT notebooks to MLX.""" def __init__(self, *args, **kwargs): kwargs = _normalize_mlx_trainer_init_args(args, kwargs) processing_class = kwargs.pop("processing_class", None) processor_from_processing_class = False if processing_class is not None: if kwargs.get("processor", None) is None: kwargs["processor"] = processing_class processor_from_processing_class = True if kwargs.get("tokenizer", None) is None: kwargs["tokenizer"] = getattr( processing_class, "tokenizer", processing_class, ) kwargs.setdefault("tokenizer", None) data_collator = kwargs.pop("data_collator", None) if data_collator is not None: if isinstance(data_collator, UnslothVisionDataCollator): collator_processor = getattr(data_collator, "processor", None) if collator_processor is not None and ( kwargs.get("processor", None) is None or processor_from_processing_class ): kwargs["processor"] = collator_processor if kwargs.get("tokenizer", None) is None: kwargs["tokenizer"] = getattr( collator_processor, "tokenizer", collator_processor, ) collator_kwargs = getattr(data_collator, "kwargs", None) or {} collator_explicit_kwargs = getattr( data_collator, "_unsloth_mlx_explicit_kwargs", set(collator_kwargs), ) collator_image_size = collator_kwargs.get( "image_size", collator_kwargs.get("resize", None), ) if isinstance(collator_image_size, list): collator_image_size = tuple(collator_image_size) if ( isinstance(collator_image_size, str) and collator_image_size.lower() == "max" ): collator_image_size = "max" if "image_size" not in kwargs and ( isinstance(collator_image_size, int) or collator_image_size == "max" or ( isinstance(collator_image_size, tuple) and len(collator_image_size) == 2 and all(isinstance(x, int) for x in collator_image_size) ) ): kwargs["image_size"] = collator_image_size for collator_key in _MLX_VISION_COLLATOR_FORWARDED_KWARGS: collator_defaulted_value = collator_key not in collator_explicit_kwargs if collator_defaulted_value and _has_mlx_training_arg_value( kwargs.get("args"), collator_key ): continue if ( collator_key in collator_kwargs and collator_key not in kwargs and collator_kwargs[collator_key] is not None ): kwargs[collator_key] = collator_kwargs[collator_key] _raise_unsupported_mlx_vision_collator_kwargs(collator_kwargs) elif _is_mlx_native_text_collator(data_collator): pass # redundant on MLX; MLXTrainer batches/masks/pads natively else: raise NotImplementedError( "Unsloth MLX: custom data_collator is not supported by " "MLXTrainer. Pass the dataset directly or use the MLX " "trainer's native batching path." ) trainer_kwargs, config_kwargs, ignored_kwargs = _split_mlx_trainer_kwargs(kwargs) _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs) trainer_kwargs = _filter_supported_mlx_trainer_kwargs(trainer_kwargs) trainer_kwargs["args"] = _coerce_mlx_training_args( trainer_kwargs.get("args"), config_kwargs, ) if getattr( trainer_kwargs["args"], "completion_only_loss", None ) is True and not _is_vlm_model(trainer_kwargs.get("model")): raise NotImplementedError( "Unsloth MLX: completion_only_loss=True is only supported " "for VLM training. For text SFT, call train_on_responses_only " "after constructing the trainer." ) if getattr( trainer_kwargs["args"], "train_on_completions", None ) is True and not _is_vlm_model(trainer_kwargs.get("model")): raise NotImplementedError( "Unsloth MLX: train_on_completions=True is only supported " "for VLM training. For text SFT, call train_on_responses_only " "after constructing the trainer." ) trainer_kwargs["args"] = _apply_unsloth_trainer_mlx_defaults( trainer_kwargs["args"], model = trainer_kwargs.get("model"), max_seq_length_explicit = (trainer_kwargs.get("max_seq_length") is not None), ) super().__init__(**trainer_kwargs) self.processing_class = ( processing_class if processing_class is not None else self.processor or self.tokenizer ) if trainer_kwargs.get("max_seq_length") is not None: _set_mlx_cuda_style_context_length( self.args, self.args.max_seq_length, ) self._unsloth_mlx_ignored_trainer_kwargs = ignored_kwargs class UnslothVisionDataCollator: def __init__( self, model = None, processor = None, *args, **kwargs, ): explicit_kwargs = set(kwargs) if len(args) > len(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS): raise TypeError( "UnslothVisionDataCollator on MLX accepts at most " f"{len(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS)} positional " "options after model and processor." ) for key, value in zip(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS, args): if key in kwargs: raise TypeError( f"UnslothVisionDataCollator got multiple values for argument {key!r}" ) kwargs[key] = value explicit_kwargs.add(key) if "completion_only_loss" not in kwargs: kwargs["completion_only_loss"] = True self.model = model self.processor = processor self.args = () self.kwargs = kwargs self._unsloth_mlx_explicit_kwargs = explicit_kwargs def __call__(self, features): raise NotImplementedError( "Unsloth: UnslothVisionDataCollator is a compatibility placeholder " "on MLX. Pass the dataset to UnslothTrainer; MLXTrainer performs " "vision batching internally." ) def get_chat_template(*args, **kwargs): """Apply an Unsloth chat template through a lazy MLX-safe import.""" from .chat_templates import get_chat_template as _get_chat_template return _get_chat_template(*args, **kwargs) def apply_chat_template(*args, **kwargs): """Format a dataset with an Unsloth chat template through a lazy import.""" from .chat_templates import apply_chat_template as _apply_chat_template return _apply_chat_template(*args, **kwargs) def standardize_data_formats(*args, **kwargs): """Normalize ShareGPT-style datasets through the shared zoo helper.""" from unsloth_zoo.dataset_utils import standardize_data_formats as _standardize_data_formats return _standardize_data_formats(*args, **kwargs) def standardize_sharegpt(*args, **kwargs): """Alias ShareGPT standardization to the shared dataset-format helper.""" return standardize_data_formats(*args, **kwargs) def train_on_responses_only(*args, **kwargs): """Mask non-response tokens through the shared zoo dataset helper.""" from unsloth_zoo.dataset_utils import train_on_responses_only as _train_on_responses_only return _train_on_responses_only(*args, **kwargs) def _safe_mlx_trl_star_exports(_trl): """Return importable TRL star exports plus the MLX SFT shims.""" exports = list(getattr(_trl, "__all__", ())) safe_exports = [] for name in exports: try: getattr(_trl, name) except Exception: continue safe_exports.append(name) for name in ("SFTConfig", "SFTTrainer"): if name not in safe_exports: safe_exports.append(name) return safe_exports # trl trainers with no MLX implementation yet. Swap them for stubs that fail # with a clear message instead of importing the real torch/CUDA trainer and # crashing deep inside it, so an unmigrated GRPO/DPO/ORPO notebook is legible. _MLX_UNSUPPORTED_TRL_TRAINERS = ( "GRPOTrainer", "DPOTrainer", "ORPOTrainer", "KTOTrainer", "PPOTrainer", "RewardTrainer", ) def _make_mlx_unsupported_trl_trainer(name): def __init__(self, *args, **kwargs): raise NotImplementedError( f"Unsloth: {name} is not yet supported on the MLX (Apple Silicon) " f"backend. Only SFT training runs on MLX today; use a CUDA/ROCm GPU " f"for {name}." ) return type(name, (), {"__init__": __init__, "_unsloth_mlx_unsupported": True}) class _MLXSFTConfig(UnslothTrainingArguments): """`trl.SFTConfig` alias that keeps TRL's default training length. TRL/HF SFTConfig defaults to num_train_epochs=3 (max_steps=-1); the native MLX config defaults to max_steps=60. An unmigrated notebook that builds SFTConfig without an explicit length would otherwise silently run 60 MLX steps under this alias, so seed the TRL epoch default when neither max_steps nor num_train_epochs is given (epoch mode is MLX-supported). """ def __init__(self, *args, **kwargs): keys = set(kwargs) if len(args) == 1 and isinstance(args[0], dict): keys |= set(args[0]) if not ({"max_steps", "num_train_epochs"} & keys): kwargs.setdefault("num_train_epochs", 3) super().__init__(*args, **kwargs) def _install_mlx_trl_sft_shim(): """Install MLX-backed TRL SFT shims without replacing the TRL module.""" _trl = _sys.modules.get("trl") if _trl is None: try: import trl as _trl except ImportError: _trl = _types.ModuleType("trl") _trl.__version__ = "0.0.0+unsloth-mlx" _trl.__package__ = "trl" _trl.__path__ = [] _trl.__spec__ = _machinery.ModuleSpec("trl", loader = None, is_package = True) _sys.modules["trl"] = _trl _trl.SFTTrainer = UnslothTrainer _trl.SFTConfig = _MLXSFTConfig # Only retarget trainers the installed trl actually exposes (don't invent # attributes); idempotent so re-importing unsloth is a no-op. # Decide what to stub from trl's declared exports (__all__) and already # materialized attrs only. A getattr probe here would trigger trl's lazy # trainer import, pulling torch and breaking `import unsloth` on torch-free # MLX just to check existence. _trl_exports = set(getattr(_trl, "__all__", ()) or ()) # Stub every non-SFT trainer trl exposes, not just a fixed list, so newer # trainers (RLOOTrainer, ...) also fail with a clear MLX message instead # of importing the real torch trainer. Names come from __all__ so we never # resolve them (that would trigger trl's lazy import and pull torch). _unsupported = set(_MLX_UNSUPPORTED_TRL_TRAINERS) | { _n for _n in _trl_exports if _n.endswith("Trainer") and _n != "SFTTrainer" } for _name in _unsupported: _current = vars(_trl).get(_name) if getattr(_current, "_unsloth_mlx_unsupported", False): continue if _name in _trl_exports or _current is not None: setattr(_trl, _name, _make_mlx_unsupported_trl_trainer(_name)) _trl.__all__ = _safe_mlx_trl_star_exports(_trl) _trl.__UNSLOTH_MLX_COMPAT__ = True def _install_mlx_unsloth_trainer_shim(): module_name = f"{__name__}.trainer" _trainer = _types.ModuleType(module_name) _trainer.__package__ = __name__ _trainer.__spec__ = _machinery.ModuleSpec(module_name, loader = None) _trainer.MLXTrainer = MLXTrainer _trainer.MLXTrainingConfig = MLXTrainingConfig _trainer.UnslothTrainer = UnslothTrainer _trainer.UnslothTrainingArguments = UnslothTrainingArguments _trainer.UnslothVisionDataCollator = UnslothVisionDataCollator _sys.modules[module_name] = _trainer globals()["trainer"] = _trainer _install_mlx_trl_sft_shim() _install_mlx_unsloth_trainer_shim() else: # GPU path: load everything from _gpu_init from ._gpu_init import * from ._gpu_init import __version__ def get_gpu_memory_stats(): """Return CUDA/ROCm/XPU device stats, peak memory, and total memory in GiB.""" try: import torch if hasattr(torch, "xpu") and torch.xpu.is_available(): props = torch.xpu.get_device_properties(0) peak = ( torch.xpu.max_memory_reserved() if hasattr(torch.xpu, "max_memory_reserved") else torch.xpu.max_memory_allocated() ) total = getattr(props, "total_memory", 0) return props, _bytes_to_gb(peak), _bytes_to_gb(total) or 1.0 if hasattr(torch, "cuda") and torch.cuda.is_available(): props = torch.cuda.get_device_properties(0) peak = torch.cuda.max_memory_reserved() total = getattr(props, "total_memory", 0) return props, _bytes_to_gb(peak), _bytes_to_gb(total) or 1.0 except Exception: pass stats = _UnslothDeviceStats("Unknown GPU", 0) return stats, 0.0, 1.0 def clear_gpu_memory(): """Clear cached GPU memory on CUDA, ROCm, or XPU when available.""" try: import torch if hasattr(torch, "xpu") and torch.xpu.is_available(): torch.xpu.empty_cache() elif hasattr(torch, "cuda") and torch.cuda.is_available(): torch.cuda.empty_cache() except Exception: pass