A fork of upstream Gefen that fixes critical gaps in the shipped release, adds modern-architecture support, and closes the loss gap to AdamW.

Gefen is intended to be a drop-in replacement for the AdamW optimizer for memory-efficient training. It keeps the familiar AdamW training recipe while cutting optimizer state to ~1 byte per parameter (measured 1.01 B/param) — an 8× reduction versus classic fp32 AdamW state, about 6.5 GiB saved per billion parameters, while maintaining AdamW-level loss. The reduced memory footprint enables training larger models and larger batch sizes for added throughput. Gefen-X (fused) currently out-performs AdamW-8bit and AdamW-4bit in throughput, memory, and loss, while staying competitive with the bf16-state fused AdamW baseline's loss at roughly a quarter of the optimizer memory.

Fork Highlights:

• Fine-tuning and pre-training — upstream Gefen was validated for pre-training; this fork makes it a drop-in optimizer for fine-tuning: full fine-tune (FFT), LoRA, and QLoRA, with a VRAM sizing guide.
• New: matches AdamW's loss out of the box at about a quarter of its optimizer memory (default factored_v_2d). See Benchmarks and the factored-v lever.
• Works on modern decoders (Qwen3, Llama-3, Mistral). Upstream loses its memory advantage on these architectures (≈9 B/param — worse than AdamW); this fork keeps the intended ≈1 B/param.
• Validated across 2026 architectures — two dozen modern LLMs, VLMs, and image/video/audio media-gen models full-fine-tuned on 24 GB GPUs, plus 20-30B MoEs via LoRA. See the compatibility matrix.
• ≈2× faster opt.step() via fused CUDA kernels with validated numerical parity.
• Whole-model Muon option (GefenMuonHybrid) with a selectable AdamW quality backup or ~1 B/param Gefen low-memory backup, plus task-specific SFT and pretraining recipes.
• Reliable native checkpoint save/resume and FSDP2 training support — broken or absent in the shipped release; see Distributed Training for the checkpoint compatibility table.
• Hardened against crashes (device/dtype guards, bounds checks, race fixes).

Basic Usage

Installation

pip install gefen-x          # imports as `gefen`

[!IMPORTANT] It's gefen-x, not gefen. This fork publishes as gefen-x and imports as gefen. pip install gefen fetches the upstream package, which does not include the fixes/improvements above.

For development and latest source, install editable from source instead:

git clone https://github.com/thad0ctor/Gefen-X && cd Gefen-X && pip install -e .

The package is pure-Python at install time, the fused kernels (CUDA required) are built on first use via PyTorch JIT (torch.utils.cpp_extension) + nvcc. The first CUDA run takes a few minutes; later runs reuse the cached build keyed on your environment. Force a rebuild with GEFEN_FORCE_REBUILD=1.

Requirements

CUDA toolkit and host compiler compatible with your PyTorch build (the JIT compiles for your GPU's compute capability; set TORCH_CUDA_ARCH_LIST to target arch or a multi-arch build).

Supported versions

The Hugging Face Trainer flow uses transformers and Accelerate, while the data-backed examples also use datasets; these are not core dependencies — install them alongside (pip install transformers 'accelerate>=1.1' datasets). Optional benchmark baselines: bitsandbytes (AdamW-8bit), torchao (AdamW-4bit).

Compatibility

Validated with full-parameter fine-tune smoke tests on 24 GB GPUs across modern LLMs/VLMs and media-generation models — Gemma 4, Qwen3-VL, Mistral 3, MiniCPM-V 4.6, LFM2.5 (+VL), SDXL, Stable Diffusion 3.5, FLUX.1-dev (12B via FSDP2), FLUX.2-klein, Z-Image (6B via FSDP2), Anima, and ACE-Step 1.5 — and on 20-30B MoEs (GPT-OSS-20B, GLM-4.7-Flash) via LoRA adapter training. Per-model results, method, and hardware footprints: COMPATIBILITY.md.

Sizing a GPU? Hardware Requirements gives measured peak VRAM per model size and method (full fine-tune with Gefen vs AdamW, LoRA, QLoRA), plus a per-card rule of thumb for other sizes.

Want a real end-to-end run? See examples/toy-finetune/ — fine-tune a small model in about four minutes on one 24 GB GPU, with a before/after check that makes success obvious.

Learning rate (when porting an AdamW config)

Use ~0.6× your AdamW learning rate (e.g. AdamW at 5e-5 → Gefen at 3e-5). That is Gefen's measured optimum on the models tested (Qwen3 0.6B–8B), and at it Gefen matches AdamW's loss. Don't creep the LR back up toward the AdamW value — Gefen is less tolerant of a too-hot LR than AdamW is.

Better than the heuristic: measure it. The LR finder picks the best LR for your model empirically:

python -m gefen.tools.find_lr --model <path> --optimizer gefen --method sweep

See tools/README.md.

One knock-on effect: weight decay in AdamW-style optimizers is applied as lr × weight_decay, so scaling lr down to 0.6× also weakens regularization by 0.6×. If your recipe relies on weight decay, scale weight_decay up by the inverse (~1.7×) or retune it.

Distributed Training

Gefen drops into standard distributed training like any other PyTorch optimizer, with either fused=True or fused=False.

FSDP2 checkpoint scope. Plain Gefen and Muon approx save and resume exactly through PyTorch's standard full-state checkpoint calls, as long as the GPU count and sharding layout are unchanged and every GPU joins the save. Anything outside that scope refuses to load instead of corrupting state — full details.

Mixed precision works out of the box: BF16 and standard AMP behave exactly as with any PyTorch optimizer, and true-FP16 GradScaler training is handled safely — an overflow step changes nothing. FSDP1 FP16 needs torch.distributed.fsdp.ShardedGradScaler.

DeepSpeed ZeRO config. Set "zero_allow_untested_optimizer": true and leave the config's optimizer section unset. With optimizer CPU-offload, also set "zero_force_ds_cpu_optimizer": false — otherwise raw DeepSpeed refuses to initialize, and accelerate-based launchers (axolotl) silently swap in DeepSpeed's own CPU Adam.

Determinism (deterministic)

Set deterministic=True when every GPU replica must produce bit-identical results on matching GPUs. It is off by default (the fastest routing), checkpoints remember the setting, and older checkpoints without it still load. One combination is rejected: plain Gefen with deterministic=True, factored_v_2d=True, and stochastic_round=True together.

optimizer = Gefen(model.named_parameters(), lr=3e-4, fused=True, deterministic=True)

It is an ordinary constructor argument, so it works the same everywhere you build the optimizer: pass it in the Hugging Face Trainer construction, or in axolotl add optim_args: { deterministic: true }.

CUDA Graphs & torch.compile (capturable)

All three optimizers accept capturable=True (same meaning as torch.optim's argument): opt.step() can then be captured in a torch.cuda.CUDAGraph or wrapped in torch.compile(mode="reduce-overhead") — the compiled hybrid step is about 10% faster than eager. Checkpoints stay correct across graph replays. Usage, caveats, and measured numbers: COMPATIBILITY.md.

Benchmarks:

Loss curves — loss over the 2000 steps. For each optimizer (distinguished by color), there are two lines: solid = held-out eval loss (the 32-example validation set, logged every 50 steps — this is the comparison metric in the table above) and dashed = train-loss EMA (exponential moving average of the training loss).

Gefen-fused (with the factored_v_2d default) sits on the AdamW cluster in each retained cohort — best in the provenance-complete 0.6B seed (1.4251 vs 1.4360 for fused AdamW) and a statistical tie within the legacy 1.7B context (1.2257 vs 1.2169; each eval number moves ±0.005–0.007 between reruns, so gaps under ~0.01 are ties). Balanced-SFT Gefen-Muon + AdamW ties fused AdamW at 0.6B (1.4368 vs 1.4360). Its current 1.7B endpoint is 1.2434, but no strict delta against the legacy 1.7B AdamW row is asserted. These are single-seed outcomes, not universal optimizer rankings.

Throughput vs peak VRAM — the speed/memory frontier (upper-left = faster and lighter is better). Gefen-fused holds the lowest-VRAM column at 0.97× fused-AdamW throughput in the complete 0.6B cohort and 0.96× within the legacy 1.7B cohort. AdamW-4-bit has the highest peak VRAM in that legacy 1.7B context despite its small optimizer state (torchao transient buffers).

Why AdamW-4-bit's peak VRAM is high (and Gefen's isn't). torchao's 4-bit step upcasts each parameter's state to fp32 to compute the update ("optimizer step calculations are always done in FP32"), so the peak is set by a transient stack of fp32 buffers sized to the largest tensor — here the tied embedding / LM head — not by the tiny 1.06 B/param packed state. With bf16 master weights the fused-AdamW baseline is already lean, so that fp32 spike pushes 4-bit's peak above even the bf16-state fused AdamW baseline (15.11 vs 14.00 GiB at 1.7B). Gefen keeps the same ~1 B/param state without the fp32 recompute and holds the lowest peak at both measured scales.

Balanced-SFT Gefen-Muon uses less peak VRAM and optimizer state than fused AdamW, but more than plain Gefen because its non-Muon subset keeps conventional AdamW state. Newton–Schulz orthogonalization also makes it the slowest optimizer in this end-to-end table.

Gefen-fused is the recommended SFT configuration overall: AdamW-level loss with the lowest peak VRAM and the lowest optimizer state in each retained cohort, at 0.97× AdamW's speed in the complete 0.6B comparison and 0.96× within the legacy 1.7B context. (AdamW-4-bit's optimizer state is also small, but its peak VRAM is the highest in the legacy 1.7B results because of transient buffers — Gefen is the real memory winner.)

Gefen-Muon + AdamW is the balanced recipe when you specifically want Muon-style updates: it ties fused AdamW's loss at 0.6B while using 55% less optimizer state, but it is 33% slower. The complete 1.7B Muon row records 1.2434 eval loss, 1373 tok/s, and 1.550 state B/param; the non-Muon 1.7B rows are legacy context, so cross-cohort speed and loss percentages are intentionally not claimed. For minimum-state Muon, use the Gefen backup recipe below; otherwise plain Gefen-fused is the stronger SFT default.

Review the raw runs: per-cell training logs (step-by-step loss, throughput, VRAM) in docs/benchmarks/logs/ · aggregated metrics as CSV and JSONL.

Gefen-X Integrations:

Using Gefen with Axolotl

Axolotl trains with Gefen through two native optimizer names, added by axolotl-ai-cloud/axolotl#3808 (install axolotl from that PR until it lands in a release):

• optimizer: gefenx — the core quantized AdamW drop-in (Gefen).
• optimizer: gefenx_muon — the whole-model Muon hybrid (GefenMuonHybrid): Muon on 2D hidden weights plus a selectable Gefen or AdamW backup for embeddings / LM head / norms / biases.

Install gefen-x (pip install gefen-x), not the PyPI gefen — which is upstream Gefen, without this fork's fixes (the modern-arch period==1 memory fallback, the ~2× fused kernels, the Muon hybrid). See Installation.

gefenx — memory-efficient AdamW drop-in (recommended):

optimizer: gefenx
learning_rate: 6.0e-6    # ≈0.6× your AdamW LR — Gefen takes larger effective steps
bf16: true
optim_args:
  fused: true            # ≈2× faster opt.step; validated numerical parity (default)
  factored_v_2d: true    # Adafactor-style 2D 2nd moment — matches AdamW loss (default)

gefenx_muon — whole-model Muon hybrid (balanced SFT):

optimizer: gefenx_muon
learning_rate: 3.0e-5    # tune for your model/data; this is the retained SFT LR
bf16: true
optim_args:
  fused: true
  backup_optimizer: adamw
  backup_lr: 3.0e-5      # full LR for embed / head / norms in the balanced recipe
  ns_schedule: tuned3
  normuon: true
  sharded_mode: exact    # "distributed" splits Newton-Schulz across GPUs under FSDP2

The Axolotl factory currently defaults to the low-memory Gefen recipe. To select it explicitly, change the balanced block above to:

optim_args:
  fused: true
  backup_optimizer: gefen
  backup_lr: 1.5e-5      # 0.5 × learning_rate
  backup_1d_period_one: true
  ns_schedule: tuned3
  normuon: true

For quality-first pretraining, keep the AdamW backup at full LR and set ns_schedule: standard plus normuon: false. In every recipe, adjust_lr_fn="match_rms_adamw" keeps Muon updates on the AdamW LR scale.

How config maps: learning_rate, weight_decay, adam_beta1/adam_beta2, and adam_epsilon forward to the selected constructor. Under optim_args, fused and deterministic apply to both optimizer names; factored_v_2d is gefenx-only; and backup_optimizer, backup_lr, ns_schedule, normuon, and sharded_mode are gefenx_muon-only. String values are coerced to type. For plain Gefen, measure the LR instead of relying only on the ~0.6× heuristic: python -m gefen.tools.find_lr --model <path> --optimizer gefen --method sweep — see tools/README.md.

Hugging Face Trainer

Until native optim="gefen" support is released in Transformers, pass Gefen to the Trainer with optimizer_cls_and_kwargs:

from gefen import Gefen
from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="outputs",
    learning_rate=3e-5,   # ≈0.6× your AdamW LR — see "Learning rate" above; don't copy an AdamW-scale LR here
    weight_decay=0.0,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    optimizer_cls_and_kwargs=(
        Gefen,
        {
            "lr": training_args.learning_rate,
            "betas": (training_args.adam_beta1, training_args.adam_beta2),
            "eps": training_args.adam_epsilon,
            "fused": True,
        },
    ),
)

PyTorch

import torch
from gefen import Gefen

device = "cuda" if torch.cuda.is_available() else "cpu"
model = torch.nn.Linear(128, 10).to(device)

# optimizer = torch.optim.AdamW(
optimizer = Gefen(  # Replace AdamW with Gefen:
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.0,
)

inputs = torch.randn(32, 128, device=device)
targets = torch.randint(0, 10, (32,), device=device)

logits = model(inputs)
loss = torch.nn.functional.cross_entropy(logits, targets)
loss.backward()

optimizer.step()
optimizer.zero_grad(set_to_none=True)

print('Finished successfully.')

Extension: Gefen-Muon

GefenMuon adds a Muon-style pseudo-orthogonalization step (Keller Jordan et al., "Muon", 2024) on the first moment (skipping the second moment), then quantizes those first moments to 8-bit with Gefen's partitioning quantization.

Scope — read this first. Like Muon, GefenMuon optimizes 2D hidden weight matrices only. It has no code path for embeddings, the LM head, or 1D parameters (norms, biases) — feeding it those will fail or misbehave. For a full model, use GefenMuonHybrid below, which routes the non-Muon parameters to your selected Gefen or AdamW backup.

Pass (name, param) pairs, not bare tensors. GefenMuon keys its 8-bit codebook cache on each parameter's name. If you strip the names (e.g. [p for _, p in pairs]) every parameter collapses to the name "none" and the cache is corrupted. Always pass the named pairs through:

from gefen import GefenMuon

# muon_named_params: list of (name, param) for your 2D hidden weight matrices.
optimizer = GefenMuon(muon_named_params, lr=lr)   # note: raw GefenMuon defaults weight_decay=0.1 (the standard Muon recipe), unlike Gefen / the hybrid, which default 0.0

Full models: GefenMuonHybrid

GefenMuonHybrid is the drop-in for training a whole model with Muon. It routes 2D hidden weight matrices to GefenMuon and everything else (embeddings, LM head, norms, biases) to a selectable backup optimizer, behind a single torch.optim.Optimizer interface:

• backup_optimizer="gefen" is the backward-compatible default and keeps both halves quantized for the smallest optimizer state.
• backup_optimizer="adamw" uses torch.optim.AdamW for the backup parameters; it costs normal AdamW state on that subset but is the measured quality choice for both SFT and pretraining.

Hand it a model and it splits and routes the parameters for you — no manual param lists:

from gefen import GefenMuonHybrid

optimizer = GefenMuonHybrid(model, lr=ADAMW_LR)               # single-argument form
optimizer = GefenMuonHybrid.from_model(model, lr=ADAMW_LR)    # equivalent explicit classmethod

Advanced Usage:

[!NOTE] Advanced tuning levers are available for advanced users. These features are largely experimental in nature and deviate from default settings - often trading tokens/sec throughput for accuracy.

Gefen Muon

Which Muon recipe should I use?

There is no single winner across tasks. The retained RTX 3090 matrix supports three distinct choices. Treat the learning rates below as roles, not universal constants: the SFT run selected 3e-5, while the small pretraining protocol selected 3e-3; tune MUON_LR for your model and data.

SFT: balanced Muon recipe

Use tuned three-step Newton–Schulz plus NorMuon and a full-LR AdamW backup:

optimizer = GefenMuonHybrid(
    model,
    lr=MUON_LR,
    backup_optimizer="adamw",
    backup_lr=MUON_LR,
    ns_schedule="tuned3",
    normuon=True,
    adjust_lr_fn="match_rms_adamw",
    weight_decay=WEIGHT_DECAY,
    fused=True,
)

The retained schedule ablation selected this combination over classic NS5 for SFT. In the curated end-to-end Qwen3-0.6B comparison it tied fused AdamW's loss with 55% less optimizer state, but was 33% slower. The provenance-complete 1.7B Muon run reached 1.2434 eval loss; its non-Muon baselines are legacy context, so no strict cross-cohort loss gap is claimed. This is the recommended way to run Muon for SFT, while plain Gefen-fused remains the stronger overall SFT default on the measured models.

Pretraining: quality first

Keep classic five-step Newton–Schulz, disable NorMuon, and use AdamW on the non-Muon parameters:

optimizer = GefenMuonHybrid(
    model,
    lr=MUON_LR,
    backup_optimizer="adamw",
    backup_lr=MUON_LR,
    ns_schedule="standard",
    ns_steps=5,
    normuon=False,
    adjust_lr_fn="match_rms_adamw",
    weight_decay=WEIGHT_DECAY,
    fused=True,
)

[!NOTE] Classic NS5 won the replicated 33M and 134M pretraining comparisons, so the faster SFT schedule is not promoted as a universal default. The evidence is small-scale and replicated, not a scaling-law claim.

Lowest optimizer memory

Keep the default Gefen backup when capacity matters more than the last quality increment:

optimizer = GefenMuonHybrid(
    model,
    lr=MUON_LR,
    backup_optimizer="gefen",      # constructor default; explicit for clarity
    backup_lr=0.5 * MUON_LR,
    backup_1d_period_one=True,
    ns_schedule="tuned3",
    normuon=True,
    adjust_lr_fn="match_rms_adamw",
    weight_decay=WEIGHT_DECAY,
    fused=True,
)

This keeps both halves on quantized Gefen state and is the low-memory recommendation, not the quality winner. Exact matrix-run figures and protocol are retained in PR #62 rather than a generated report in the repository.

backup_optimizer and fused are independent. The former selects Gefen or AdamW for non-Muon parameters. fused=True selects each chosen optimizer's fused CUDA implementation; it does not mean “use Gefen as the backup.”

Manual control. For a custom split, gefen.split_params_for_muon(model) returns the (muon_named, backup_named) (name, param) lists; pass them as the two positional arguments (GefenMuonHybrid(muon_named, backup_named, lr=...)) and adjust as needed. The single-argument and from_model forms cover the common case.

Muon learning rate — keep adjust_lr_fn="match_rms_adamw" (the default). Muon-native update scaling is not on the same footing as either backup. The default rescale makes an AdamW-scale lr meaningful across the whole hybrid, while muon_lr and backup_lr remain available for deliberate splits:

• "match_rms_adamw" (the default) rescales Muon updates to AdamW-equivalent size, so one AdamW-scale lr works for the whole model.
• None keeps Muon's native scaling — an AdamW-sized lr then under-trains the Muon matrices, and no single lr suits both halves. Don't combine None with an AdamW-scale lr.

Optional knobs (GefenMuonHybrid):

It supports step(), zero_grad(), state_dict()/load_state_dict(), and LR schedulers (e.g. torch.optim.lr_scheduler.StepLR(optimizer, ...)) like any optimizer. Because it splits params at construction rather than taking a single iterable, build it yourself and hand it to the Hugging Face Trainer via optimizers= (not optimizer_cls_and_kwargs):

optimizer = GefenMuonHybrid(
    model,
    lr=training_args.learning_rate,
    backup_optimizer="adamw",
    backup_lr=training_args.learning_rate,
    ns_schedule="tuned3",
    normuon=True,
)
trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset,
                  optimizers=(optimizer, None))  # (optimizer, lr_scheduler)

To prove save/resume in a real Trainer run, use benchmarks/trainer_resume/: it compares an uninterrupted run against a save-and-resume run and requires them to match exactly. The run command and covered scope are in COMPATIBILITY.md.

Quality Lever: factored second moment on 2D params (factored_v_2d)

[!NOTE] On by default — this is what makes plain Gefen match AdamW's loss. Set factored_v_2d=False for the legacy behavior: ~4% faster end-to-end at 1.7B (on par with fused AdamW's throughput), about 0.06 worse eval loss, and a lower best learning rate (2e-5 instead of 3e-5 at the benchmarked scales).

Gefen historically shared one adaptive step size across each block of parameters, where AdamW keeps one per parameter — that granularity difference was the ~0.06 loss gap. factored_v_2d gives every parameter of a weight matrix its own step size by tracking just one statistic per row and one per column (the Adafactor idea) and combining them on the fly inside the update kernel. The stored state is only rows + columns per matrix, so the ~1 B/param memory story is unchanged.

opt = Gefen(model.named_parameters(), lr=0.6 * ADAMW_LR, fused=True)  # factored-v is the default

Chart from the original factored-v ablation run. The table combines the current retained optimizer rows with that ablation; current per-cell logs and the retained Qwen3-1.7B legacy-mode log live in docs/benchmarks/logs/.

Experimental Lever: faster Newton-Schulz (ns_schedule, fp8_ns)

[!NOTE] GefenMuonHybrid defaults to ns_schedule="tuned3" for throughput. Keep it for the balanced SFT recipe; use ns_schedule="standard" for quality-first pretraining. fp8_ns is off by default.

The orthogonalization step is ~90% of a Muon step, so it has two speed levers:

• ns_schedule — fewer Newton–Schulz iterations for more throughput. "tuned3" (the hybrid default) is ~1.6× faster in the optimizer kernel. It carries a measurable pretraining-quality trade, while adding NorMuon recovers the classic endpoint/tail in the retained SFT run.
• fp8_ns — run the matrix math in fp8 on GPUs that support it (RTX 40-series/Hopper and newer) for matrices large enough to benefit; everywhere else it silently falls back to normal precision, so it's always safe to enable.

opt = GefenMuonHybrid(
    muon_named_params, backup_named_params,
    lr=5e-5, adjust_lr_fn="match_rms_adamw", fused=True,
    ns_schedule="tuned4",   # "tuned3" (hybrid default, ~1.6x) | "tuned4" (equal-quality, ~1.2x) | "standard" (classic quintic)
    fp8_ns=True,            # sm_89+ large matrices only; a no-op bf16 fallback elsewhere
)

Measured: fp8_ns pays off on larger models (~1.2× end-to-end at Qwen3-4B) and its built-in size gate avoids the penalty on small ones. tuned3 improves end-to-end throughput, but its quality is task-dependent: pair it with NorMuon for SFT, and retain classic NS5 for quality-first pretraining. The exact retained schedule comparison is documented in PR #62.

Quality Lever: per-neuron 2nd moment on the Newton-Schulz output (normuon)

[!NOTE] On by default in GefenMuonHybrid; free on speed and memory. Disable with normuon=False.

Muon's orthogonalization step leaves some output neurons taking slightly bigger steps than others. normuon (after NorMuon) evens that out by tracking a running average per neuron and normalizing against it — at essentially zero cost in speed or memory.

Measured in SFT: a small, consistent recovery for tuned3, culminating in an endpoint/tail tie with classic NS5 on the retained 2,000-update frontier. The quality-first pretraining recipe instead uses classic NS5 with NorMuon off. A related experiment (cautious=True, sign-agreement masking) clearly hurt and ships off.

Experimental Lever: per-element Gefen-backup state on embed / LM-head (backup_2d_period_one)

[!NOTE] Opt-in, default off. Unlike the levers above, this one costs memory — it's a deliberate loss-for-memory trade.

With a Gefen backup, this knob gives embeddings and the LM head full per-element statistics, like AdamW keeps everywhere. It is a deliberate loss/memory experiment and has no effect in AdamW-backup mode because AdamW is already per-element.

opt = GefenMuonHybrid(
    model,
    backup_optimizer="gefen",
    lr=5e-5, backup_lr=2.5e-5, adjust_lr_fn="match_rms_adamw",
    backup_1d_period_one=True,
    backup_2d_period_one=True,   # per-element v on embed/LM-head; costs memory
)

A Qwen3-1.7B sweep closed a small AdamW gap at 2.45 B/param. In the newer controlled matrix, the 2D variants made only small loss changes while clearly increasing state, peak VRAM, and step cost. Leave it off for the low-memory recipe; select the AdamW backup directly when quality is the objective.

A related opt-in, stochastic_round=True, switches the 8-bit momentum to unbiased rounding. It's free and validated but measured loss-neutral, so it stays optional.

Experimental Lever: sharded Newton-Schulz under FSDP2 (sharded_mode)

[!NOTE] Opt-in. Default is sharded_mode="exact" (single-GPU parity).

Under FSDP2, Muon's orthogonalization needs each full weight matrix, but every GPU only holds a slice. Three strategies:

• "exact" (default) — every GPU rebuilds every matrix and does the same work. It matches the single-GPU oracle within the tested BF16 tolerance, but the cross-rank GEMM path is not promised bitwise identical; the work is redundant.
• "distributed" (experimental) — each matrix is assigned by its stable position in the full distributed parameter set to one GPU, which does the work and shares the result. Bit-identical to "exact" on homogeneous GPUs and faster as you add them. Momentum ownership is stable when the active gradient set varies, and state_dict() collectively gathers owner-local momentum so rank 0 can write a complete checkpoint after all ranks call it.
• "approx" — each GPU works on just its slice. Fastest, but results genuinely differ — an accuracy trade.

from gefen import GefenMuonHybrid

opt = GefenMuonHybrid(
    muon_named_params, backup_named_params,
    lr=5e-5, adjust_lr_fn="match_rms_adamw", fused=True,
    sharded_mode="distributed",  # "exact" (default, parity) | "distributed" (experimental, collective checkpoint) | "approx" (fastest, non-parity)
)
# only takes effect under FSDP2 (DTensor params); no-op single-GPU

"distributed" checkpointing is collective. state_dict() gathers each owner's momentum across ranks, so every rank must call it (as in a standard FSDP full-state-dict flow). Calling state_dict() on rank 0 only — e.g. a rank-0-only save loop — deadlocks. Save and load also transiently materialize the full unsharded momentum on every rank, so peak memory at checkpoint time approaches "exact" mode's. The saved checkpoint is complete on every rank and resumes under any world size, including a single-process optimizer; incomplete or inconsistent owner state fails closed before any state is touched.

Measured (Qwen3-0.6B, 2 and 4 GPUs): "distributed" matched "exact" in the retained convergence runs while improving throughput by 1.06× / 1.12×; "approx" is faster still (1.25× / 1.39×) but visibly costs training quality, and the cost grows with GPU count. The "distributed" win grows with model size.

Current Gaps and Roadmap:

Param Groups and Integrations

Gefen, GefenMuon, and GefenMuonHybrid preserve the parameter groups you pass to the optimizer, so list-indexed layer-wise LR recipes, per-group LR logging, and state_dict()["param_groups"] see the same group boundaries as conventional torch.optim optimizers. Per-parameter names are stored in optimizer state and mirrored as each group's param_names list for integrations that need name-level routing. Checkpoints from the older one-group-per-parameter layout are migrated on load when the total parameter order still matches and the old per-param hyperparameters can be represented by the new group layout.

Known limitations

• Hybrid checkpoint schema. GefenMuonHybrid's state_dict() uses its own nested {"muon": ..., "backup": ..., "backup_optimizer": "gefen" | "adamw"} layout. Resume from a checkpoint the hybrid itself saved—not one consolidated or converted to the flat torch {state, param_groups} layout. Cross-backend loads are rejected before either child is mutated; legacy untagged hybrid checkpoints are interpreted as Gefen-backed.
• FSDP2 optimizer checkpoints don't reshard. Plain Gefen and Muon approx resume only on the same GPU count and layout; changing either refuses to load. Model weights are unaffected — details.
• True-FP16 overflow skips are invisible to Accelerate's step_was_skipped flag. BF16 and standard AMP are unaffected and are the recommended modes in Trainer/Accelerate.

Troubleshooting

Roadmap

Credit (Upstream Gefen):

Citation (Gefen Upstream):

If you found this library useful, please consider citing our work:

@article{benedek2026gefen,
  title={Gefen: Optimized Stochastic Optimizer},
  author={Benedek, Nadav and Koren, Tomer and Fried, Ohad},
  journal={arXiv preprint arXiv:2606.13894},
  year={2026}
}

Paper (Gefen Upstream):

Gefen: Optimized Stochastic Optimizer