unsloth/tests/python/test_gefenx_optimizer.py

"""Unit tests for the Gefen-X optimizer integration (``unsloth.optimizers.gefenx``).

These tests exercise the config→constructor mapping and parameter routing WITHOUT
requiring a GPU, the real ``gefen`` package, or a full ``import unsloth`` (which
triggers GPU init). The helper module has no relative imports, so it is loaded by
file path and handed a fake ``gefen`` module that records constructor arguments.
"""

import importlib.util
import pathlib
import sys
import types
from dataclasses import dataclass, field
from typing import List, Optional

import pytest

# --- Load unsloth/optimizers/gefenx.py standalone (no unsloth package import) ---
_MODULE_PATH = (
    pathlib.Path(__file__).resolve().parents[2] / "unsloth" / "optimizers" / "gefenx.py"
)
_spec = importlib.util.spec_from_file_location("_unsloth_gefenx_under_test", _MODULE_PATH)
gefenx = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(gefenx)


# --- Minimal stand-ins so the tests need neither torch nor unsloth.trainer -------
@dataclass
class _GefenXConfig:
    fused: bool = True
    factored_v_2d: bool = True
    force_1d_period_one: bool = False
    force_2d_period_one: bool = False
    period_one_substrings: tuple = ()
    codebook_refresh_every: int = 0
    stochastic_round: bool = False
    capturable: bool = False
    betas: Optional[tuple] = None
    eps: Optional[float] = None
    extra_kwargs: dict = field(default_factory=dict)


@dataclass
class _GefenXMuonConfig:
    fused: bool = True
    adjust_lr_fn: str = "match_rms_adamw"
    backup_lr_scale: Optional[float] = 0.5
    backup_lr: Optional[float] = None
    muon_lr: Optional[float] = None
    muon_weight_decay: Optional[float] = None
    backup_weight_decay: Optional[float] = None
    backup_1d_period_one: bool = True
    backup_2d_period_one: bool = False
    momentum: float = 0.95
    nesterov: bool = True
    ns_steps: int = 5
    ns_schedule: str = "tuned3"
    sharded_mode: str = "exact"
    fp8_ns: bool = False
    stochastic_round: bool = False
    normuon: bool = True
    cautious: bool = False
    capturable: bool = False
    backup_substrings: Optional[List[str]] = None
    betas: Optional[tuple] = None
    eps: Optional[float] = None
    extra_kwargs: dict = field(default_factory=dict)


class _FakeParam:
    """Enough of an nn.Parameter for make_gefenx_param_groups / grouping."""

    def __init__(self, requires_grad=True):
        self.requires_grad = requires_grad


class _FakeModel:
    def __init__(self, named):
        self._named = named

    def named_parameters(self):
        return list(self._named)


@pytest.fixture
def fake_gefen(monkeypatch):
    """Install a fake ``gefen`` module that records what got constructed."""
    captured = {}

    class _FakeGefen:
        def __init__(self, params, **kwargs):
            captured["gefen"] = {"params": params, "kwargs": kwargs}
            # Mimic torch.optim.Optimizer.param_groups shape for downstream code.
            self.param_groups = params

    def _from_model(model, *, backup_substrings=None, **kwargs):
        captured["muon"] = {
            "model": model,
            "backup_substrings": backup_substrings,
            "kwargs": kwargs,
        }
        return "MUON_OPTIMIZER"

    class _FakeHybrid:
        from_model = staticmethod(_from_model)

    module = types.ModuleType("gefen")
    module.Gefen = _FakeGefen
    module.GefenMuonHybrid = _FakeHybrid
    monkeypatch.setitem(sys.modules, "gefen", module)
    return captured


# --------------------------------------------------------------------------- #
# coerce_optim_arg
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize(
    "raw,expected",
    [
        ("true", True),
        ("False", False),
        ("none", None),
        ("null", None),
        ("5", 5),
        ("6.0e-6", 6.0e-6),
        ("match_rms_adamw", "match_rms_adamw"),
        (True, True),      # non-strings pass through unchanged
        (0.5, 0.5),
        (None, None),
    ],
)
def test_coerce_optim_arg(raw, expected):
    assert gefenx.coerce_optim_arg(raw) == expected


# --------------------------------------------------------------------------- #
# make_gefenx_param_groups
# --------------------------------------------------------------------------- #
def test_param_groups_single_group_without_embedding_lr():
    model = _FakeModel(
        [
            ("model.layers.0.self_attn.q_proj.weight", _FakeParam()),
            ("model.embed_tokens.modules_to_save.default.weight", _FakeParam()),
            ("model.layers.0.frozen.weight", _FakeParam(requires_grad=False)),
        ]
    )
    groups = gefenx.make_gefenx_param_groups(model, lr=1e-4, weight_decay=0.01)
    # No embedding_lr => one group; frozen params excluded.
    assert len(groups) == 1
    assert len(groups[0]["params"]) == 2
    assert groups[0]["lr"] == 1e-4
    assert groups[0]["weight_decay"] == 0.01


def test_param_groups_splits_embedding_lr():
    model = _FakeModel(
        [
            ("model.layers.0.self_attn.q_proj.weight", _FakeParam()),
            ("model.embed_tokens.modules_to_save.default.weight", _FakeParam()),
        ]
    )
    groups = gefenx.make_gefenx_param_groups(
        model, lr=1e-4, weight_decay=0.0, embedding_lr=5e-6
    )
    assert len(groups) == 2
    non_embed, embed = groups
    assert non_embed["lr"] == 1e-4 and len(non_embed["params"]) == 1
    assert embed["lr"] == 5e-6 and len(embed["params"]) == 1


# --------------------------------------------------------------------------- #
# build_gefenx_optimizer
# --------------------------------------------------------------------------- #
def test_build_gefenx_forwards_config_and_falls_back_to_trainer_betas(fake_gefen):
    model = _FakeModel([("w", _FakeParam())])
    config = _GefenXConfig(fused=True, factored_v_2d=False, stochastic_round=True)
    opt = gefenx.build_gefenx_optimizer(
        model, config, lr=1e-4, weight_decay=0.01,
        betas=(0.9, 0.95), eps=1e-8,
    )
    kw = fake_gefen["gefen"]["kwargs"]
    assert kw["fused"] is True
    assert kw["factored_v_2d"] is False
    assert kw["stochastic_round"] is True
    # betas/eps not set on the config => inherit the trainer's AdamW values.
    assert kw["betas"] == (0.9, 0.95)
    assert kw["eps"] == 1e-8
    # Empty period_one_substrings is dropped, not forwarded as ().
    assert "period_one_substrings" not in kw
    assert opt.param_groups == fake_gefen["gefen"]["params"]


def test_build_gefenx_config_betas_override_and_extra_kwargs(fake_gefen):
    model = _FakeModel([("w", _FakeParam())])
    config = _GefenXConfig(
        betas=(0.8, 0.9), eps=1e-6,
        period_one_substrings=("embed", "lm_head"),
        extra_kwargs={"codebook_refresh_every": "100"},  # string coerced to int
    )
    gefenx.build_gefenx_optimizer(
        model, config, lr=1e-4, weight_decay=0.0, betas=(0.9, 0.999), eps=1e-8,
    )
    kw = fake_gefen["gefen"]["kwargs"]
    assert kw["betas"] == (0.8, 0.9)
    assert kw["eps"] == 1e-6
    assert kw["period_one_substrings"] == ("embed", "lm_head")
    assert kw["codebook_refresh_every"] == 100


# --------------------------------------------------------------------------- #
# build_gefenx_muon_optimizer
# --------------------------------------------------------------------------- #
def test_build_gefenx_muon_applies_recommended_recipe(fake_gefen):
    model = _FakeModel([("w", _FakeParam())])
    config = _GefenXMuonConfig()  # defaults encode the recipe
    opt = gefenx.build_gefenx_muon_optimizer(
        model, config, lr=1e-4, weight_decay=0.0, betas=(0.9, 0.999), eps=1e-8,
    )
    assert opt == "MUON_OPTIMIZER"
    call = fake_gefen["muon"]
    assert call["model"] is model
    kw = call["kwargs"]
    assert kw["backup_1d_period_one"] is True
    assert kw["adjust_lr_fn"] == "match_rms_adamw"
    assert kw["fused"] is True
    # backup_lr defaults to backup_lr_scale (0.5) * lr.
    assert kw["backup_lr"] == pytest.approx(0.5 * 1e-4)
    assert kw["lr"] == 1e-4


def test_build_gefenx_muon_explicit_backup_lr_wins(fake_gefen):
    model = _FakeModel([("w", _FakeParam())])
    config = _GefenXMuonConfig(backup_lr=3e-5, backup_lr_scale=0.5)
    gefenx.build_gefenx_muon_optimizer(
        model, config, lr=1e-4, weight_decay=0.0, betas=(0.9, 0.999), eps=1e-8,
    )
    assert fake_gefen["muon"]["kwargs"]["backup_lr"] == 3e-5


def test_build_gefenx_muon_backup_lr_scale_none_leaves_backup_lr_unset(fake_gefen):
    model = _FakeModel([("w", _FakeParam())])
    config = _GefenXMuonConfig(backup_lr_scale=None)
    gefenx.build_gefenx_muon_optimizer(
        model, config, lr=1e-4, weight_decay=0.0, betas=(0.9, 0.999), eps=1e-8,
    )
    # No scale and no explicit backup_lr => gefen uses its own default (None).
    assert fake_gefen["muon"]["kwargs"].get("backup_lr") is None


def test_build_gefenx_muon_passes_lr_weight_decay_and_backup_substrings(fake_gefen):
    model = _FakeModel([("w", _FakeParam())])
    config = _GefenXMuonConfig(backup_substrings=["router", "gate"])
    gefenx.build_gefenx_muon_optimizer(
        model, config, lr=2e-4, weight_decay=0.05, betas=(0.9, 0.999), eps=1e-8,
    )
    call = fake_gefen["muon"]
    assert call["backup_substrings"] == ["router", "gate"]
    assert call["kwargs"]["lr"] == 2e-4
    assert call["kwargs"]["weight_decay"] == 0.05


# --------------------------------------------------------------------------- #
# Device gate: NVIDIA CUDA only (AMD/ROCm and Intel XPU rejected)
# --------------------------------------------------------------------------- #
def test_gate_rejects_rocm_hip_build(fake_gefen, monkeypatch):
    torch = pytest.importorskip("torch")
    # Simulate an AMD/ROCm PyTorch build by tagging torch.version.hip.
    monkeypatch.setattr(torch.version, "hip", "6.0.0", raising=False)
    model = _FakeModel([("w", _FakeParam())])
    with pytest.raises(RuntimeError, match="ROCm|HIP|CUDA"):
        gefenx.build_gefenx_optimizer(
            model, _GefenXConfig(), lr=1e-4, weight_decay=0.0,
            betas=(0.9, 0.999), eps=1e-8,
        )
    with pytest.raises(RuntimeError, match="ROCm|HIP|CUDA"):
        gefenx.build_gefenx_muon_optimizer(
            model, _GefenXMuonConfig(), lr=1e-4, weight_decay=0.0,
            betas=(0.9, 0.999), eps=1e-8,
        )
    # The gate fires before gefen is even imported/constructed.
    assert "gefen" not in fake_gefen and "muon" not in fake_gefen


def test_gate_allows_non_hip(monkeypatch):
    torch = pytest.importorskip("torch")
    # A normal (non-HIP) build must pass the gate without raising.
    monkeypatch.setattr(torch.version, "hip", None, raising=False)
    gefenx._require_nvidia_cuda()  # should not raise


# --------------------------------------------------------------------------- #
# End-to-end against the REAL gefen package + REAL Unsloth trainer path.
#
# These construct real gefen optimizers on a real torch model, take a real step,
# and assert every trainable parameter actually moved (a no-op step would fail).
# --------------------------------------------------------------------------- #
def _cuda_available():
    try:
        import torch

        return torch.cuda.is_available()
    except Exception:
        return False


_CUDA = _cuda_available()


def _tiny_model(torch, device="cpu"):
    # Embedding + 2D Linears + LayerNorm exercises all of gefen's routing buckets:
    # 2D hidden weights (Muon half), embedding/1D norm/bias (Gefen backup half).
    return torch.nn.Sequential(
        torch.nn.Embedding(16, 8),
        torch.nn.Linear(8, 8),
        torch.nn.LayerNorm(8),
        torch.nn.Linear(8, 8),
    ).to(device)


def _forward_backward(torch, model, device="cpu"):
    ids = torch.arange(4, device=device)
    out = model[0](ids)
    out = model[1](out)
    out = model[2](out)
    out = model[3](out)
    out.sum().backward()


def _num_changed(torch, before, model):
    return sum(1 for a, p in zip(before, model.parameters()) if not torch.equal(a, p))


def test_real_gefenx_updates_all_params_cpu():
    torch = pytest.importorskip("torch")
    pytest.importorskip("gefen")
    from gefen import Gefen

    model = _tiny_model(torch)
    before = [p.detach().clone() for p in model.parameters()]
    opt = gefenx.build_gefenx_optimizer(
        model, _GefenXConfig(fused=False),
        lr=1e-3, weight_decay=0.0, betas=(0.9, 0.999), eps=1e-8,
    )
    assert isinstance(opt, Gefen)
    _forward_backward(torch, model)
    opt.step()
    opt.zero_grad()
    assert _num_changed(torch, before, model) == len(before)


def test_real_gefenx_muon_updates_all_params_cpu():
    torch = pytest.importorskip("torch")
    pytest.importorskip("gefen")
    from gefen import GefenMuonHybrid

    model = _tiny_model(torch)
    before = [p.detach().clone() for p in model.parameters()]
    opt = gefenx.build_gefenx_muon_optimizer(
        model, _GefenXMuonConfig(fused=False),
        lr=1e-3, weight_decay=0.0, betas=(0.9, 0.999), eps=1e-8,
    )
    assert isinstance(opt, GefenMuonHybrid)
    _forward_backward(torch, model)
    opt.step()
    opt.zero_grad()
    assert _num_changed(torch, before, model) == len(before)


@pytest.mark.skipif(not _CUDA, reason="requires CUDA for the fused gefen kernels")
def test_real_gefenx_cuda_fused_updates_all_params():
    import torch
    pytest.importorskip("gefen")

    model = _tiny_model(torch, "cuda")
    before = [p.detach().clone() for p in model.parameters()]
    opt = gefenx.build_gefenx_optimizer(
        model, _GefenXConfig(fused=True),
        lr=1e-3, weight_decay=0.0, betas=(0.9, 0.999), eps=1e-8,
    )
    _forward_backward(torch, model, "cuda")
    opt.step()
    opt.zero_grad()
    torch.cuda.synchronize()
    assert _num_changed(torch, before, model) == len(before)


@pytest.mark.skipif(not _CUDA, reason="requires CUDA for the fused gefen kernels")
def test_real_gefenx_muon_cuda_fused_updates_all_params():
    import torch
    pytest.importorskip("gefen")

    model = _tiny_model(torch, "cuda")
    before = [p.detach().clone() for p in model.parameters()]
    opt = gefenx.build_gefenx_muon_optimizer(
        model, _GefenXMuonConfig(fused=True),
        lr=1e-3, weight_decay=0.0, betas=(0.9, 0.999), eps=1e-8,
    )
    _forward_backward(torch, model, "cuda")
    opt.step()
    opt.zero_grad()
    torch.cuda.synchronize()
    assert _num_changed(torch, before, model) == len(before)


# --------------------------------------------------------------------------- #
# Full Unsloth trainer path: the REAL GefenXConfig / GefenXMuonConfig dataclasses
# carried through the REAL UnslothTrainingArguments and dispatched by the REAL
# UnslothTrainer.create_optimizer (not the standalone build_* helpers).
# --------------------------------------------------------------------------- #
def test_trainer_create_optimizer_dispatches_gefenx(tmp_path):
    torch = pytest.importorskip("torch")
    pytest.importorskip("gefen")
    pytest.importorskip("unsloth")
    from unsloth import GefenXConfig
    from unsloth.trainer import UnslothTrainer, UnslothTrainingArguments
    from gefen import Gefen

    args = UnslothTrainingArguments(
        output_dir=str(tmp_path / "gx"),
        gefenx_config=GefenXConfig(fused=False),
        learning_rate=1e-3, weight_decay=0.0, report_to="none",
    )
    # Config plumbing on the real dataclass-typed argument.
    assert args.gefenx_config is not None
    assert args.gefenx_muon_config is None

    trainer = UnslothTrainer.__new__(UnslothTrainer)  # skip heavy SFTTrainer.__init__
    trainer.model = _tiny_model(torch)
    trainer.args = args
    trainer.optimizer = None

    before = [p.detach().clone() for p in trainer.model.parameters()]
    opt = trainer.create_optimizer()
    assert isinstance(opt, Gefen)
    _forward_backward(torch, trainer.model)
    opt.step()
    opt.zero_grad()
    assert _num_changed(torch, before, trainer.model) == len(before)


def test_trainer_create_optimizer_dispatches_gefenx_muon(tmp_path):
    torch = pytest.importorskip("torch")
    pytest.importorskip("gefen")
    pytest.importorskip("unsloth")
    from unsloth import GefenXMuonConfig
    from unsloth.trainer import UnslothTrainer, UnslothTrainingArguments
    from gefen import GefenMuonHybrid

    args = UnslothTrainingArguments(
        output_dir=str(tmp_path / "gm"),
        gefenx_muon_config=GefenXMuonConfig(fused=False),
        learning_rate=1e-3, weight_decay=0.0, report_to="none",
    )
    assert args.gefenx_muon_config is not None

    trainer = UnslothTrainer.__new__(UnslothTrainer)
    trainer.model = _tiny_model(torch)
    trainer.args = args
    trainer.optimizer = None

    before = [p.detach().clone() for p in trainer.model.parameters()]
    opt = trainer.create_optimizer()
    assert isinstance(opt, GefenMuonHybrid)
    _forward_backward(torch, trainer.model)
    opt.step()
    opt.zero_grad()
    assert _num_changed(torch, before, trainer.model) == len(before)