feat: add Gefen-X (gefenx / gefenx_muon) optimizer integration

Adds the Gefen-X optimizers via Unsloth's standard config-object pattern
(mirroring QGaloreConfig): GefenXConfig wraps gefen.Gefen (≈1 byte/param AdamW
replacement) and GefenXMuonConfig wraps gefen.GefenMuonHybrid (Muon on 2D hidden
weights, Gefen on embeddings/heads/norms/biases). Pass either via
UnslothTrainingArguments; UnslothTrainer.create_optimizer dispatches to the new
_create_gefenx_optimizer / _create_gefenx_muon_optimizer builders.

- unsloth/optimizers/gefenx.py: config->constructor mapping, param routing,
  the axolotl recommended recipe defaults for the Muon hybrid
  (backup_1d_period_one, adjust_lr_fn=match_rms_adamw, fused, backup_lr=0.5*lr),
  and an NVIDIA-CUDA-only gate that rejects AMD/ROCm (HIP) and Intel XPU
  (gefen ships CUDA-only kernels). gefen is imported lazily.
- unsloth/trainer.py: GefenXConfig / GefenXMuonConfig dataclasses, argument
  plumbing on UnslothTrainingArguments, create_optimizer dispatch, __all__.
- tests: 26 tests — config mapping, param routing, the device gate, and real
  end-to-end runs against gefen (CPU + fused CUDA, plain + muon) plus the full
  UnslothTrainer.create_optimizer dispatch, all asserting parameters update.

MLX is unaffected (its separate trainer has no Gefen-X path).

tests/python/test_gefenx_optimizer.py

@@ -0,0 +1,467 @@
+"""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)

unsloth/optimizers/gefenx.py

@@ -0,0 +1,258 @@
+"""Gefen-X optimizer integration for Unsloth.
+
+Two optimizers are exposed through Unsloth's config-object pattern (mirroring
+``QGaloreConfig`` / the Muon integration):
+
+* ``gefenx``       — ``gefen.Gefen``: an AdamW-family optimizer that keeps its
+  optimizer state at roughly 1 byte per parameter (8-bit quantized momentum +
+  factored / block-shared second moment).
+* ``gefenx_muon``  — ``gefen.GefenMuonHybrid``: Muon (Newton-Schulz
+  orthogonalization) on the 2D hidden weight matrices, plain Gefen on
+  embeddings / heads / norms / biases, at the same ≈1 B/param footprint.
+
+The optimizers themselves live entirely in the upstream ``gefen`` (``gefen-x``)
+package; this module only maps Unsloth's ``GefenXConfig`` / ``GefenXMuonConfig``
+to the right constructor plus the parameter routing Unsloth already uses for its
+other optimizers. ``gefen`` is imported lazily so ``import unsloth`` does not
+require the package unless a Gefen-X optimizer is actually requested.
+
+The dataclasses live in ``unsloth.trainer`` (next to ``QGaloreConfig``); this
+module reads their fields by attribute, so it stays decoupled from their exact
+definition.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+
+
+def _require_nvidia_cuda() -> None:
+    """Gate the Gefen-X optimizers to NVIDIA CUDA.
+
+    ``gefen`` ships CUDA-only kernels (``*_cuda`` ops, ``device.type != "cuda"``
+    guards), so AMD/ROCm (HIP) and Intel XPU are unsupported. On ROCm PyTorch
+    ``torch.cuda.is_available()`` is ``True`` and device type reports ``"cuda"``,
+    which would silently route into kernels compiled for NVIDIA — so detect HIP by
+    build tag and fail fast with a clear message instead. (Apple MLX never reaches
+    here: Unsloth's MLX build uses a separate trainer that has no Gefen-X path.)
+
+    A CPU-only box (no CUDA, no HIP/XPU) is allowed — ``gefen`` transparently
+    downgrades its fused path to the pure-PyTorch step there, which keeps the
+    optimizers usable for tests and small CPU experiments.
+    """
+    import torch
+
+    hip = getattr(getattr(torch, "version", None), "hip", None)
+    if hip:
+        raise RuntimeError(
+            "Unsloth: the Gefen-X optimizers require NVIDIA CUDA, but this is an "
+            f"AMD/ROCm (HIP {hip}) build of PyTorch. gefen's CUDA kernels do not "
+            "support ROCm."
+        )
+    if (
+        hasattr(torch, "xpu")
+        and torch.xpu.is_available()
+        and not (hasattr(torch, "cuda") and torch.cuda.is_available())
+    ):
+        raise RuntimeError(
+            "Unsloth: the Gefen-X optimizers require NVIDIA CUDA; Intel XPU is "
+            "unsupported by gefen's kernels."
+        )
+
+
+def coerce_optim_arg(value: Any) -> Any:
+    """Restore native types on string-form optimizer args.
+
+    String ``optim_args`` (``key=value``) arrive as strings; turn the obvious
+    ``true``/``false``/``none``/int/float literals back into real values. Mirrors
+    the axolotl Gefen-X integration so a config authored either way behaves the
+    same.
+    """
+    if not isinstance(value, str):
+        return value
+    lowered = value.strip().lower()
+    if lowered in ("true", "false"):
+        return lowered == "true"
+    if lowered in ("none", "null"):
+        return None
+    try:
+        return int(value)
+    except ValueError:
+        pass
+    try:
+        return float(value)
+    except ValueError:
+        return value
+
+
+# Fields on GefenXConfig that map 1:1 onto gefen.Gefen keyword arguments.
+_GEFENX_FIELDS: Tuple[str, ...] = (
+    "fused",
+    "factored_v_2d",
+    "force_1d_period_one",
+    "force_2d_period_one",
+    "period_one_substrings",
+    "codebook_refresh_every",
+    "stochastic_round",
+    "capturable",
+)
+
+# Fields on GefenXMuonConfig that map 1:1 onto gefen.GefenMuonHybrid kwargs.
+# `backup_lr_scale`, `backup_substrings`, `betas` and `eps` are handled specially.
+_GEFENX_MUON_FIELDS: Tuple[str, ...] = (
+    "fused",
+    "adjust_lr_fn",
+    "muon_lr",
+    "backup_lr",
+    "muon_weight_decay",
+    "backup_weight_decay",
+    "backup_1d_period_one",
+    "backup_2d_period_one",
+    "momentum",
+    "nesterov",
+    "ns_steps",
+    "ns_schedule",
+    "sharded_mode",
+    "fp8_ns",
+    "stochastic_round",
+    "normuon",
+    "cautious",
+    "capturable",
+)
+
+
+def _collect_kwargs(config, fields: Tuple[str, ...]) -> Dict[str, Any]:
+    """Pull the named attributes off `config`, plus its `extra_kwargs` escape hatch."""
+    kwargs: Dict[str, Any] = {}
+    for name in fields:
+        if hasattr(config, name):
+            value = getattr(config, name)
+            # An empty period_one_substrings means "use the gefen default"; drop it
+            # so we don't override with () when the user never set it.
+            if name == "period_one_substrings" and not value:
+                continue
+            kwargs[name] = value
+    extra = getattr(config, "extra_kwargs", None)
+    if extra:
+        kwargs.update({k: coerce_optim_arg(v) for k, v in extra.items()})
+    return kwargs
+
+
+def make_gefenx_param_groups(
+    model,
+    lr: float,
+    weight_decay: float,
+    embedding_lr: Optional[float] = None,
+) -> List[Dict[str, Any]]:
+    """Group trainable params for plain ``gefen.Gefen``.
+
+    Matches Unsloth's ``_create_unsloth_optimizer`` embedding split: params saved
+    via PEFT ``modules_to_save`` (embeddings / heads trained at full rank) get the
+    dedicated ``embedding_lr`` when one is provided; everything else shares ``lr``.
+    """
+    non_embeddings: List[Any] = []
+    embeddings: List[Any] = []
+
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if embedding_lr is not None and name.endswith("modules_to_save.default.weight"):
+            partial_name = name[: -len(".modules_to_save.default.weight")]
+            partial_name = partial_name[partial_name.rfind(".") + 1 :]
+            print(
+                f"Unsloth: Setting lr = {embedding_lr:.2e} instead of {lr:.2e} for {partial_name}."
+            )
+            embeddings.append(param)
+        else:
+            non_embeddings.append(param)
+
+    param_groups: List[Dict[str, Any]] = [
+        {"params": non_embeddings, "weight_decay": weight_decay, "lr": lr},
+    ]
+    if embeddings:
+        param_groups.append(
+            {"params": embeddings, "weight_decay": weight_decay, "lr": embedding_lr}
+        )
+    return param_groups
+
+
+def build_gefenx_optimizer(
+    model,
+    config,
+    *,
+    lr: float,
+    weight_decay: float,
+    betas: Tuple[float, float],
+    eps: float,
+    embedding_lr: Optional[float] = None,
+):
+    """Construct ``gefen.Gefen`` from a ``GefenXConfig`` and the trainer's hyperparams."""
+    _require_nvidia_cuda()
+    from gefen import Gefen
+
+    kwargs = _collect_kwargs(config, _GEFENX_FIELDS)
+    # Trainer's AdamW betas/eps are the fallback; a config override wins.
+    config_betas = getattr(config, "betas", None)
+    config_eps = getattr(config, "eps", None)
+    kwargs.setdefault("betas", tuple(config_betas) if config_betas is not None else betas)
+    kwargs.setdefault("eps", config_eps if config_eps is not None else eps)
+
+    param_groups = make_gefenx_param_groups(model, lr, weight_decay, embedding_lr)
+    return Gefen(param_groups, lr=lr, weight_decay=weight_decay, **kwargs)
+
+
+def build_gefenx_muon_optimizer(
+    model,
+    config,
+    *,
+    lr: float,
+    weight_decay: float,
+    betas: Tuple[float, float],
+    eps: float,
+):
+    """Construct ``gefen.GefenMuonHybrid`` from a ``GefenXMuonConfig``.
+
+    ``GefenMuonHybrid.from_model`` splits the model's params (2D hidden weights →
+    Muon, everything else → Gefen backup), validates the split covers every
+    trainable parameter exactly once, and constructs the hybrid. The Gefen-X
+    recommended recipe (``backup_1d_period_one`` / ``adjust_lr_fn`` /
+    ``fused`` / ``backup_lr = backup_lr_scale * lr``) is applied as defaults, all
+    overridable through the config.
+    """
+    _require_nvidia_cuda()
+    from gefen import GefenMuonHybrid
+
+    kwargs = _collect_kwargs(config, _GEFENX_MUON_FIELDS)
+
+    # Recommended recipe defaults (only fill in what the config left unset).
+    kwargs.setdefault("backup_1d_period_one", True)
+    kwargs.setdefault("adjust_lr_fn", "match_rms_adamw")
+    kwargs.setdefault("fused", True)
+
+    # backup_lr defaults to backup_lr_scale * lr (validated loss lever) unless the
+    # config set an explicit backup_lr.
+    if kwargs.get("backup_lr") is None:
+        scale = getattr(config, "backup_lr_scale", 0.5)
+        if scale is not None:
+            kwargs["backup_lr"] = scale * lr
+
+    config_betas = getattr(config, "betas", None)
+    config_eps = getattr(config, "eps", None)
+    if config_betas is not None:
+        kwargs["betas"] = tuple(config_betas)
+    else:
+        kwargs.setdefault("betas", betas)
+    if config_eps is not None:
+        kwargs["eps"] = config_eps
+    else:
+        kwargs.setdefault("eps", eps)
+
+    backup_substrings = getattr(config, "backup_substrings", None)
+
+    return GefenMuonHybrid.from_model(
+        model,
+        lr=lr,
+        weight_decay=weight_decay,
+        backup_substrings=backup_substrings,
+        **kwargs,
+    )

unsloth/trainer.py

@@ -50,6 +50,8 @@ __all__ = [
     "_patch_trl_trainer",
     "UnslothVisionDataCollator",
     "QGaloreConfig",
+    "GefenXConfig",
+    "GefenXMuonConfig",
     "check_dataset_for_missing_videos",
 ]
 
@@ -191,15 +193,85 @@ class QGaloreConfig:
     target_modules: Optional[List[str]] = None
 
 
+@dataclass
+class GefenXConfig:
+    """Configuration for the Gefen-X (``gefenx``) optimizer.
+
+    Wraps ``gefen.Gefen`` — an AdamW-family optimizer that keeps its optimizer
+    state at roughly 1 byte per parameter (8-bit quantized momentum + factored /
+    block-shared second moment). Pass an instance via
+    ``UnslothTrainingArguments(gefenx_config=...)`` to enable it. Fields map onto
+    ``gefen.Gefen`` keyword arguments; ``extra_kwargs`` forwards anything not
+    surfaced here. Requires the ``gefen-x`` package (``pip install gefen-x``).
+    """
+
+    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
+    # None → fall back to the trainer's adam_beta1/2 and adam_epsilon.
+    betas: Optional[tuple] = None
+    eps: Optional[float] = None
+    extra_kwargs: dict = field(default_factory=dict)
+
+
+@dataclass
+class GefenXMuonConfig:
+    """Configuration for the Gefen-X Muon hybrid (``gefenx_muon``) optimizer.
+
+    Wraps ``gefen.GefenMuonHybrid`` — Muon (Newton-Schulz orthogonalization) on
+    the 2D hidden weight matrices, plain Gefen on embeddings / heads / norms /
+    biases, at the same ≈1 byte/param footprint. Pass an instance via
+    ``UnslothTrainingArguments(gefenx_muon_config=...)``. Defaults encode the
+    Gefen-X recommended recipe (``backup_1d_period_one`` /
+    ``adjust_lr_fn="match_rms_adamw"`` / ``fused`` / ``backup_lr = 0.5 * lr``);
+    every field is overridable. Requires the ``gefen-x`` package.
+    """
+
+    fused: bool = True
+    adjust_lr_fn: str = "match_rms_adamw"
+    # backup_lr = backup_lr_scale * lr, unless backup_lr is set explicitly.
+    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
+    # None → fall back to the trainer's adam_beta1/2 and adam_epsilon.
+    betas: Optional[tuple] = None
+    eps: Optional[float] = None
+    extra_kwargs: dict = field(default_factory=dict)
+
+
 class UnslothTrainingArguments(TrainingArguments):
     def __init__(
         self,
         embedding_learning_rate: float = None,
         q_galore_config: Optional[QGaloreConfig] = None,
+        gefenx_config: Optional[GefenXConfig] = None,
+        gefenx_muon_config: Optional[GefenXMuonConfig] = None,
         *args,
         **kwargs,
     ):
         self.q_galore_config = q_galore_config
+        self.gefenx_config = gefenx_config
+        self.gefenx_muon_config = gefenx_muon_config
         self.embedding_learning_rate = embedding_learning_rate
         super().__init__(*args, **kwargs)
         self.embedding_learning_rate = embedding_learning_rate
@@ -256,6 +328,15 @@ class UnslothTrainer(SFTTrainer):
             embedding_lr = getattr(self.args, "embedding_learning_rate", None)
             return self._create_q_galore_optimizer(q_galore_config, embedding_lr)
 
+        # --- Gefen-X optimizers ---
+        gefenx_config = getattr(self.args, "gefenx_config", None)
+        if gefenx_config is not None and self.optimizer is None:
+            return self._create_gefenx_optimizer(gefenx_config)
+
+        gefenx_muon_config = getattr(self.args, "gefenx_muon_config", None)
+        if gefenx_muon_config is not None and self.optimizer is None:
+            return self._create_gefenx_muon_optimizer(gefenx_muon_config)
+
         # --- Embedding-LR optimizer ---
         embedding_learning_rate = getattr(self.args, "embedding_learning_rate", None)
         if embedding_learning_rate is None:
@@ -371,6 +452,52 @@ class UnslothTrainer(SFTTrainer):
 
         return self.optimizer
 
+    def _create_gefenx_optimizer(self, config: "GefenXConfig"):
+        """Build the Gefen-X optimizer (``gefen.Gefen``) from a GefenXConfig."""
+        from unsloth.optimizers.gefenx import build_gefenx_optimizer
+
+        embedding_lr = getattr(self.args, "embedding_learning_rate", None)
+        self.optimizer = build_gefenx_optimizer(
+            self.model,
+            config,
+            lr = self.args.learning_rate,
+            weight_decay = self.args.weight_decay,
+            betas = (self.args.adam_beta1, self.args.adam_beta2),
+            eps = self.args.adam_epsilon,
+            embedding_lr = embedding_lr,
+        )
+        n_params = sum(
+            len(g["params"]) for g in self.optimizer.param_groups
+        )
+        print(
+            f"🦥 Unsloth: Gefen-X optimizer enabled — "
+            f"{n_params} trainable params at ≈1 byte/param optimizer state."
+        )
+        return self.optimizer
+
+    def _create_gefenx_muon_optimizer(self, config: "GefenXMuonConfig"):
+        """Build the Gefen-X Muon hybrid (``gefen.GefenMuonHybrid``) from a GefenXMuonConfig."""
+        from unsloth.optimizers.gefenx import build_gefenx_muon_optimizer
+
+        if getattr(self.args, "embedding_learning_rate", None) is not None:
+            print(
+                "Unsloth: embedding_learning_rate is ignored by gefenx_muon — the "
+                "hybrid already routes embeddings to its Gefen backup half."
+            )
+        self.optimizer = build_gefenx_muon_optimizer(
+            self.model,
+            config,
+            lr = self.args.learning_rate,
+            weight_decay = self.args.weight_decay,
+            betas = (self.args.adam_beta1, self.args.adam_beta2),
+            eps = self.args.adam_epsilon,
+        )
+        print(
+            "🦥 Unsloth: Gefen-X Muon hybrid enabled — Muon on 2D hidden weights, "
+            "Gefen on embeddings / heads / norms / biases (≈1 byte/param)."
+        )
+        return self.optimizer
+
 
 # From `trl>=0.13.0`, they changed how to pass several params to the trainer
 # We need to patch to make the transition smooth