HeavyBall PyTorch Optimizer Library

The HeavyBall Benchmark is a diagnostic tool designed to expose hidden optimizer weaknesses, specifically targeting silent failures. A silent failure occurs when an optimizer converges without error but settles in a suboptimal solution (a suboptimal basin), making the training process appear successful when it is actually underperforming.

Unlike traditional leaderboards that provide a single score, the HeavyBall Benchmark uses a suite of over 150 independent, pass/fail tests. Each test targets a specific optimization challenge (e.g., escaping saddle points). This provides a granular map of an optimizer's strengths and weaknesses rather than a simple comparative ranking.

HeavyBall uses 'pinned upstream oracles' as accuracy and speed baselines for its benchmarking matrix (benchmarks/matrix.py). These are shallow clones of external open-source implementations.

Important for users:

• These directories are dev-time dependencies only.
• They are not imported by the heavyball package.
• They are not included in any official release.
• If you are using the heavyball library in a production or standard environment, you do not need to install or interact with these upstream repositories.

ECC (Error-Compensated Compression) reduces optimizer state memory usage by storing tensors as a bf16 value plus an int8 correction term (3 bytes total vs 4 bytes for fp32). This is particularly effective for first-order optimizers, providing roughly 25% memory savings.

Available modes:

• bf16+8
• bf16+16
• fp16+8
• fp16+16

You can apply this to the optimizer state using ecc or to both state and parameters using param_ecc.

opt = AdamW(model.parameters(), lr=1e-3, ecc="bf16+8")
opt = Muon(model.parameters(), lr=0.02, ecc="bf16+8", param_ecc="bf16+8")

When using PSGDKron, you can switch to CachedPSGDKron to reduce the per-step computational overhead.

How it works:

• PSGDKron computes the preconditioning matrix on the fly based on the triangular Q at every step.
• CachedPSGDKron precomputes the preconditioner and reuses it across steps.

Trade-offs:

• Compute: Significantly reduces per-step overhead.
• Memory: Doubles the memory overhead compared to standard PSGDKron.

Hybrid Configuration: If the memory cost of CachedPSGDKron is too high, you can combine it with triu_as_line=True. This reduces the total memory cost from 2x Q to 1.5x Q.

The benchmark evaluates optimizers using a binary pass/fail mechanism for each task. To ensure raw capability is tested rather than just default settings, the experimental setup uses a generous budget for each optimizer:

• Hyperparameter Trials: 1,000 trials per task.
• Step Limit: Each trial runs for up to 1,000,000 steps.

Key Metrics in Results:

• Success Rate: The percentage of trials that successfully solved the task.
• Attempts: The average number of trials required to achieve success. A lower number indicates the optimizer is easier to tune for that specific problem.

HeavyBall supports DDP and FSDP. For second-order methods (Muon, SOAP, PSGD) using FSDP, HeavyBall automatically detects sharded parameters and performs a metadata-first all_to_all_single exchange to repartition parameters so the update can be computed on a single rank.

Note: When using FSDP, you must set use_orig_params=True for shape detection to work.

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from heavyball import Muon

model = FSDP(model, use_orig_params=True)
opt = Muon(model.parameters(), lr=0.02)

Install HeavyBall via pip. Requires PyTorch >= 2.2.

pip install heavyball

When upgrading to 2.0.0, follow these steps:

1. Re-test orthogonalization-sensitive optimizers: If using Muon or similar, you may need to set utils.zeroth_power_mode="qr" to maintain previous behavior.
2. Migrate checkpoints: Use scripts/migrate_optimizer_state.py to convert old state dicts to the new indexed format.
3. Update custom tooling: Ensure any custom state-dict or checkpointing logic can handle transform-indexed keys (e.g., exp_avg_3 instead of exp_avg).

Due to breaking changes in how optimizer state keys are stored (they now include a per-transform index, e.g., exp_avg_3), old checkpoints are incompatible with 2.x.

To migrate a checkpoint, use the provided migration script: python scripts/migrate_optimizer_state.py <checkpoint_path> heavyball.<OptimizerClass>

python scripts/migrate_optimizer_state.py /path/to/checkpoint heavyball.AdamC

Use the migrate_optimizer_state.py script to convert 1.x or 2.x checkpoints to the 3.x format. The script automatically resolves old class names (including aliases) and handles the foreach to multi_tensor key rename in parameter groups.

python scripts/migrate_optimizer_state.py <checkpoint.pt> <OptimizerClass>

HeavyBall provides optimizers with an API compatible with torch.optim. You can use standard replacements like AdamW, specialized preconditioned optimizers like SOAP or LATHER, or orthogonal optimizers like Muon. By default, HeavyBall consumes gradients during step() and clears p.grad. If your training loop requires gradients to remain attached after the optimizer step, pass consume_grad=False.

from heavyball import AdamW

opt = AdamW(model.parameters(), lr=1e-3)