fix: address Gemini Code Assist review on the Gefen-X integration

- optimizers/gefenx.py: expand _RESERVED_EXTRA_KWARGS with model / embedding_lr
  / embedding_learning_rate so a stray extra_kwargs entry can't collide with the
  positional model arg or land as an unexpected keyword; skip None-valued config
  fields in _collect_kwargs so gefen keeps its own runtime defaults instead of
  being handed an explicit None (muon_lr / muon_weight_decay / backup_weight_decay).
- tests: cover the new reserved keys (model / embedding_*) and the None-skip
  behaviour. 31 tests pass.

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(

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: