| Name | heavyball JSON |
| Version |
2.2.0
JSON |
| download |
| home_page | None |
| Summary | Efficient Optimizers |
| upload_time | 2025-10-12 21:25:44 |
| maintainer | None |
| docs_url | None |
| author | None |
| requires_python | >=3.9 |
| license | None |
| keywords |
torch
optimizer
muon
soap
psgd
|
| VCS |
 |
| bugtrack_url |
|
| requirements |
No requirements were recorded.
|
| Travis-CI |
No Travis.
|
| coveralls test coverage |
No coveralls.
|
# heavyball
[][pypi]
[][license]
_High-performance, extensible, chainable optimizers for PyTorch._
## Why heavyball
- **Lightning-Fast Training**: Batched `foreach` operations deliver significant speedups on large models.
- **Adaptive & Extensible**: Built-in AdamW, RMSprop, Schedule-Free algorithms, and PaLM-inspired schedules.
- **Plug-and-Play**: Drop-in replacements for `torch.optim` with seamless integration.
- **Customizable**: Chainable API lets you compose optimizers and transforms (MARS correction, cautious updates, orthogonal updates).
- **Battle-Tested**: Extensive benchmarks and real-world examples included.
## Key Features
- Foreach-based optimizers: `ForeachAdamW`, `ForeachRMSprop`, `ForeachSFAdamW`, `Muon`, `ADOPT`, `MSAM` (Momentum SAM), …
- Schedule-Free optimizers with dynamic learning rate adaptation.
- Advanced update rules: MARS correction, cautious updates, PaLM beta2 scheduling.
- Chainable transforms for custom optimization recipes.
- Comprehensive benchmark suite packaged separately as LightBench (`../LightBench`).
- Detailed documentation and example-driven tutorials.
## Quickstart
**Install:**
```bash
pip install heavyball
```
**Basic usage:**
```python
import torch
from torch import nn
from heavyball import ForeachAdamW
model = nn.Sequential(
nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10)
)
optimizer = ForeachAdamW(model.parameters(), lr=1e-3)
for data, target in dataloader:
optimizer.zero_grad()
output = model(data)
loss = torch.nn.functional.cross_entropy(output, target)
loss.backward()
optimizer.step()
```
## Benchmarks
> Reproduce benchmarks with LightBench (install it via `pip install -e ../LightBench` from the repo root):
> ```bash
> python3 -m lightbench.run_all_benchmarks --opt ForeachSOAP --opt LaProp --opt AdamW --opt Muon --opt ForeachCachedNewtonPSGD --opt RMSprop --opt OrthoLaProp --opt ForeachSFAdamW --opt ForeachADOPT --opt LaPropOrtho --opt CachedPSGDKron --opt SignLaProp --opt ForeachSOLP --opt PSGDLRA --opt NewtonPSGDLRA --opt NewtonHybrid2PSGDKron --opt NewtonHybrid2PSGDLRA --opt mars-NewtonHybrid2PSGDLRA --opt MSAMLaProp --opt mars-adaptive-NewtonHybrid2PSGDKron --opt mars-ortho-NewtonHybrid2PSGDKron --opt MuonLaProp --opt mars-unscaled-NewtonHybrid2PSGDKron --opt mars-NewtonHybrid2PSGDKron --opt cautious-AdamW --opt unscaled_cautious-AdamW --opt mars-AdamW --dtype float32 --steps 1000000 --trials 1000 --parallelism 256 --seeds 1 --difficulties trivial --difficulties easy --difficulties medium --difficulties hard --difficulties extreme --difficulties nightmare --timeout 2880
> ```
## Migrating from HeavyBall 1.x
- Read the detailed [2.0.0 migration notes](docs/heavyball2.md) for an end-to-end checklist, including guidance for restoring legacy behaviour when needed.
- Use `scripts/migrate_optimizer_state.py` to rewrite pre-2.0 optimizer checkpoints:
```bash
python scripts/migrate_optimizer_state.py path/to/checkpoint.pt heavyball.ForeachAdamW --state-key optimizer
```
The utility renames legacy state entries, fans them out per parameter view, and injects the HeavyBall metadata block expected by 2.0.0.
## Contributing
We welcome contributions! Please check the [issue tracker][tracker] and follow these steps:
1. Fork the repo and create a feature branch.
2. Install dev dependencies: `pip install -e .[dev]`.
3. Run tests: `pytest`.
4. Submit a pull request.
## License
BSD 3-Clause — see the [LICENSE](LICENSE) file.
---
<p align="center">
Made by the HeavyBall team.
</p>
[pypi]: https://pypi.org/project/heavyball/
[license]: LICENSE
[tracker]: https://github.com/HomebrewML/HeavyBall/issues
Raw data
{
"_id": null,
"home_page": null,
"name": "heavyball",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.9",
"maintainer_email": null,
"keywords": "torch, optimizer, muon, soap, psgd",
"author": null,
"author_email": "HeavyBall Authors <github.heavyball@nestler.sh>",
"download_url": "https://files.pythonhosted.org/packages/7d/9a/822a7ead2e9362eb5b796979ec1f9cec8df8c9bacbab40b939572f8d87d4/heavyball-2.2.0.tar.gz",
"platform": null,
"description": "# heavyball\n\n[][pypi]\n[][license]\n\n_High-performance, extensible, chainable optimizers for PyTorch._\n\n## Why heavyball\n\n- **Lightning-Fast Training**: Batched `foreach` operations deliver significant speedups on large models.\n- **Adaptive & Extensible**: Built-in AdamW, RMSprop, Schedule-Free algorithms, and PaLM-inspired schedules.\n- **Plug-and-Play**: Drop-in replacements for `torch.optim` with seamless integration.\n- **Customizable**: Chainable API lets you compose optimizers and transforms (MARS correction, cautious updates, orthogonal updates).\n- **Battle-Tested**: Extensive benchmarks and real-world examples included.\n\n## Key Features\n\n- Foreach-based optimizers: `ForeachAdamW`, `ForeachRMSprop`, `ForeachSFAdamW`, `Muon`, `ADOPT`, `MSAM` (Momentum SAM), \u2026\n- Schedule-Free optimizers with dynamic learning rate adaptation.\n- Advanced update rules: MARS correction, cautious updates, PaLM beta2 scheduling.\n- Chainable transforms for custom optimization recipes.\n- Comprehensive benchmark suite packaged separately as LightBench (`../LightBench`).\n- Detailed documentation and example-driven tutorials.\n\n## Quickstart\n\n**Install:**\n```bash\npip install heavyball\n```\n\n**Basic usage:**\n```python\nimport torch\nfrom torch import nn\nfrom heavyball import ForeachAdamW\n\nmodel = nn.Sequential(\n nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10)\n)\noptimizer = ForeachAdamW(model.parameters(), lr=1e-3)\n\nfor data, target in dataloader:\n optimizer.zero_grad()\n output = model(data)\n loss = torch.nn.functional.cross_entropy(output, target)\n loss.backward()\n optimizer.step()\n```\n\n## Benchmarks\n\n> Reproduce benchmarks with LightBench (install it via `pip install -e ../LightBench` from the repo root):\n> ```bash\n> python3 -m lightbench.run_all_benchmarks --opt ForeachSOAP --opt LaProp --opt AdamW --opt Muon --opt ForeachCachedNewtonPSGD --opt RMSprop --opt OrthoLaProp --opt ForeachSFAdamW --opt ForeachADOPT --opt LaPropOrtho --opt CachedPSGDKron --opt SignLaProp --opt ForeachSOLP --opt PSGDLRA --opt NewtonPSGDLRA --opt NewtonHybrid2PSGDKron --opt NewtonHybrid2PSGDLRA --opt mars-NewtonHybrid2PSGDLRA --opt MSAMLaProp --opt mars-adaptive-NewtonHybrid2PSGDKron --opt mars-ortho-NewtonHybrid2PSGDKron --opt MuonLaProp --opt mars-unscaled-NewtonHybrid2PSGDKron --opt mars-NewtonHybrid2PSGDKron --opt cautious-AdamW --opt unscaled_cautious-AdamW --opt mars-AdamW --dtype float32 --steps 1000000 --trials 1000 --parallelism 256 --seeds 1 --difficulties trivial --difficulties easy --difficulties medium --difficulties hard --difficulties extreme --difficulties nightmare --timeout 2880\n> ```\n\n## Migrating from HeavyBall 1.x\n\n- Read the detailed [2.0.0 migration notes](docs/heavyball2.md) for an end-to-end checklist, including guidance for restoring legacy behaviour when needed.\n- Use `scripts/migrate_optimizer_state.py` to rewrite pre-2.0 optimizer checkpoints:\n ```bash\n python scripts/migrate_optimizer_state.py path/to/checkpoint.pt heavyball.ForeachAdamW --state-key optimizer\n ```\n The utility renames legacy state entries, fans them out per parameter view, and injects the HeavyBall metadata block expected by 2.0.0.\n\n\n## Contributing\n\nWe welcome contributions! Please check the [issue tracker][tracker] and follow these steps:\n1. Fork the repo and create a feature branch.\n2. Install dev dependencies: `pip install -e .[dev]`.\n3. Run tests: `pytest`.\n4. Submit a pull request.\n\n## License\n\nBSD 3-Clause \u2014 see the [LICENSE](LICENSE) file.\n\n---\n<p align=\"center\">\n Made by the HeavyBall team.\n</p>\n\n[pypi]: https://pypi.org/project/heavyball/\n[license]: LICENSE\n[tracker]: https://github.com/HomebrewML/HeavyBall/issues\n",
"bugtrack_url": null,
"license": null,
"summary": "Efficient Optimizers",
"version": "2.2.0",
"project_urls": {
"source": "https://github.com/HomebrewML/HeavyBall",
"tracker": "https://github.com/HomebrewML/HeavyBall/issues"
},
"split_keywords": [
"torch",
" optimizer",
" muon",
" soap",
" psgd"
],
"urls": [
{
"comment_text": null,
"digests": {
"blake2b_256": "d8d795cd6b919477c1446f1678b5664342c4237edf22c528deefd52097063230",
"md5": "1b1aa8cfe41a095803e627f23afe1495",
"sha256": "6c1fb3d34e6ec6a42f5795314c1f3e08ac106266b6b4fa5c3e535149df33fbbb"
},
"downloads": -1,
"filename": "heavyball-2.2.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "1b1aa8cfe41a095803e627f23afe1495",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.9",
"size": 53420,
"upload_time": "2025-10-12T21:25:43",
"upload_time_iso_8601": "2025-10-12T21:25:43.429395Z",
"url": "https://files.pythonhosted.org/packages/d8/d7/95cd6b919477c1446f1678b5664342c4237edf22c528deefd52097063230/heavyball-2.2.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": null,
"digests": {
"blake2b_256": "7d9a822a7ead2e9362eb5b796979ec1f9cec8df8c9bacbab40b939572f8d87d4",
"md5": "8cd60709a9d72b61d02251e69b4df2ba",
"sha256": "90d0473887de709ceab486ae593cc8f8df39973e3c28697981a4523a2c213ae3"
},
"downloads": -1,
"filename": "heavyball-2.2.0.tar.gz",
"has_sig": false,
"md5_digest": "8cd60709a9d72b61d02251e69b4df2ba",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.9",
"size": 74733,
"upload_time": "2025-10-12T21:25:44",
"upload_time_iso_8601": "2025-10-12T21:25:44.882057Z",
"url": "https://files.pythonhosted.org/packages/7d/9a/822a7ead2e9362eb5b796979ec1f9cec8df8c9bacbab40b939572f8d87d4/heavyball-2.2.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-10-12 21:25:44",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "HomebrewML",
"github_project": "HeavyBall",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "heavyball"
}
