From d91824583452f8d1faf3973a15d3dc4ef5a334ac Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Fri, 3 Jul 2026 06:02:26 +0800 Subject: [PATCH] Add MLX-aware public Unsloth trainer API (#6462) * feat: add mlx public trainer api * test: cover mlx public trainer api * fix: preserve mlx epoch trainer configs * fix: pass mlx warmup ratio through config * fix: align mlx trainer dataset order * fix: keep mlx chat templates import-light * fix: infer mlx trainer context length * fix: mirror cuda mlx context defaults * fix: align mlx notebook trainer defaults * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: keep mlx public helpers import-light * refactor: reuse mlx optimizer normalization * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: address mlx review feedback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: tighten mlx training argument parity * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: align mlx trainer eos default * Fix MLX trainer to accept DataCollatorForSeq2Seq and handle TokenizerWrapper in get_chat_template * Trim redundant docstrings on internal MLX helpers * MLX review fixes: Studio optimizer import-safe on non-MLX hosts, preserve explicit max_length, skip MLX tests before import * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX review round 2: defer max_length to model context, optimizer alias fallback for older zoo, skip non-MLX test on missing GPU deps * MLX review round 3: keep chat_templates importable without torch on MLX * fix: preserve MLX trainer notebook shims * fix: ignore CUDA tokenizer moves on MLX * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: harden MLX trainer shims * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: unwrap MLX scheduler enum args * fix: coerce integral MLX epoch counts * fix: spoof CUDA compatibility APIs on MLX * fix: harden MLX notebook compatibility shims * MLX: add torch.cuda.mem_get_info to the compatibility shim Notebook memory cells call torch.cuda.mem_get_info()[0] directly (not gated by is_available), so on MLX it raises without a shim. Return (free, total) bytes from the MLX device stats, consistent with the other torch.cuda compat helpers, and add a matching assertion to the compat-API test. * MLX: use active memory for mem_get_info; fix BatchEncoding.to keyword device Address review on the MLX compatibility shim: - torch.cuda.mem_get_info() now derives free bytes from current active MLX memory instead of the peak high-water mark, so a capacity check stays accurate after a transient spike or a prior run. - BatchEncoding.to(device=...) passed by keyword no longer forwards a positional None alongside the keyword (which raised "multiple values for 'device'"), so non-CUDA keyword moves like .to(device="cpu") delegate correctly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX: accept preserve_dataset_order; stub RL trainers with a clear error Two fixes so unmigrated notebooks behave predictably on MLX (torch present): - preserve_dataset_order is a real MLXTrainingConfig field but was missing from the extra-argument allowlist, so passing it (as a config or trainer kwarg) could be rejected as unknown on a zoo without the field. Add it to _MLX_IMPLEMENTED_EXTRA_ARGUMENTS so the documented no-shuffle path is reachable. - GRPO/DPO/ORPO (and KTO/PPO/Reward) have no MLX trainer yet. Retarget the ones the installed trl exposes to a stub that raises a clear 'not supported on MLX' error instead of importing the real torch/CUDA trainer and crashing deep inside it. Only existing trainers are retargeted (no invented attributes), idempotent across re-imports. * MLX: make RL-trainer stubbing import-safe; back current-memory APIs with active memory Address review on the MLX shims: - The RL-trainer stub loop probed trl with getattr(_trl, name), which triggers trl's lazy trainer import and pulls torch -- that can crash import unsloth on a torch-free MLX install just to check existence. Decide what to stub from trl.__all__ + already-materialized attrs (vars) instead; never resolve the real trainer. All trl trainer names are in __all__, so they are still stubbed (even torch-free), and the probe no longer imports torch. - torch.cuda.memory_reserved / memory_allocated (the current, non-max APIs) were aliased to peak max_memory_reserved. Back them with current active MLX memory so cleanup / capacity checks see live usage; max_* keep the peak high-water mark. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX: keep TRL's SFTConfig epoch default under the trl.SFTConfig alias Unmigrated notebooks import SFTConfig from trl, which the MLX build aliases to the public training-args class. TRL/HF SFTConfig defaults to num_train_epochs=3 (max_steps=-1); the native MLX config defaults to max_steps=60. So an SFTConfig built without an explicit length silently ran 60 MLX steps instead of TRL's 3 epochs under the alias. Alias trl.SFTConfig to a thin subclass that seeds the TRL epoch default only when neither max_steps nor num_train_epochs is given; explicit lengths pass through untouched, and the native public args class keeps its MLX default. Epoch mode is supported by the MLX trainer. * MLX CI: keep the GGUF reload smoke under the job timeout The RELOAD-GGUF-via-llama-cli step timed out at 300s. BF16 GGUF decode is CPU-bound on the macOS runner (~10s+/token), so generating 24 tokens landed right on the 300s cliff and killed the process. This step is a save/reload integrity smoke (it only needs a few chars of output), so the token count is incidental: generate 8 tokens with explicit threads and a small headroom on the subprocess timeout, all env-tunable (UNSLOTH_GGUF_RELOAD_N / _THREADS / _TIMEOUT). Cuts the reload well under the 25 minute job budget. * MLX: broaden trainer stubs, real peak-memory reset, fix shim tests Address review on the MLX public API: - The SFTConfig identity tests asserted trl.SFTConfig is UnslothTrainingArguments, but the alias now points at the _MLXSFTConfig subclass that preserves TRL's epoch default, so the MLX suite failed before testing the shim. Assert issubclass instead. - torch.cuda.reset_peak_memory_stats was a no-op, so max_memory_reserved kept earlier model-load peaks across a scoped run. Wire it to mx.reset_peak_memory with the same core/metal fallback used for the reads. - The unsupported-trainer stubs were a fixed list, so trainers outside it (a newer RLOOTrainer) still routed to the real torch trainer. Derive the set from trl.__all__ (every non-SFT *Trainer) so all non-SFT surfaces fail with a clear MLX message; names come from __all__ so trl is never resolved. - The non-MLX export smoke skipped only on missing bitsandbytes/triton; other absent GPU deps (numpy/torch/unsloth-zoo, or _gpu_init re-raising ImportError) made it fail on CPU hosts. Skip on any ImportError. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: keep MLX notebook compatibility minimal * MLX CI: force CPU + small context for the GGUF reload smoke The RELOAD-GGUF-via-llama-cli step timed out even at 8 tokens (>420s), so it is a fixed hang, not per-token cost: on the paravirtual macOS runner GPU llama.cpp's Metal backend stalls, and the gemma3 GGUF advertises a 32768 context that llama-cli would otherwise fully allocate. Run llama-cli CPU-only (-ngl 0) with a small context (-c 256); keep generation short. All env-tunable (UNSLOTH_GGUF_RELOAD_NGL / _CTX / _N / _THREADS / _TIMEOUT). Also print llama.cpp's partial stdout/stderr on timeout so a future hang is diagnosable instead of an opaque TimeoutExpired. * MLX CI: export the reload-smoke GGUF as q8_0, not bf16 The GGUF reload via llama-cli timed out on the runner even CPU-only with a tiny context and 8 tokens. Root cause is the format, not the flags: the smoke exported quantization_method='not_quantized', which maps to a bf16 GGUF, and llama.cpp's bf16 CPU decode is unusably slow on the paravirtual macOS runner. Export q8_0 (fast_quantized, the exporter default and what users deploy) instead -- llama.cpp has optimized q8_0 CPU kernels, so the fresh-process reload loads and generates in seconds. The reload stays CPU-only (-ngl 0) with a small context. * test: clear TRL shim before availability check --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 --- studio/backend/core/training/worker.py | 58 +- .../tests/test_mlx_training_worker_config.py | 2 +- tests/python/test_mlx_public_trainer_api.py | 1206 ++++++++++++++++ tests/studio/run_real_mlx_smoke.py | 54 +- unsloth/__init__.py | 1270 ++++++++++++++++- unsloth/chat_templates.py | 67 +- 6 files changed, 2595 insertions(+), 62 deletions(-) create mode 100644 tests/python/test_mlx_public_trainer_api.py diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 610af2472e..17dc1299ca 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1253,32 +1253,48 @@ def _adapt_for_mlx_vlm( return adapted -_MLX_STUDIO_OPTIM_MAP = { - "adamw_8bit": "adamw", - "paged_adamw_8bit": "adamw", - "adamw_bnb_8bit": "adamw", - "paged_adamw_32bit": "adamw", - "adamw_torch": "adamw", - "adamw_torch_fused": "adamw", - "adamw": "adamw", - "adafactor": "adafactor", - "sgd": "sgd", - "adam": "adam", - "muon": "muon", - "lion": "lion", -} _MLX_STUDIO_LR_SCHEDULERS = {"linear", "cosine", "constant"} +# Fallback alias map mirroring unsloth_zoo._normalize_mlx_optimizer_name, used +# only when mlx (Apple Silicon) is not importable so Studio config validation +# still works on non-MLX hosts. The zoo function stays the source of truth. +_MLX_STUDIO_ADAMW_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", + ) +) +_MLX_STUDIO_NATIVE_OPTIMIZERS = ("adafactor", "adamw", "adam", "sgd", "muon", "lion") + + def _normalize_mlx_studio_optimizer(value): - raw = str(value or "adamw_8bit").strip().lower() try: - return _MLX_STUDIO_OPTIM_MAP[raw] - except KeyError: - supported = ", ".join(sorted(_MLX_STUDIO_OPTIM_MAP)) - raise ValueError( - f"Unsupported optimizer for MLX training: {value!r}. " f"Supported values: {supported}." - ) + from unsloth_zoo.mlx.trainer import _normalize_mlx_optimizer_name + return _normalize_mlx_optimizer_name(value or "adamw_8bit") + except (ImportError, ValueError): + # Missing mlx, or an older unsloth-zoo whose normalizer lacks CUDA/TRL + # aliases: map common adamw_* names locally so notebook defaults work. + opt = str(getattr(value, "value", value) or "adamw_8bit").strip().lower() + opt = opt.rsplit(".", 1)[-1].replace("-", "_") + if opt in _MLX_STUDIO_ADAMW_ALIASES: + opt = "adamw" + if opt not in _MLX_STUDIO_NATIVE_OPTIMIZERS: + supported = ", ".join(_MLX_STUDIO_NATIVE_OPTIMIZERS) + raise ValueError( + f"Unsupported optimizer for MLX training: {value!r}. " + f"Supported optimizers: {supported}." + ) + return opt def _normalize_mlx_studio_scheduler(value): diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index dce5e27c08..14fc0933d0 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -76,7 +76,7 @@ def test_mlx_studio_optimizer_aliases_are_explicit(): def test_mlx_studio_rejects_unknown_optimizer(): - with pytest.raises(ValueError, match = "Unsupported optimizer for MLX training"): + with pytest.raises(ValueError, match = "Supported"): _normalize_mlx_studio_optimizer("adamw_typo") diff --git a/tests/python/test_mlx_public_trainer_api.py b/tests/python/test_mlx_public_trainer_api.py new file mode 100644 index 0000000000..2c33f86af1 --- /dev/null +++ b/tests/python/test_mlx_public_trainer_api.py @@ -0,0 +1,1206 @@ +"""Tests for the MLX public trainer compatibility surface.""" + +from __future__ import annotations + +import builtins +import importlib +import importlib.util +import platform +import sys +import types +import warnings + +import pytest + +_MLX_SKIP_REASON = "MLX public trainer API is only active on the MLX backend" + + +def _import_mlx_unsloth(): + """Import unsloth and skip when the current platform is not using MLX.""" + # Skip before importing unsloth so non-MLX hosts missing optional GPU deps + # (e.g. bitsandbytes) skip cleanly instead of erroring at collection. + if not ( + platform.system() == "Darwin" + and platform.machine() == "arm64" + and importlib.util.find_spec("mlx") is not None + ): + pytest.skip(_MLX_SKIP_REASON) + unsloth = importlib.import_module("unsloth") + if getattr(unsloth, "DEVICE_TYPE", None) != "mlx": + pytest.skip(_MLX_SKIP_REASON) + return unsloth + + +class _DummyModel: + """Small model stub that satisfies MLXTrainer constructor probes.""" + + def trainable_parameters(self): + """Return no trainable parameters for constructor-only tests.""" + return {} + + +class _DummyVLMModel(_DummyModel): + """Small VLM model stub for MLX vision trainer constructor probes.""" + + _is_vlm_model = True + + +def test_mlx_exports_unsloth_trainer_api(): + """MLX imports should expose the public Unsloth trainer API.""" + unsloth = _import_mlx_unsloth() + from unsloth import ( + RawTextDataLoader, + TextPreprocessor, + UnslothTrainer, + UnslothTrainingArguments, + clear_gpu_memory, + get_gpu_memory_stats, + ) + + assert RawTextDataLoader is unsloth.RawTextDataLoader + assert TextPreprocessor is unsloth.TextPreprocessor + assert UnslothTrainer is unsloth.UnslothTrainer + assert UnslothTrainingArguments is unsloth.UnslothTrainingArguments + assert get_gpu_memory_stats is unsloth.get_gpu_memory_stats + assert clear_gpu_memory is unsloth.clear_gpu_memory + assert issubclass(UnslothTrainer, unsloth.MLXTrainer) + assert issubclass(UnslothTrainingArguments, unsloth.MLXTrainingConfig) + assert importlib.util.find_spec("unsloth.memory") is None + + +def test_non_mlx_exports_public_trainer_api_when_available(): + """GPU/ROCm imports should keep exporting the public Unsloth trainer API.""" + try: + unsloth = importlib.import_module("unsloth") + except ImportError as exc: + # Non-MLX import pulls the optional GPU stack (numpy/torch/unsloth-zoo, + # bitsandbytes/triton, and _gpu_init can re-raise missing deps as + # ImportError). Skip when any of it is unavailable rather than failing + # collection on CPU/ROCm/XPU review hosts. + pytest.skip(f"non-MLX import dependency unavailable: {exc}") + if getattr(unsloth, "DEVICE_TYPE", None) == "mlx": + pytest.skip("non-MLX export smoke test only runs on GPU/ROCm backends") + + assert unsloth.UnslothTrainer is not None + assert unsloth.UnslothTrainingArguments is not None + assert callable(unsloth.get_gpu_memory_stats) + assert callable(unsloth.clear_gpu_memory) + assert importlib.util.find_spec("unsloth.memory") is None + + +def test_mlx_training_arguments_accept_trl_style_kwargs(): + """TRL/SFTConfig-style kwargs should normalize without breaking MLX config.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "bf16.*dataset_kwargs"): + args = unsloth.UnslothTrainingArguments( + max_length = 123, + max_steps = 10, + warmup_ratio = 0.2, + remove_unused_columns = False, + dataset_kwargs = {"skip_prepare_dataset": True}, + bf16 = True, + ) + + assert args.max_seq_length == 123 + assert args.warmup_steps == 2 + assert args.remove_unused_columns is False + assert args.dataset_kwargs == {"skip_prepare_dataset": True} + assert args.bf16 is True + assert args.warmup_ratio == 0.2 + assert args._unsloth_mlx_max_seq_length_explicit is False + assert args._unsloth_mlx_warmup_steps_explicit is False + + +def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras(): + """Implemented and falsey inert compatibility kwargs should stay quiet.""" + unsloth = _import_mlx_unsloth() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + args = unsloth.UnslothTrainingArguments( + warmup_ratio = 0.2, + max_steps = 10, + padding_free = False, + remove_unused_columns = False, + assistant_only_loss = False, + completion_only_loss = False, + ) + + assert args.warmup_steps == 2 + assert args.padding_free is False + assert args.remove_unused_columns is False + assert args.completion_only_loss is False + assert caught == [] + + +def test_mlx_training_arguments_prefer_canonical_max_seq_length(): + """Canonical MLX config fields should win over compatibility aliases.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(max_seq_length = 456, max_length = 123) + dict_args = unsloth.UnslothTrainingArguments( + {"max_length": 123, "max_seq_length": 456}, + ) + + assert args.max_seq_length == 456 + assert args.max_length == 456 + assert args._unsloth_mlx_max_length_value == 456 + assert dict_args.max_seq_length == 456 + assert dict_args.max_length == 456 + assert dict_args._unsloth_mlx_max_length_value == 456 + assert args._unsloth_mlx_max_seq_length_explicit is True + assert dict_args._unsloth_mlx_max_seq_length_explicit is True + + +def test_mlx_training_arguments_preserve_explicit_positive_warmup_steps(): + """Explicit warmup_steps should take precedence over warmup_ratio.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments( + max_steps = 10, + warmup_steps = 5, + warmup_ratio = 0.1, + ) + + assert args.warmup_steps == 5 + assert args._unsloth_mlx_warmup_steps_explicit is True + + +def test_mlx_clear_gpu_memory_uses_metal_fallback(monkeypatch): + """Older MLX releases expose cache clearing under mx.metal.clear_cache.""" + unsloth = _import_mlx_unsloth() + import mlx.core as mx + + called = [] + metal = getattr(mx, "metal", None) or type("Metal", (), {})() + monkeypatch.delattr(mx, "clear_cache", raising = False) + monkeypatch.setattr(mx, "metal", metal, raising = False) + monkeypatch.setattr(metal, "clear_cache", lambda: called.append("metal"), raising = False) + + unsloth.clear_gpu_memory() + + assert called == ["metal"] + + +def test_mlx_training_arguments_preserve_explicit_epoch_training(): + """Epoch-based configs should not inherit the MLX max_steps default.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(num_train_epochs = 1, warmup_ratio = 0.1) + default_args = unsloth.UnslothTrainingArguments() + + assert args.num_train_epochs == 1 + assert args.max_steps == -1 + assert args.warmup_ratio == 0.1 + assert args._unsloth_mlx_warmup_steps_explicit is False + assert default_args.max_steps == unsloth.MLXTrainingConfig.max_steps + + +def test_mlx_training_arguments_keep_mlx_dataset_order_default(): + """Training arguments alone should not override MLX's native data order.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(max_steps = 1) + explicit_default = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "default", + ) + + assert args.dataset_order == "default" + assert args._unsloth_mlx_dataset_order_explicit is False + assert args._unsloth_mlx_max_seq_length_explicit is False + assert explicit_default.dataset_order == "default" + assert explicit_default._unsloth_mlx_dataset_order_explicit is True + + +def test_mlx_training_arguments_warn_on_meaningful_inert_kwargs(): + """Unsupported TrainingArguments knobs should not be silently ignored.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "push_to_hub.*save_strategy"): + args = unsloth.UnslothTrainingArguments( + save_strategy = "steps", + push_to_hub = True, + padding_free = False, + ) + + assert args.save_strategy == "steps" + assert args.push_to_hub is True + assert args.padding_free is False + + +def test_mlx_training_arguments_reject_unknown_kwargs(): + """Unknown SFTConfig flags should fail instead of becoming inert attributes.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "assistant_only_loss"): + unsloth.UnslothTrainingArguments(assistant_only_loss = True) + + completion_args = unsloth.UnslothTrainingArguments(completion_only_loss = True) + assert completion_args.completion_only_loss is True + + +def test_mlx_training_arguments_reject_unsupported_object_flags(): + """Object-style SFTConfig flags should not be silently dropped.""" + unsloth = _import_mlx_unsloth() + + class ArgsObject: + max_steps = 1 + assistant_only_loss = True + + with pytest.raises(NotImplementedError, match = "assistant_only_loss"): + unsloth._coerce_mlx_training_args(ArgsObject()) + + class CompletionArgsObject: + max_steps = 1 + completion_only_loss = True + + completion_args = unsloth._coerce_mlx_training_args(CompletionArgsObject()) + assert completion_args.completion_only_loss is True + + +def test_mlx_training_arguments_accept_output_dir_positional(): + """A single positional output_dir should match TrainingArguments behavior.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments("custom-outputs", max_steps = 3) + + assert args.output_dir == "custom-outputs" + assert args.max_steps == 3 + + +def test_mlx_training_arguments_normalize_optim_and_object_aliases(): + """Common notebook optimizer names and object aliases should normalize.""" + unsloth = _import_mlx_unsloth() + + class Scheduler: + value = "cosine" + + class ArgsObject: + optim = "adamw_8bit" + eval_steps = None + lr_scheduler_type = Scheduler() + max_length = 321 + max_steps = 10 + num_train_epochs = 3.0 + save_steps = 500 + save_strategy = "no" + warmup_ratio = 0.1 + warmup_steps = 0 + + with pytest.warns(RuntimeWarning, match = "save_strategy"): + args = unsloth._coerce_mlx_training_args(ArgsObject()) + + assert args.optim == "adamw" + assert args.eval_steps == 0 + assert args.lr_scheduler_type == "cosine" + assert args.max_seq_length == 321 + assert args.num_train_epochs == 3 + assert type(args.num_train_epochs) is int + assert args.save_steps == 0 + assert args.warmup_steps == 1 + assert args._unsloth_mlx_warmup_steps_explicit is False + + +def test_mlx_training_arguments_accept_supported_notebook_kwargs(): + """Supported SFT notebooks should be able to pass their current args.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns( + RuntimeWarning, + match = "bf16.*dataset_kwargs.*gradient_checkpointing_kwargs.*save_strategy", + ): + args = unsloth.UnslothTrainingArguments( + bf16 = True, + dataset_kwargs = {"skip_prepare_dataset": True}, + dataset_num_proc = 4, + dataset_text_field = "text", + embedding_learning_rate = 5e-5, + fp16 = False, + gradient_accumulation_steps = 8, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = {"use_reentrant": False}, + learning_rate = 1e-4, + logging_steps = 2, + lr_scheduler_type = "cosine", + max_grad_norm = 0.3, + max_length = 1024, + max_steps = 10, + num_train_epochs = 1, + optim = "paged_adamw_8bit", + output_dir = "outputs", + padding_free = False, + per_device_train_batch_size = 1, + remove_unused_columns = False, + report_to = "none", + save_strategy = "steps", + seed = 123, + warmup_ratio = 0.1, + weight_decay = 0.01, + ) + + assert args.dataset_num_proc == 4 + assert args.dataset_text_field == "text" + assert args.embedding_learning_rate == 5e-5 + assert args.gradient_accumulation_steps == 8 + assert args.gradient_checkpointing is True + assert args.learning_rate == 1e-4 + assert args.logging_steps == 2 + assert args.lr_scheduler_type == "cosine" + assert args.max_grad_norm == 0.3 + assert args.max_seq_length == 1024 + assert args.max_steps == 10 + assert args.num_train_epochs == 1 + assert args.optim == "adamw" + assert args.output_dir == "outputs" + assert args.per_device_train_batch_size == 1 + assert args.report_to == "none" + assert args.seed == 123 + assert args.warmup_ratio == 0.1 + assert args.warmup_steps == 1 + assert args.weight_decay == 0.01 + assert args.dataset_kwargs == {"skip_prepare_dataset": True} + assert args.gradient_checkpointing_kwargs == {"use_reentrant": False} + assert args.save_strategy == "steps" + + +def test_mlx_training_arguments_honor_direct_no_save_strategy(): + """Direct kwargs should map save_strategy=no to save_steps=0.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "save_strategy"): + args = unsloth.UnslothTrainingArguments( + save_strategy = "no", + save_steps = 500, + ) + + assert args.save_steps == 0 + + +def test_mlx_trainer_accepts_common_sft_kwargs(): + """UnslothTrainer should accept common SFTTrainer kwargs on MLX.""" + unsloth = _import_mlx_unsloth() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + dataset_num_proc = 8, + max_length = 456, + optim = "adamw_bnb_8bit", + processing_class = object(), + ) + + assert trainer.args.max_steps == 1 + assert trainer.args.dataset_num_proc == 8 + assert trainer.args.max_seq_length == 456 + assert trainer.args.max_grad_norm == 1.0 + assert trainer.args.optim == "adamw" + assert trainer.args.dataset_order == "torch_randperm" + assert trainer._unsloth_mlx_ignored_trainer_kwargs == {} + assert caught == [] + + +def test_mlx_trainer_preserves_explicit_dataset_order(): + """UnslothTrainer should only set torch_randperm when order is implicit.""" + unsloth = _import_mlx_unsloth() + + explicit_default = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "default", + ), + ) + explicit_sequential = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "sequential", + ), + ) + implicit_with_override = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + dataset_num_proc = 4, + ) + implicit_streaming = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, streaming = True), + ) + explicit_no_clip = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + max_grad_norm = 0.0, + ), + ) + + assert explicit_default.args.dataset_order == "default" + assert explicit_sequential.args.dataset_order == "sequential" + assert implicit_with_override.args.dataset_order == "torch_randperm" + assert implicit_streaming.args.dataset_order == "default" + assert implicit_with_override.args.max_grad_norm == 1.0 + assert explicit_no_clip.args.max_grad_norm == 0.0 + + +def test_mlx_trainer_uses_model_context_length_when_implicit(): + """UnslothTrainer should mirror CUDA's max_length bridge precedence.""" + unsloth = _import_mlx_unsloth() + model = _DummyModel() + model.max_seq_length = 321 + max_length_model = _DummyModel() + max_length_model.max_seq_length = 321 + none_model = _DummyModel() + none_model.max_seq_length = 321 + explicit_seq_model = _DummyModel() + explicit_seq_model.max_seq_length = 321 + clamped_seq_model = _DummyModel() + clamped_seq_model.max_seq_length = 321 + model_max_length = _DummyModel() + model_max_length.max_length = 777 + metadata_model = _DummyModel() + metadata_model.config = type("Config", (), {"max_position_embeddings": 888})() + metadata_tokenizer = type("Tokenizer", (), {"model_max_length": 999})() + explicit_max_length_no_model = _DummyModel() + trainer_override_model = _DummyModel() + trainer_override_model.max_seq_length = 321 + config_override_model = _DummyModel() + config_override_model.max_seq_length = 432 + + implicit = unsloth.UnslothTrainer( + model = model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + max_length_args = unsloth.UnslothTrainer( + model = max_length_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_length = 123), + ) + none_args = unsloth.UnslothTrainer( + model = none_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = None), + ) + explicit_seq = unsloth.UnslothTrainer( + model = explicit_seq_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = 123), + ) + clamped_seq = unsloth.UnslothTrainer( + model = clamped_seq_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = 654), + ) + model_max_length_only = unsloth.UnslothTrainer( + model = model_max_length, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + metadata_ignored = unsloth.UnslothTrainer( + model = metadata_model, + tokenizer = metadata_tokenizer, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + explicit_max_length = unsloth.UnslothTrainer( + model = explicit_max_length_no_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_length = 123), + ) + trainer_override = unsloth.UnslothTrainer( + model = trainer_override_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + max_seq_length = 654, + ) + config_with_override = unsloth.UnslothTrainer( + model = config_override_model, + tokenizer = None, + train_dataset = [], + args = unsloth.MLXTrainingConfig(max_steps = 1), + dataset_num_proc = 4, + ) + + assert implicit.args.max_seq_length == 321 + assert implicit.args.max_length == 321 + assert max_length_args.args.max_seq_length == 321 + assert max_length_args.args.max_length == 321 + assert none_args.args.max_seq_length == 321 + assert none_args.args.max_length == 321 + assert explicit_seq.args.max_seq_length == 123 + assert explicit_seq.args.max_length == 123 + assert clamped_seq.args.max_seq_length == 321 + assert clamped_seq.args.max_length == 321 + assert model_max_length_only.args.max_seq_length == 777 + assert model_max_length_only.args.max_length == 777 + assert metadata_ignored.args.max_seq_length == 1024 + assert metadata_ignored.args.max_length == 1024 + assert explicit_max_length.args.max_seq_length == 123 + assert explicit_max_length.args.max_length == 123 + assert trainer_override.args.max_seq_length == 654 + assert trainer_override.args.max_length == 654 + assert config_with_override.args.max_seq_length == 432 + assert config_with_override.args.max_length == 432 + + +def test_mlx_trainer_processing_class_overrides_explicit_none_tokenizer(): + """TRL passes tokenizer=None while processing_class carries the tokenizer.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + + class Processor: + pass + + processor = Processor() + processor.tokenizer = tokenizer + + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processing_class = processor, + ) + + assert trainer.processor is processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_vision_collator_processor_overrides_processing_class(): + """Vision notebooks pass the tokenizer as processing_class and processor in collator.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + + class Processor: + pass + + processor = Processor() + processor.tokenizer = tokenizer + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), processor) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processing_class = tokenizer, + data_collator = collator, + ) + + assert trainer.processor is processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_preserves_explicit_processor_over_vision_collator(): + """Explicit processor kwargs should stay authoritative over collator metadata.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + explicit_processor = object() + + class Processor: + pass + + collator_processor = Processor() + collator_processor.tokenizer = tokenizer + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), collator_processor) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processor = explicit_processor, + processing_class = tokenizer, + data_collator = collator, + ) + + assert trainer.processor is explicit_processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_forwards_vision_collator_positional_defaults(): + """Vision collator CUDA-style positionals should route into MLX args.""" + unsloth = _import_mlx_unsloth() + collator = unsloth.UnslothVisionDataCollator( + _DummyVLMModel(), + object(), + 2048, + None, + "max", + -100, + False, + None, + None, + True, + None, + False, + ) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + data_collator = collator, + ) + + assert trainer.args.max_seq_length == 2048 + assert trainer.args.image_size == "max" + assert trainer.args.completion_only_loss is False + + +def test_mlx_vision_collator_default_does_not_override_explicit_args(): + """Implicit collator defaults should not override explicit trainer args.""" + unsloth = _import_mlx_unsloth() + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), object()) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = False, + ), + data_collator = collator, + ) + + assert trainer.args.completion_only_loss is False + + +def test_mlx_trainer_rejects_unsafe_unsupported_sft_kwargs(): + """Unsupported kwargs that change training semantics should fail on MLX.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "peft_config"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + peft_config = object(), + ) + + +def test_mlx_trainer_rejects_metrics_and_callbacks(): + """Trainer hooks should fail because MLXTrainer cannot honor them yet.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "callbacks"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [object()], + ) + with pytest.raises(NotImplementedError, match = "compute_metrics"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + compute_metrics = lambda *_: None, + ) + + +def test_mlx_trainer_rejects_custom_data_collator(): + """MLXTrainer owns batching; custom SFT data collators must not be ignored.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "data_collator"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + data_collator = object(), + ) + + +def test_mlx_trainer_rejects_text_completion_only_loss(): + """Text MLX training should not silently ignore completion_only_loss=True.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "completion_only_loss=True"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = True, + ), + ) + + +def test_mlx_trainer_allows_vlm_completion_only_loss(): + """VLM MLX training supports completion_only_loss during collation.""" + unsloth = _import_mlx_unsloth() + + class VLMModel(_DummyModel): + _is_vlm_model = True + + trainer = unsloth.UnslothTrainer( + model = VLMModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = True, + ), + ) + + assert trainer.args.completion_only_loss is True + + +def test_mlx_trainer_accepts_trl_style_positional_args(): + """TRL-style positional `(model, args, ...)` should not be read as tokenizer.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments("trl-outputs", max_steps = 2) + trainer = unsloth.UnslothTrainer( + _DummyModel(), + args, + train_dataset = [], + tokenizer = None, + ) + + assert trainer.args is args + assert trainer.args.output_dir == "trl-outputs" + assert trainer.train_dataset == [] + + +def test_mlx_trainer_accepts_trl_none_placeholder_positionals(): + """Explicit TRL default placeholders should preserve later positional args.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + processing_class = object() + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + dataset, + None, + processing_class, + ) + + assert getattr(trainer.train_dataset, "_dataset", trainer.train_dataset) is dataset + assert getattr(trainer, "_mlx_train_dataset_for_batches", dataset) is dataset + assert trainer.tokenizer is processing_class + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_accepts_short_trl_none_placeholder_positionals(): + """Short TRL placeholder calls should keep the fourth arg as train_dataset.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + dataset, + ) + + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_accepts_short_trl_placeholders_with_keyword_dataset(): + """Short TRL placeholders should not conflict with keyword train_dataset.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + train_dataset = dataset, + ) + + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_preserves_mlx_positional_schema_with_none_tokenizer(): + """MLX-style `(model, tokenizer, train_dataset, ...)` should still work.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + dataset, + None, + ) + + assert trainer.tokenizer is None + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + + +def test_mlx_compatibility_shims_are_installed(): + """Old notebook imports should resolve to the MLX public API after unsloth import.""" + unsloth = _import_mlx_unsloth() + + trl = importlib.import_module("trl") + trainer_module = importlib.import_module("unsloth.trainer") + chat_templates = importlib.import_module("unsloth.chat_templates") + dataset_utils = importlib.import_module("unsloth_zoo.dataset_utils") + + assert importlib.util.find_spec("trl") is not None + assert importlib.util.find_spec("unsloth.trainer") is not None + assert unsloth.trainer is trainer_module + assert unsloth.chat_templates is chat_templates + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trainer_module.UnslothTrainer is unsloth.UnslothTrainer + assert trainer_module.UnslothVisionDataCollator is unsloth.UnslothVisionDataCollator + assert chat_templates.train_on_responses_only is dataset_utils.train_on_responses_only + assert callable(unsloth.train_on_responses_only) + + +def test_mlx_trl_shim_preserves_existing_trl_module(monkeypatch): + """The MLX TRL shim should patch, not replace, an already-loaded TRL module.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.existing_marker = object() + trl.ExistingExport = object() + trl.__all__ = ["ExistingExport", "BrokenExport"] + + def _raise_for_broken_export(name): + if name == "BrokenExport": + raise RuntimeError("optional dependency missing") + raise AttributeError(name) + + trl.__getattr__ = _raise_for_broken_export + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + assert sys.modules["trl"] is trl + assert trl.__path__ == ["real-trainer-package"] + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trl.__UNSLOTH_MLX_COMPAT__ is True + assert "ExistingExport" in trl.__all__ + assert "BrokenExport" not in trl.__all__ + assert "SFTTrainer" in trl.__all__ + assert "SFTConfig" in trl.__all__ + + +def test_mlx_trl_shim_installs_real_trl_or_stub(monkeypatch): + """The MLX TRL shim should prefer real TRL and stub only if unavailable.""" + unsloth = _import_mlx_unsloth() + monkeypatch.delitem(sys.modules, "trl", raising = False) + real_trl_available = importlib.util.find_spec("trl") is not None + + unsloth._install_mlx_trl_sft_shim() + trl = importlib.import_module("trl") + + if real_trl_available: + assert trl.__version__ != "0.0.0+unsloth-mlx" + else: + assert trl.__version__ == "0.0.0+unsloth-mlx" + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trl.__UNSLOTH_MLX_COMPAT__ is True + + +def test_mlx_trl_star_import_exports_public_shims(): + """Existing `from trl import *` callers should receive MLX SFT shims.""" + unsloth = _import_mlx_unsloth() + namespace = {} + + exec("from trl import *", namespace) + + assert namespace["SFTTrainer"] is unsloth.UnslothTrainer + assert issubclass(namespace["SFTConfig"], unsloth.UnslothTrainingArguments) + + +def test_mlx_rl_trainers_stub_with_clear_error(monkeypatch): + """GRPO/DPO/ORPO trainers have no MLX path, so the shim retargets the ones trl + exposes to a clear NotImplementedError instead of a confusing CUDA crash, and + never invents trainers trl does not have.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + + class _RealTrainer: + def __init__(self, *args, **kwargs): + raise AssertionError("the real torch/CUDA trainer must not run on MLX") + + trl.GRPOTrainer = _RealTrainer + trl.DPOTrainer = _RealTrainer + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + for name in ("GRPOTrainer", "DPOTrainer"): + assert getattr(trl, name) is not _RealTrainer + with pytest.raises(NotImplementedError) as exc: + getattr(trl, name)(model = None, args = None) + assert "MLX" in str(exc.value) and name in str(exc.value) + # trainers trl never exposed must not be invented + assert not hasattr(trl, "PPOTrainer") + # idempotent: a second install keeps the same stub + stub = trl.GRPOTrainer + unsloth._install_mlx_trl_sft_shim() + assert trl.GRPOTrainer is stub + + +def test_mlx_rl_trainer_stub_is_lazy_import_safe(monkeypatch): + """Stubbing unsupported trl trainers must not resolve them: trl lazy-imports + pull torch, so on a torch-free MLX install a getattr probe would crash + `import unsloth`. The shim reads __all__/vars metadata and never triggers + trl's __getattr__ for a trainer it is about to replace.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.__all__ = ["SFTTrainer", "SFTConfig", "GRPOTrainer", "DPOTrainer"] + resolved = [] + + def _lazy_getattr(name): + resolved.append(name) + raise ImportError(f"lazy import of {name} would pull torch") + + trl.__getattr__ = _lazy_getattr + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() # must not raise despite the lazy trl + + # trainers declared in __all__ are stubbed WITHOUT ever resolving the real one + assert resolved == [] + for name in ("GRPOTrainer", "DPOTrainer"): + with pytest.raises(NotImplementedError): + getattr(trl, name)(model = None) + + +def test_mlx_stubs_trl_trainers_outside_fixed_set(monkeypatch): + """Any non-SFT trainer trl exports (e.g. a newer RLOOTrainer not in the fixed + list) must be stubbed too, so no torch trainer slips through on MLX.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.__all__ = ["SFTTrainer", "SFTConfig", "RLOOTrainer"] + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + with pytest.raises(NotImplementedError) as exc: + trl.RLOOTrainer(model = None) + assert "MLX" in str(exc.value) and "RLOOTrainer" in str(exc.value) + # SFT stays usable; only non-SFT trainers are stubbed + assert trl.SFTTrainer is unsloth.UnslothTrainer + + +def test_mlx_preserve_dataset_order_is_accepted(): + """preserve_dataset_order=True must be accepted (it is a real MLX config field), + not rejected as an unknown/unsupported argument.""" + unsloth = _import_mlx_unsloth() + args = unsloth.UnslothTrainingArguments( + output_dir = "mlx-out", + max_steps = 10, + preserve_dataset_order = True, + ) + assert getattr(args, "preserve_dataset_order", False) is True + + +def test_mlx_sftconfig_alias_keeps_trl_epoch_default(monkeypatch): + """`trl.SFTConfig` (aliased on MLX) keeps TRL's default training length: with + no explicit max_steps/num_train_epochs it runs TRL's 3 epochs, not the native + MLX 60-step default. An explicit length is authoritative and untouched.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + # no explicit length -> TRL epoch default (3 epochs, step cap disabled) + cfg = trl.SFTConfig(output_dir = "mlx-out") + assert cfg.num_train_epochs == 3 + assert cfg.max_steps == -1 + # explicit step / epoch counts stay exactly as written + assert trl.SFTConfig(output_dir = "mlx-out", max_steps = 17).max_steps == 17 + assert trl.SFTConfig(output_dir = "mlx-out", num_train_epochs = 2).num_train_epochs == 2 + + +def test_mlx_vision_collator_is_constructor_compatible(): + """Vision notebooks should be able to instantiate the collator placeholder.""" + unsloth = _import_mlx_unsloth() + + collator = unsloth.UnslothVisionDataCollator("model", "processor", flag = True) + + assert collator.model == "model" + assert collator.processor == "processor" + assert collator.kwargs == {"completion_only_loss": True, "flag": True} + + +def test_mlx_train_on_responses_only_returns_shared_mask_function(): + """The MLX public shim should expose the shared response-mask helper.""" + unsloth = _import_mlx_unsloth() + + class Tokenizer: + def __call__( + self, + text, + add_special_tokens = False, + ): + return types.SimpleNamespace( + input_ids = { + "": [1], + "": [2], + }[text] + ) + + def convert_tokens_to_ids(self, token): + return token + + mask_fn = unsloth.train_on_responses_only( + None, + instruction_part = "", + response_part = "", + tokenizer = Tokenizer(), + return_function = True, + ) + masked = mask_fn( + { + "input_ids": [[1, 10, 2, 20, 21, 1, 11]], + } + ) + + assert masked == {"labels": [[-100, -100, -100, 20, 21, -100, -100]]} + + last_mask_fn = unsloth.train_on_responses_only( + None, + instruction_part = "", + response_part = "", + tokenizer = Tokenizer(), + return_function = True, + last_response_only = True, + ) + last_masked = last_mask_fn( + { + "input_ids": [[1, 10, 2, 20, 1, 11, 2, 30]], + } + ) + + assert last_masked == {"labels": [[-100, -100, -100, -100, -100, -100, -100, 30]]} + + +def test_mlx_get_chat_template_uses_light_tokenizer_patch(monkeypatch): + """MLX notebooks should not import CUDA-heavy tokenizer/save helpers.""" + _import_mlx_unsloth() + from unsloth.chat_templates import get_chat_template + import unsloth_zoo.tokenizer_utils as tokenizer_utils + + class Tokenizer: + is_fast = True + padding_side = "right" + eos_token = "" + bos_token = "" + unk_token = "" + pad_token = "" + added_tokens_decoder = {} + + def fake_patch_tokenizer(model, tokenizer): + return model, tokenizer + + real_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + if name.startswith("unsloth.models") or name.startswith("unsloth.save"): + raise AssertionError(f"unexpected CUDA-heavy import: {name}") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(tokenizer_utils, "patch_tokenizer", fake_patch_tokenizer) + monkeypatch.setattr(builtins, "__import__", guarded_import) + + tokenizer = get_chat_template( + Tokenizer(), + chat_template = ("{{ messages }}", ""), + ) + + assert tokenizer.chat_template == "{{ messages }}" + assert tokenizer.padding_side == "right" + + +def test_mlx_gpu_memory_stats_helper_shape(): + """The portable memory helper should return CUDA-shaped values.""" + unsloth = _import_mlx_unsloth() + + stats, used, total = unsloth.get_gpu_memory_stats() + + assert isinstance(stats.name, str) + assert hasattr(stats, "total_memory") + assert isinstance(used, float) + assert total > 0 + + +def test_mlx_torch_cuda_compatibility_shim(): + """Existing CUDA memory and move calls should run on MLX.""" + unsloth = _import_mlx_unsloth() + torch = pytest.importorskip("torch") + from transformers.tokenization_utils_base import BatchEncoding + + stats, used, total = unsloth.get_gpu_memory_stats() + cuda_stats = torch.cuda.get_device_properties(0) + + assert cuda_stats.name == stats.name + assert cuda_stats.total_memory == stats.total_memory + assert torch.cuda.get_device_name(0) == stats.name + assert torch.cuda.max_memory_reserved() == int(used * 1024 * 1024 * 1024) + assert torch.cuda.max_memory_allocated() == torch.cuda.max_memory_reserved() + # current (non-max) APIs report live active memory, not the peak high-water + # mark, and never exceed it. + assert 0 <= torch.cuda.memory_reserved() <= torch.cuda.max_memory_reserved() + assert torch.cuda.memory_allocated() == torch.cuda.memory_reserved() + assert torch.cuda.device_count() == 1 + assert torch.cuda.current_device() == 0 + assert torch.cuda.get_device_capability() == (0, 0) + assert total > 0 + + free_bytes, total_bytes = torch.cuda.mem_get_info() + assert total_bytes == int(total * 1024 * 1024 * 1024) + assert 0 <= free_bytes <= total_bytes + + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + torch.cuda.set_device(0) + + tensor = torch.tensor([1, 2, 3]) + assert tensor.to("cuda") is tensor + assert tensor.cuda() is tensor + assert tensor.to(device = "cuda") is tensor + assert tensor.to("cuda", dtype = torch.float32).dtype == torch.float32 + + batch = BatchEncoding({"input_ids": tensor}) + assert batch.to("cuda") is batch + assert batch.to(device = "cuda") is batch diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 7a63dcfb85..275fe7ac57 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -403,10 +403,14 @@ def cmd_train(args) -> int: metrics["gguf_dir"] = str(gguf_dir) with Phase("save_gguf", metrics): try: + # q8_0 (the exporter default), not bf16: llama.cpp has optimized q8_0 + # CPU kernels, whereas bf16 CPU decode is unusably slow on the runner + # and made the fresh-process llama-cli reload below time out. q8_0 is + # also what users deploy by default. model.save_pretrained_gguf( str(gguf_dir), tokenizer = tokenizer, - quantization_method = "not_quantized", + quantization_method = "fast_quantized", ) gguf_files = sorted(gguf_dir.glob("*.gguf")) if not gguf_files: @@ -565,31 +569,36 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int: raise SystemExit(f"no .gguf files in {save_dir}") gguf_path = gguf_files[0] - # This is a save/reload-integrity smoke; a few generated tokens are enough. - # Keep llama.cpp bounded on macOS runners where BF16 GGUF decode is CPU-bound. + # Save/reload-integrity smoke (assert below only needs a few chars). The GGUF is + # exported q8_0 (see save_gguf) because llama.cpp bf16 CPU decode is unusably slow + # on the runner. Run CPU-only (-ngl 0), cap the context (-c 256, the model + # advertises 32768), and keep generation short; all env-tunable. n_predict = os.environ.get("UNSLOTH_GGUF_RELOAD_N", "8") n_threads = os.environ.get("UNSLOTH_GGUF_RELOAD_THREADS", str(os.cpu_count() or 4)) + n_ctx = os.environ.get("UNSLOTH_GGUF_RELOAD_CTX", "256") + n_gpu_layers = os.environ.get("UNSLOTH_GGUF_RELOAD_NGL", "0") reload_timeout = int(os.environ.get("UNSLOTH_GGUF_RELOAD_TIMEOUT", "420")) - + argv = [ + str(llama_cli), + "-m", + str(gguf_path), + "-p", + PROMPT, + "-n", + n_predict, + "-t", + n_threads, + "-c", + n_ctx, + "-ngl", + n_gpu_layers, + "--temp", + "0", + "--seed", + str(SEED), + "--no-warmup", + ] with Phase("reload_gguf", metrics): - argv = [ - str(llama_cli), - "-m", - str(gguf_path), - "-p", - PROMPT, - "-n", - n_predict, - "-t", - n_threads, - "--temp", - "0", - "--seed", - str(SEED), - "-c", - "256", - "--no-warmup", - ] try: proc = subprocess.run( argv, @@ -606,6 +615,7 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int: return stream.decode("utf-8", errors = "replace") return stream or "" + print(f" [reload:gguf] TIMEOUT running: {' '.join(argv)}", flush = True) print(f" [reload:gguf] TIMEOUT stdout:\n{_decode(exc.stdout)[:1000]}", flush = True) print(f" [reload:gguf] TIMEOUT stderr:\n{_decode(exc.stderr)[:1000]}", flush = True) raise diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 8202195ca8..04cc600725 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -33,6 +33,27 @@ if platform.system() == "Windows": 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 @@ -66,7 +87,12 @@ if _IS_MLX: # 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 + 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( @@ -75,6 +101,53 @@ if _IS_MLX: "`pip install -U unsloth-zoo` or rerun install.sh." ) from _e + import dataclasses as _dataclasses + import importlib.machinery as _machinery + import sys as _sys + import types as _types + import warnings as _warnings + + __version__ = unsloth_zoo.__version__ + DEVICE_TYPE = "mlx" + + 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 @@ -89,9 +162,6 @@ if _IS_MLX: TextPreprocessor = _raw_text.TextPreprocessor del _raw_text, _raw_text_spec, _raw_text_path, _Path - __version__ = unsloth_zoo.__version__ - DEVICE_TYPE = "mlx" - class FastLanguageModel: @staticmethod def from_pretrained(*args, **kwargs): @@ -141,14 +211,1202 @@ if _IS_MLX: 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", + ) + _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 _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["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( - "Unsloth: UnslothVisionDataCollator is not used on MLX. " - "Use the MLX trainer/data path instead." + 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 diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 60eb8de5d7..169b2dbd0e 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -27,18 +27,25 @@ __all__ = [ "test_construct_chat_template", ] -from transformers import StoppingCriteria, StoppingCriteriaList -from torch import LongTensor, FloatTensor -from transformers.models.llama.modeling_llama import logger +from transformers.utils import logging +try: + from torch import LongTensor, FloatTensor +except ImportError: + LongTensor = FloatTensor = None +logger = logging.get_logger(__name__) import os import shutil -from .tokenizer_utils import * import re from .ollama_template_mappers import OLLAMA_TEMPLATES -from unsloth_zoo.dataset_utils import ( - train_on_responses_only, - standardize_data_formats, -) +try: + from unsloth_zoo.dataset_utils import ( + train_on_responses_only, + standardize_data_formats, + ) +except ImportError: + # dataset_utils pulls torch; keep chat_templates importable on torch-free + # (MLX) hosts, which expose these via the backend-specific wrappers instead. + train_on_responses_only = standardize_data_formats = None standardize_sharegpt = standardize_data_formats CHAT_TEMPLATES = {} DEFAULT_SYSTEM_MESSAGE = {} @@ -1838,11 +1845,24 @@ def get_chat_template( map_eos_token = True, system_message = None, patch_saving = True, - use_zoo_tokenizer_patch = False, + use_zoo_tokenizer_patch = None, ): assert(type(map_eos_token) is bool) + import sys + is_mlx_backend = getattr(sys.modules.get("unsloth"), "DEVICE_TYPE", None) == "mlx" + if use_zoo_tokenizer_patch is None: + use_zoo_tokenizer_patch = is_mlx_backend old_tokenizer = tokenizer + # mlx-lm's TokenizerWrapper._tokenizer is the HF tokenizer, not the Rust + # backend the vocab-edit paths below need; unwrap here, re-wrap before return. + _mlx_tokenizer_wrapper = None + if is_mlx_backend and tokenizer.__class__.__name__ == "TokenizerWrapper": + _inner_tokenizer = getattr(tokenizer, "_tokenizer", None) + if _inner_tokenizer is not None and hasattr(_inner_tokenizer, "is_fast"): + _mlx_tokenizer_wrapper = tokenizer + tokenizer = _inner_tokenizer + IS_GEMMA = False if tokenizer.__class__.__name__.startswith("Gemma"): if chat_template == "chatml": chat_template = "gemma_chatml" @@ -1952,6 +1972,7 @@ def get_chat_template( pass # Must fix the sentence piece tokenizer since there's no tokenizer.model file! + from .tokenizer_utils import fix_sentencepiece_tokenizer tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,) else: pass @@ -1997,6 +2018,7 @@ def get_chat_template( # Must fix the sentence piece tokenizer since there's no tokenizer.model file! token_mapping = { old_eos_token : stop_word, } + from .tokenizer_utils import fix_sentencepiece_tokenizer tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,) pass @@ -2057,13 +2079,25 @@ def get_chat_template( # stopping_criteria = create_stopping_criteria(tokenizer, stop_word) # Patch saving functions - if patch_saving: + if patch_saving and not is_mlx_backend: from .save import patch_saving_functions tokenizer = patch_saving_functions(tokenizer) # Add Ollama tokenizer._ollama_modelfile = ollama_modelfile tokenizer._system_message = system_message + + # Re-wrap so the trainer gets the same TokenizerWrapper type back. + if _mlx_tokenizer_wrapper is not None: + _mlx_tokenizer_wrapper._tokenizer = tokenizer + eos_token_id = getattr(tokenizer, "eos_token_id", None) + if eos_token_id is not None: + _mlx_tokenizer_wrapper._eos_token_ids = {eos_token_id} + _mlx_tokenizer_wrapper._chat_template = None + _mlx_tokenizer_wrapper.has_chat_template = ( + getattr(tokenizer, "chat_template", None) is not None + ) + tokenizer = _mlx_tokenizer_wrapper return tokenizer#, stopping_criteria @@ -2749,6 +2783,15 @@ extra_eos_tokens = None, def create_stopping_criteria(tokenizer, stop_word = "eos_token"): + try: + import torch + from transformers import StoppingCriteria, StoppingCriteriaList + except ImportError as exc: + raise ImportError( + "Unsloth: create_stopping_criteria requires PyTorch and is only " + "supported on Torch backends." + ) from exc + class StoppingCriteriaSub(StoppingCriteria): __slots__ = "stop_token", "single_match", "length", @@ -2828,10 +2871,10 @@ def test_chat_templates(): for j in range(len(messages)-1): correct_prompt.append_message(correct_prompt.roles[j%2==1], messages[j+1]["content"]) correct_prompt.append_message(correct_prompt.roles[1], "") - correct_prompt = tokenizer.bos_token + correct_prompt.get_prompt() template = vicuna_template correct_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") + correct_prompt = correct_tokenizer.bos_token + correct_prompt.get_prompt() correct_tokenizer.chat_template = template our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) assert(correct_prompt == our_prompt) @@ -2845,10 +2888,10 @@ def test_chat_templates(): for j in range(len(messages)-1): correct_prompt.append_message(correct_prompt.roles[j%2==1], messages[j+1]["content"]) correct_prompt.append_message(correct_prompt.roles[1], "") - correct_prompt = tokenizer.bos_token + correct_prompt.get_prompt() template = vicuna_old_template correct_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") + correct_prompt = correct_tokenizer.bos_token + correct_prompt.get_prompt() correct_tokenizer.chat_template = template our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) # We add ourselves