diff --git a/tests/python/test_mlx_public_trainer_api.py b/tests/python/test_mlx_public_trainer_api.py index 2c33f86af1..89f304c76d 100644 --- a/tests/python/test_mlx_public_trainer_api.py +++ b/tests/python/test_mlx_public_trainer_api.py @@ -115,6 +115,9 @@ def test_mlx_training_arguments_accept_trl_style_kwargs(): 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() + supported_eval_kwargs = {} + if "eval_strategy" in unsloth._MLX_TRAINING_CONFIG_FIELDS: + supported_eval_kwargs = {"eval_strategy": "no", "eval_delay": 1} with warnings.catch_warnings(record = True) as caught: warnings.simplefilter("always") @@ -125,12 +128,16 @@ def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras(): remove_unused_columns = False, assistant_only_loss = False, completion_only_loss = False, + **supported_eval_kwargs, ) 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 + if supported_eval_kwargs: + assert args.eval_strategy == "no" + assert args.eval_delay == 1 assert caught == [] @@ -705,17 +712,10 @@ def test_mlx_trainer_rejects_unsafe_unsupported_sft_kwargs(): ) -def test_mlx_trainer_rejects_metrics_and_callbacks(): - """Trainer hooks should fail because MLXTrainer cannot honor them yet.""" +def test_mlx_trainer_rejects_compute_metrics(): + """compute_metrics is still unsupported by MLXTrainer.""" 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(), @@ -725,6 +725,46 @@ def test_mlx_trainer_rejects_metrics_and_callbacks(): ) +def test_mlx_trainer_accepts_callbacks(): + """Callbacks are routed to MLXTrainer when the zoo backend supports them.""" + unsloth = _import_mlx_unsloth() + from transformers import TrainerCallback + + if not unsloth._mlx_trainer_supports_kwarg("callbacks"): + pytest.skip("requires unsloth-zoo MLXTrainer callback support") + + class Callback(TrainerCallback): + pass + + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [Callback()], + ) + assert any(isinstance(cb, Callback) for cb in trainer.callback_handler.callbacks) + + +def test_mlx_trainer_rejects_callbacks_with_old_zoo(monkeypatch): + """Older unsloth-zoo builds should fail clearly instead of TypeError.""" + unsloth = _import_mlx_unsloth() + from transformers import TrainerCallback + + monkeypatch.setattr( + unsloth, + "_mlx_trainer_supports_kwarg", + lambda name: name != "callbacks", + ) + + with pytest.raises(NotImplementedError, match = "callbacks require"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [TrainerCallback()], + ) + + def test_mlx_trainer_rejects_custom_data_collator(): """MLXTrainer owns batching; custom SFT data collators must not be ignored.""" unsloth = _import_mlx_unsloth() diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 04cc600725..de5cd0f61b 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -102,6 +102,7 @@ if _IS_MLX: ) from _e import dataclasses as _dataclasses + import inspect as _inspect import importlib.machinery as _machinery import sys as _sys import types as _types @@ -109,6 +110,30 @@ if _IS_MLX: __version__ = unsloth_zoo.__version__ DEVICE_TYPE = "mlx" + _MLX_TRAINER_ACCEPTS_VAR_KWARGS = False + _MLX_TRAINER_SUPPORTED_KWARGS = frozenset() + try: + _MLX_TRAINER_INIT_PARAMETERS = _inspect.signature(MLXTrainer.__init__).parameters + _MLX_TRAINER_ACCEPTS_VAR_KWARGS = any( + param.kind is _inspect.Parameter.VAR_KEYWORD + for param in _MLX_TRAINER_INIT_PARAMETERS.values() + ) + _MLX_TRAINER_SUPPORTED_KWARGS = frozenset( + name + for name, param in _MLX_TRAINER_INIT_PARAMETERS.items() + if name != "self" + and param.kind + in ( + _inspect.Parameter.POSITIONAL_OR_KEYWORD, + _inspect.Parameter.KEYWORD_ONLY, + ) + ) + except (TypeError, ValueError): + pass + + def _mlx_trainer_supports_kwarg(name): + """Return whether the installed zoo MLXTrainer accepts a kwarg.""" + return _MLX_TRAINER_ACCEPTS_VAR_KWARGS or name in _MLX_TRAINER_SUPPORTED_KWARGS def _is_mlx_cuda_device_target(device): """Return True when a torch .to/.cuda target asks for CUDA on MLX.""" @@ -966,6 +991,7 @@ if _IS_MLX: "args", "formatting_func", "processor", + "callbacks", ) _TRL_SFT_TRAINER_POSITIONAL_KWARGS = ( "model", @@ -985,6 +1011,29 @@ if _IS_MLX: ) _MLX_TRAINER_KWARGS = frozenset(_MLX_TRAINER_POSITIONAL_KWARGS) + def _filter_supported_mlx_trainer_kwargs(trainer_kwargs): + """Drop inert/empty kwargs unsupported by this zoo MLXTrainer.""" + unsupported = { + key: value + for key, value in trainer_kwargs.items() + if not _mlx_trainer_supports_kwarg(key) + } + names = sorted( + key for key, value in unsupported.items() if _is_meaningful_mlx_extra_value(value) + ) + if names: + subject = ", ".join(names) + verb = "requires" if len(names) == 1 else "require" + raise NotImplementedError( + "Unsloth MLX: " + f"{subject} {verb} an unsloth-zoo build with " + "matching MLXTrainer support. Upgrade unsloth-zoo together " + "with unsloth." + ) + for key in unsupported: + trainer_kwargs.pop(key, None) + return trainer_kwargs + def _is_mlx_native_text_collator(collator): """HF pad/copy collators are redundant on MLX; match by class name.""" for klass in type(collator).__mro__: @@ -1162,6 +1211,7 @@ if _IS_MLX: trainer_kwargs, config_kwargs, ignored_kwargs = _split_mlx_trainer_kwargs(kwargs) _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs) + trainer_kwargs = _filter_supported_mlx_trainer_kwargs(trainer_kwargs) trainer_kwargs["args"] = _coerce_mlx_training_args( trainer_kwargs.get("args"), config_kwargs,