diff --git a/tests/python/test_gefenx_optimizer.py b/tests/python/test_gefenx_optimizer.py index 55f2f12967..89fce7a4d3 100644 --- a/tests/python/test_gefenx_optimizer.py +++ b/tests/python/test_gefenx_optimizer.py @@ -238,6 +238,47 @@ def test_muon_extra_kwargs_reserved_backup_substrings_dropped(fake_gefen): assert kw["ns_steps"] == 7 +def test_extra_kwargs_reserved_model_and_embedding_keys_dropped(fake_gefen): + # model / embedding_lr / embedding_learning_rate would collide (model) or be + # unexpected keywords; they are dropped from extra_kwargs, and the real model + # is still routed through from_model positionally. + model = _FakeModel([("w", _FakeParam())]) + config = _GefenXMuonConfig( + extra_kwargs = {"model": "oops", "embedding_lr": 1e-5, "embedding_learning_rate": 1e-5} + ) + with pytest.warns(UserWarning, match = "reserved key"): + gefenx.build_gefenx_muon_optimizer( + model, + config, + lr = 1e-4, + weight_decay = 0.0, + betas = (0.9, 0.999), + eps = 1e-8, + ) + kw = fake_gefen["muon"]["kwargs"] + assert "model" not in kw and "embedding_lr" not in kw + assert "embedding_learning_rate" not in kw + assert fake_gefen["muon"]["model"] is model + + +def test_none_config_fields_are_not_forwarded(fake_gefen): + # None config fields are "unset" -> not passed, so gefen keeps its own defaults. + model = _FakeModel([("w", _FakeParam())]) + config = _GefenXMuonConfig(muon_lr = None, muon_weight_decay = None, backup_weight_decay = None) + gefenx.build_gefenx_muon_optimizer( + model, + config, + lr = 1e-4, + weight_decay = 0.0, + betas = (0.9, 0.999), + eps = 1e-8, + ) + kw = fake_gefen["muon"]["kwargs"] + assert "muon_lr" not in kw + assert "muon_weight_decay" not in kw + assert "backup_weight_decay" not in kw + + def test_build_gefenx_config_betas_override_and_extra_kwargs(fake_gefen): model = _FakeModel([("w", _FakeParam())]) config = _GefenXConfig( diff --git a/unsloth/optimizers/gefenx.py b/unsloth/optimizers/gefenx.py index 5817a05fc4..0cc284e126 100644 --- a/unsloth/optimizers/gefenx.py +++ b/unsloth/optimizers/gefenx.py @@ -27,12 +27,23 @@ import warnings from typing import Any, Dict, List, Optional, Tuple # Keys the builders pass explicitly (positionally or by keyword) to the gefen -# constructors. Allowing them through the ``extra_kwargs`` escape hatch would -# duplicate a keyword argument (TypeError) or shadow ``learning_rate``, so they -# are dropped from ``extra_kwargs`` with a warning — set them via the dedicated -# config field / ``learning_rate`` / ``weight_decay`` instead. +# constructors, or that Unsloth consumes elsewhere (embedding LR / param routing). +# Allowing them through the ``extra_kwargs`` escape hatch would duplicate a keyword +# argument (TypeError), shadow ``learning_rate``, or land as an unexpected keyword, +# so they are dropped from ``extra_kwargs`` with a warning — set them via the +# dedicated config field / ``learning_rate`` / ``weight_decay`` / +# ``embedding_learning_rate`` instead. _RESERVED_EXTRA_KWARGS = frozenset( - {"params", "lr", "weight_decay", "backup_substrings", "backup_lr_scale"} + { + "params", + "model", + "lr", + "weight_decay", + "backup_substrings", + "backup_lr_scale", + "embedding_lr", + "embedding_learning_rate", + } ) @@ -141,6 +152,11 @@ def _collect_kwargs(config, fields: Tuple[str, ...]) -> Dict[str, Any]: # so we don't override with () when the user never set it. if name == "period_one_substrings" and not value: continue + # None means "unset" in our configs (e.g. muon_lr / *_weight_decay) — + # skip it so gefen applies its own runtime default instead of being + # handed an explicit None. + if value is None: + continue kwargs[name] = value extra = getattr(config, "extra_kwargs", None) if extra: