feat(mlx): route trainer callbacks (#6929)

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()

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,