HeavyBall — optimizer library and diagnostic benchmark for silent optimizer failures

TL;DR

HeavyBall is a PyTorch optimizer library by Lucas Nestler (@_clashluke) where every optimizer is assembled from composable, torch.compile(fullgraph=True)-fused building blocks — AdamW, Muon, SOAP, PSGD (Kronecker / low-rank), LATHER, ADOPT, Schedule-Free, LaProp, plus orthogonal/hybrid variants — with chainable flags (MARS variance reduction, cautious updates, PaLM-style β₂, ECC state compression) that compose across optimizers. The headline artifact is the HeavyBall Benchmark / LightBench suite: ~150 independent pass/fail diagnostic tests designed to expose silent optimizer failures, where an optimizer converges without error but settles in a clearly suboptimal basin. Results show no single optimizer dominates — PSGDKron and AdamW+Cautious sit at the top (77.0% / 75.7% pass rates), Muon alone clears only 51.0%.

Key claims

A “silent failure” is when an optimizer converges without error but lands in a markedly worse basin than alternatives at the same task; loss curves alone cannot distinguish this from success, so silent failures only surface via expensive downstream evaluation [docs/benchmark.md “The Problem of Silent Failures”].
The benchmark replaces leaderboard-style scoring with a granular pass/fail diagnostic map: ~150 tests, each targeting a specific known optimization challenge (e.g. saddle-point escape), no partial credit, producing a per-(optimizer × task) capability matrix [docs/benchmark.md “The HeavyBall Diagnostic Benchmark”, benchmark_matrix figure].
Per-task budget is 1,000 hyperparameter trials × up to 1,000,000 steps; the reported Attempts column is the average trials-to-success and is interpreted as a tunability measure (lower = easier to tune) [docs/benchmark.md “Experimental Setup”].
Headline results: PSGDKron and PSGDKron (Newton) tie at 77.0% pass; AdamW + Cautious 75.7%; SOAP 72.5%; plain AdamW 72.3%; MuonLaProp 68.2%; RMSprop 55.6%; plain Muon 51.0% [docs/benchmark.md Results table].
The Cautious flag — masking update components that disagree with the current gradient direction — raises AdamW from 72.3% → 75.7% pass and lowers trials-to-success from 107.8 → 61.2, i.e. it both reduces silent failures and improves tunability [docs/benchmark.md “Case Study: The AdamW Family”].
Saddle-point navigation is included as an explicit diagnostic test; first-order gradient-only methods reliably stall in such regions, providing a sharp pass/fail signal for momentum / preconditioning capability [docs/benchmark.md “Case Study: Escaping the Saddle Point”].
The library’s composability is the engineering claim: LaProp(..., ecc="bf16+8", mars=True, caution=True, palm=True) is valid; flags are implemented as FunctionTransform chains that fuse into the same compiled graph, so MARS, ECC, cautious masking, weight decay, and stochastic rounding all collapse into a single Triton kernel for fused first-order optimizers — AdamW step time reportedly drops from 14 reads + 9 writes across O(N) kernels to 4 reads + 3 writes in one kernel (≈3× speedup vs. uncompiled), and from 10.63 ms in HeavyBall 2 to 4.15 ms in HeavyBall 3 [README “How it compiles”, “Benchmarks”].
ECC state compression (bf16 base + int8 correction, ≈3 bytes vs fp32’s 4) is implemented as a composable flag attached as call-time attributes rather than per-optimizer code; cites FlashOptim (arXiv 2602.23349) as the source of the bf16+8 scheme [README “ECC”].
For second-order optimizers under FSDP, HeavyBall auto-detects sharded parameters on first step and repartitions via all_to_all_single: each weight matrix is deterministically owned by one rank, the owner reconstructs the full parameter, computes the update once, and returns updated shards — saving redundant compute/memory at the cost of communication versus DDP-style replicated updates [README “Distributed Training”].

Method

HeavyBall implements optimizers as FunctionTransform chains over ~100 building-block functions in utils.py, each wrapped with torch.compile(fullgraph=True) so calls inline into a single graph and fuse into Triton kernels. Features that elsewhere ship as separate libraries — MARS gradient variance reduction, cautious-direction update masking, ECC compressed state (bf16 + int8 correction terms), PaLM β₂ scheduling, gradient/update clipping, warmup — are exposed as kwargs that compose freely across optimizers, because BaseOpt handles them once and dispatches to the optimizer-specific transform. A Parallel transform supports grafted optimizers (e.g. grafted-Adam = Adam-direction × SGD-magnitude). Second-order methods (Muon Newton-Schulz, SOAP/PSGD Kronecker factor updates) compile their preconditioning steps as separate compiled regions to avoid bad fused codepaths, while their elementwise portions still fuse with the rest of the graph.

The HeavyBall Benchmark (LightBench, hosted at github.com/HomebrewML/LightBench) frames silent failures as the central evaluation problem and replaces continuous loss curves with binary pass/fail outcomes per (optimizer, task) cell. With 1,000-trial / 1M-step budgets, the benchmark is explicitly testing raw capability rather than default-config behavior — an optimizer that requires careful tuning to pass a task still counts as a pass.

Results

Optimizer	Cautious	MARS	Success	Attempts	Avg runtime (s)
PSGDKron	–	–	77.0%	73.2	8240
PSGDKron (Newton)	–	–	77.0%	80.5	9052
AdamW	✓	–	75.7%	61.2	8072
SOAP	–	–	72.5%	77.9	7827
AdamW	–	–	72.3%	107.8	10029
MuonLaProp	–	–	68.2%	82.7	10141
RMSprop	–	–	55.6%	114.4	10725
Muon	–	–	51.0%	129.1	14525

Two takeaways the docs surface explicitly: (i) PSGD with full Kronecker preconditioning tops the chart, and SOAP places ahead of plain Muon by a substantial margin, suggesting the “Muon is sufficient” framing common in modded-NanoGPT speedrun discourse may not generalize to harder optimization landscapes; (ii) the Cautious flag is a Pareto improvement on AdamW — higher pass rate, lower trials-to-success, lower runtime — and the docs argue this is a diagnostic signal that plain AdamW silently fails on tasks Cautious AdamW solves.

Why it’s interesting

This is the first wiki-filed artifact that takes silent optimizer failure as a first-class evaluation problem rather than reporting a single best-LR loss curve, and the diagnostic-matrix framing directly complicates the “just use Muon” narrative implicit in the speedrun line (Muon is Not That Special: Random or Inverted Spectra Work Just as Well, The Newton-Muon Optimizer, MuonH (hyperball) sets new modded-NanoGPT speedrun record at 3325 steps — Keller Jordan announcing @wen_kaiyue's result): plain Muon clears only 51.0% of tasks here, well below AdamW+Cautious and PSGD. The complementary story is engineering — HeavyBall’s compose-then-compile architecture is the most aggressive version of the “treat optimizer features as composable transforms” pattern, and its ECC implementation directly imports the bf16+8 quantization scheme from FlashOptim (FlashOptim: drop-in memory-efficient optimizer replacements (Adam/SGD/Lion) via weight-splitting + ECC + companded 8-bit state), generalizing the same idea across any optimizer in the library rather than just Adam/SGD/Lion. Connects to the architecture-vs-optimizer fork mapped in Training stability at scale — if architecture (QK-Norm, GatedNorm, mHC) is doing most of the stability work in modern pretraining, then this benchmark’s relevance to from-scratch LLM runs is unclear; but if the optimizer-side fix (Controlled LLM Training on Spectral Sphere, SSO) is the right framing, then diagnostic capability matrices like this one are exactly the missing evaluation layer.