Install & Compatibility
Where this runs
tested against v0.14.5 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
2159MB installed
● package 2159MB
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
AdamW
✓ from heavyball import AdamW
✗ from heavyball import ForeachAdamW
In HeavyBall v3.0.0, `Foreach*` prefixes were removed from optimizer class names to simplify the public API; use the short, canonical names instead.
SOAP
✓ from heavyball import SOAP
Muon
✓ from heavyball import Muon
This quickstart demonstrates how to use a HeavyBall optimizer, such as `AdamW`, with a simple PyTorch model and a basic training loop. It covers model and data preparation, optimizer initialization, and the standard `zero_grad()`, `backward()`, and `step()` calls. HeavyBall optimizers are designed as drop-in replacements for `torch.optim` classes.
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from heavyball import AdamW # Or any other HeavyBall optimizer
# 1. Define a dummy model
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 1)
def forward(self, x):
return self.linear(x)
model = SimpleModel()
# 2. Prepare dummy data
X = torch.randn(100, 10)
y = torch.randn(100, 1)
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=16)
# 3. Initialize the HeavyBall optimizer
optimizer = AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()
# 4. Training loop (simplified)
num_epochs = 5
for epoch in range(num_epochs):
for batch_X, batch_y in dataloader:
optimizer.zero_grad()
output = model(batch_X)
loss = loss_fn(output, batch_y)
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")
Debug
Known issues
breakingHeavyBall v3.0.0 removed `Foreach*` prefixes from optimizer class names (e.g., `ForeachAdamW` is now `AdamW`). Code relying on the old naming convention will break.fixUpdate optimizer imports and instantiations to use the new, shorter class names (e.g., `from heavyball import AdamW`).
affects: >=3.0.0
breakingHeavyBall v2.2.0 introduced changes to the SOAP optimizer infrastructure. Custom SOAP variants created for earlier versions may not work out-of-the-box.fixRefer to the v2.2.0 release notes for guidance on converting existing SOAP configurations to the new infrastructure. Updates are often trivial.
affects: >=2.2.0
gotchaHeavyBall's default division backend (`eps_clamp`) differs from the industry standard (`eps_add`) used by PyTorch and Optax, potentially leading to meaningfully different numerical results if not accounted for.fixTo align with standard behavior, set `heavyball.utils.default_division_backend = "eps_add"` early in your script. Other options like `atan2` are also available.
affects: >=2.2.1
gotchaWhen using ECC (Error Correction Code) with `torch.compile`, earlier versions (pre-v2.3.1) could experience `torch.compile` fusing away crucial ECC math, leading to incorrect results, particularly with stochastic rounding.fixUpgrade to HeavyBall v2.3.1 or later to leverage the internal fixes (manual bit arithmetic). If using older versions, avoid ECC with stochastic rounding.
affects: <2.3.1
breakingHeavyBall v2.0.0 introduced significant numerical stability improvements, SVD computation accuracy, and a reworked chainable backend, impacting checkpointing. Optimizer checkpoints saved with HeavyBall v1.x are not directly compatible.fixUse the provided `scripts/migrate_optimizer_state.py` utility to convert pre-2.0 optimizer checkpoints. Consult the v2.0.0 and v3.0.0 migration guides for detailed instructions.
affects: >=2.0.0 (from 1.x)
gotchaHeavyBall optimizers, by default, consume gradients during `step()` and clear `p.grad`. If your training loop requires gradients to remain attached after the optimizer step (e.g., for gradient accumulation or logging), they will be cleared.fixInitialize your optimizer with `consume_grad=False` (e.g., `AdamW(model.parameters(), lr=1e-3, consume_grad=False)`) to prevent gradients from being cleared automatically.
affects: All versions
Errors
Common errors & fixes
AttributeError: module 'heavyball' has no attribute 'ForeachAdamW'
Attempting to use an optimizer name with the `Foreach*` prefix (e.g., `ForeachAdamW`) after upgrading to HeavyBall v3.0.0 or later, where these prefixes were removed.
fixUpdate your import statements and optimizer instantiations to use the simplified, shorter class names. For `ForeachAdamW`, use `from heavyball import AdamW`.
RuntimeError: Error(s) in loading state_dict for SimpleModel: Unexpected key(s) in state_dict: "optimizer_states.0.state.step".
Loading a model or optimizer checkpoint saved with an older version of HeavyBall (e.g., v1.x or v2.x) into a newer version (v2.0.0+ or v3.0.0+) without applying necessary migration steps due to changes in internal state representation.
fixFor checkpoints saved with HeavyBall v1.x, use the `scripts/migrate_optimizer_state.py` utility provided in the repository. For v2.x checkpoints, consult the v3.0.0 migration guide for specific conversion steps if any are needed.
Optimizer step produces significantly different or worse convergence compared to `torch.optim` or Optax.
The default division backend used by HeavyBall (`eps_clamp`) for calculating adaptive learning rates or update norms differs from the `eps_add` method commonly used in `torch.optim` and Optax, leading to numerical discrepancies.
fixTo match the standard behavior, set the division backend globally before initializing optimizers: `import heavyball.utils; heavyball.utils.default_division_backend = "eps_add"`.
ModuleNotFoundError: No module named 'heavyball.optimizers'
Incorrect import path for optimizers. HeavyBall optimizers are typically available directly under the `heavyball` namespace, not a nested `heavyball.optimizers` module.
fixChange your import statement from `from heavyball.optimizers import AdamW` to `from heavyball import AdamW`.
Upgrade
Version history
3.2.0latest on PyPI · released May 13, 2026
Audit
Dependencies
torchrequiredHeavyBall is a PyTorch optimizer library and requires PyTorch >= 2.2 for optimal functionality, especially with `torch.compile` features.