diff --git a/tests/python/test_gefenx_optimizer.py b/tests/python/test_gefenx_optimizer.py index 376eadfe4d..fe18d22dea 100644 --- a/tests/python/test_gefenx_optimizer.py +++ b/tests/python/test_gefenx_optimizer.py @@ -190,6 +190,41 @@ def test_build_gefenx_forwards_config_and_falls_back_to_trainer_betas(fake_gefen assert opt.param_groups == fake_gefen["gefen"]["params"] +def test_extra_kwargs_reserved_keys_are_dropped(fake_gefen): + # Reserved keys (would collide with the builders' explicit args -> TypeError) + # are dropped from extra_kwargs with a warning; non-reserved keys pass through. + model = _FakeModel([("w", _FakeParam())]) + config = _GefenXConfig( + extra_kwargs={"lr": 9.9, "weight_decay": 9.9, "codebook_refresh_every": "50"} + ) + with pytest.warns(UserWarning, match="reserved key"): + gefenx.build_gefenx_optimizer( + model, config, lr=1e-4, weight_decay=0.01, + betas=(0.9, 0.999), eps=1e-8, + ) + kw = fake_gefen["gefen"]["kwargs"] + # Build succeeds (no duplicate-keyword TypeError) and the reserved extra_kwargs + # values did NOT override the real builder args (9.9 dropped, 1e-4 / 0.01 win). + assert kw["lr"] == 1e-4 + assert kw["weight_decay"] == 0.01 + assert kw["codebook_refresh_every"] == 50 # allowed key -> coerced + + +def test_muon_extra_kwargs_reserved_backup_substrings_dropped(fake_gefen): + model = _FakeModel([("w", _FakeParam())]) + config = _GefenXMuonConfig(extra_kwargs={"backup_substrings": ["x"], "ns_steps": "7"}) + 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"] + # backup_substrings is passed as an explicit arg, not via **kwargs. + assert "backup_substrings" not in kw + assert fake_gefen["muon"]["backup_substrings"] is None + assert kw["ns_steps"] == 7 + + def test_build_gefenx_config_betas_override_and_extra_kwargs(fake_gefen): model = _FakeModel([("w", _FakeParam())]) config = _GefenXConfig( @@ -298,7 +333,10 @@ def _cuda_available(): try: import torch - return torch.cuda.is_available() + # NVIDIA CUDA only. On ROCm/HIP torch.cuda.is_available() is also True, but + # the Gefen-X gate rejects HIP — so the fused tests must skip there too, + # otherwise they'd error on _require_nvidia_cuda() instead of skipping. + return torch.cuda.is_available() and getattr(torch.version, "hip", None) is None except Exception: return False @@ -366,7 +404,7 @@ def test_real_gefenx_muon_updates_all_params_cpu(): assert _num_changed(torch, before, model) == len(before) -@pytest.mark.skipif(not _CUDA, reason="requires CUDA for the fused gefen kernels") +@pytest.mark.skipif(not _CUDA, reason="requires NVIDIA CUDA for the fused gefen kernels") def test_real_gefenx_cuda_fused_updates_all_params(): import torch pytest.importorskip("gefen") @@ -384,7 +422,7 @@ def test_real_gefenx_cuda_fused_updates_all_params(): assert _num_changed(torch, before, model) == len(before) -@pytest.mark.skipif(not _CUDA, reason="requires CUDA for the fused gefen kernels") +@pytest.mark.skipif(not _CUDA, reason="requires NVIDIA CUDA for the fused gefen kernels") def test_real_gefenx_muon_cuda_fused_updates_all_params(): import torch pytest.importorskip("gefen") @@ -465,3 +503,18 @@ def test_trainer_create_optimizer_dispatches_gefenx_muon(tmp_path): opt.step() opt.zero_grad() assert _num_changed(torch, before, trainer.model) == len(before) + + +def test_conflicting_optimizer_configs_raise(tmp_path): + pytest.importorskip("unsloth") + from unsloth import GefenXConfig, GefenXMuonConfig + from unsloth.trainer import UnslothTrainingArguments + + # Setting both Gefen-X configs is ambiguous (dispatch would silently pick one). + with pytest.raises(ValueError, match="mutually exclusive"): + UnslothTrainingArguments( + output_dir=str(tmp_path / "conflict"), + gefenx_config=GefenXConfig(), + gefenx_muon_config=GefenXMuonConfig(), + report_to="none", + ) diff --git a/unsloth/optimizers/gefenx.py b/unsloth/optimizers/gefenx.py index ee7c8b9719..2950f27852 100644 --- a/unsloth/optimizers/gefenx.py +++ b/unsloth/optimizers/gefenx.py @@ -23,8 +23,18 @@ definition. from __future__ import annotations +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. +_RESERVED_EXTRA_KWARGS = frozenset( + {"params", "lr", "weight_decay", "backup_substrings", "backup_lr_scale"} +) + def _require_nvidia_cuda() -> None: """Gate the Gefen-X optimizers to NVIDIA CUDA. @@ -134,7 +144,14 @@ def _collect_kwargs(config, fields: Tuple[str, ...]) -> Dict[str, Any]: kwargs[name] = value extra = getattr(config, "extra_kwargs", None) if extra: - kwargs.update({k: coerce_optim_arg(v) for k, v in extra.items()}) + for key, value in extra.items(): + if key in _RESERVED_EXTRA_KWARGS: + warnings.warn( + f"Unsloth: ignoring reserved key {key!r} in Gefen-X extra_kwargs; " + f"set it via the config field / learning_rate / weight_decay instead." + ) + continue + kwargs[key] = coerce_optim_arg(value) return kwargs diff --git a/unsloth/trainer.py b/unsloth/trainer.py index aa8f2c6c99..b7223af2d9 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -273,6 +273,26 @@ class UnslothTrainingArguments(TrainingArguments): self.gefenx_config = gefenx_config self.gefenx_muon_config = gefenx_muon_config self.embedding_learning_rate = embedding_learning_rate + + # q_galore_config, gefenx_config and gefenx_muon_config are mutually + # exclusive — create_optimizer dispatches to the first one set, so + # supplying more than one would silently ignore the rest. Fail loudly. + _optimizer_configs = [ + name + for name, cfg in ( + ("q_galore_config", q_galore_config), + ("gefenx_config", gefenx_config), + ("gefenx_muon_config", gefenx_muon_config), + ) + if cfg is not None + ] + if len(_optimizer_configs) > 1: + raise ValueError( + f"Unsloth: only one optimizer config may be set, but got " + f"{_optimizer_configs}. q_galore_config, gefenx_config and " + f"gefenx_muon_config are mutually exclusive." + ) + super().__init__(*args, **kwargs) self.embedding_learning_rate = embedding_learning_rate @@ -467,7 +487,7 @@ class UnslothTrainer(SFTTrainer): embedding_lr = embedding_lr, ) n_params = sum( - len(g["params"]) for g in self.optimizer.param_groups + p.numel() for g in self.optimizer.param_groups for p in g["params"] ) print( f"🦥 Unsloth: Gefen-X optimizer enabled — "