diff --git a/tests/python/test_gefenx_optimizer.py b/tests/python/test_gefenx_optimizer.py new file mode 100644 index 0000000000..376eadfe4d --- /dev/null +++ b/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) diff --git a/unsloth/optimizers/gefenx.py b/unsloth/optimizers/gefenx.py new file mode 100644 index 0000000000..ee7c8b9719 --- /dev/null +++ b/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, + ) diff --git a/unsloth/trainer.py b/unsloth/trainer.py index 83cb1758f0..aa8f2c6c99 100644 --- a/unsloth/trainer.py +++ b/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