SpectOptim

Replace AdamW in ~20 lines. Get Muon / NorMuon / paper Aurora / EigenSOAP / DiagKron with auto param grouping, Newton–Schulz engine, neuron-death telemetry, HF Trainer hooks, and FSDP helpers.

Part of the Spectral Training Stack with OptiScale and OrthoLab.

Install

pip install -e ".[dev]"

Quick start

import torch.nn as nn
from spectoptim import SpectOptim, diagnose, TelemetryLogger

model = MyTransformer()
opt = SpectOptim(model, variant="normuon", lr_spectral=0.02, lr_aux=3e-4)

logger = TelemetryLogger("telemetry.csv", every=10)
for step, batch in enumerate(loader):
    opt.zero_grad(set_to_none=True)
    loss = model(batch)
    loss.backward()
    opt.step()
    logger.log_step(step, model, opt, extras={"loss": float(loss)})

print(diagnose(model, opt))

Presets:

from spectoptim import from_pretrained_defaults, swap_variant

opt = from_pretrained_defaults(model, "nanochat")   # NorMuon defaults
swap_variant(opt, "aurora")  # tall-MLP row-uniform fix, same groups

Variants

Variant	Role
muon	Polar / Newton–Schulz + aspect-ratio scale √max(1, rows/cols)
normuon	Row second-moment after orthogonalization + paper-ish RMS scale 0.2√(mn)/‖Ô‖_F
aurora	Paper leverage-uniform polar (Tilde Research aurora_polar)
aurora_heuristic	Legacy post-hoc row rebalance
muown	Experimental row-gain sketch
soap / eigen_soap	Eigenbasis Kronecker SOAP-style preconditioner
diag_kron	Cheap diagonal Kronecker baseline

Scale note: Muon uses aspect-ratio scaling; NorMuon uses a separate RMS formula. Do not treat their LRs as interchangeable without retuning.

from spectoptim.hf_trainer import create_optimizer
from spectoptim import canzona_overlap_spectral_updates, newton_schulz_compiled

Embeddings, norms, biases, and heads stay on AdamW automatically. FSDP: spectral_update_with_gather / canzona_overlap_spectral_updates for gather↔NS overlap (NS+aspect kernel only).

Features (v0.1)

• Auto grouping + GroupRules (LoRA / embed / MoE-aware skips)
• NS engine: Muon / classic / Chebyshev coeffs, Gram-NS, adaptive early stop
• Telemetry: update norms, row-norm std (death signal), soft/hard rank (sampled), condition
• diagnose() one-shot health report with actionable tips
• FSDP helper: gather → NS + aspect → scatter micro-API (spectoptim.fsdp_utils)
• Tuning: LR range test, lr×wd plateau map, muP transfer checklist
• Microbench: python bench/microbench.py

Docs

• Migration from AdamW
• FSDP notes
• Mixed precision pitfalls

Microbench

python bench/microbench.py --steps 100

Citations

Muon (Jordan et al.), NorMuon, Aurora, Muown, SOAP/Shampoo literature — see docs/migration.md.

License

MIT