heavyball 1.7.2__py3-none-any.whl → 2.0.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- heavyball/__init__.py +276 -37
- heavyball/chainable.py +419 -206
- heavyball/helpers.py +808 -0
- heavyball/utils.py +1062 -315
- heavyball-2.0.0.dist-info/METADATA +122 -0
- heavyball-2.0.0.dist-info/RECORD +9 -0
- {heavyball-1.7.2.dist-info → heavyball-2.0.0.dist-info}/WHEEL +1 -1
- heavyball-1.7.2.dist-info/METADATA +0 -939
- heavyball-1.7.2.dist-info/RECORD +0 -8
- {heavyball-1.7.2.dist-info → heavyball-2.0.0.dist-info}/licenses/LICENSE +0 -0
- {heavyball-1.7.2.dist-info → heavyball-2.0.0.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,122 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: heavyball
|
|
3
|
+
Version: 2.0.0
|
|
4
|
+
Summary: Efficient Optimizers
|
|
5
|
+
Author-email: HeavyBall Authors <github.heavyball@nestler.sh>
|
|
6
|
+
Project-URL: source, https://github.com/HomebrewML/HeavyBall
|
|
7
|
+
Project-URL: tracker, https://github.com/HomebrewML/HeavyBall/issues
|
|
8
|
+
Keywords: torch,optimizer,muon,soap,psgd
|
|
9
|
+
Classifier: Intended Audience :: Developers
|
|
10
|
+
Classifier: Intended Audience :: Science/Research
|
|
11
|
+
Classifier: License :: OSI Approved :: BSD License
|
|
12
|
+
Classifier: Natural Language :: English
|
|
13
|
+
Classifier: Operating System :: OS Independent
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Requires-Python: >=3.9
|
|
16
|
+
Description-Content-Type: text/markdown
|
|
17
|
+
License-File: LICENSE
|
|
18
|
+
Requires-Dist: opt-einsum>=3.4.0
|
|
19
|
+
Requires-Dist: torch>=2.7.0
|
|
20
|
+
Requires-Dist: numpy
|
|
21
|
+
Provides-Extra: dev
|
|
22
|
+
Requires-Dist: pre-commit; extra == "dev"
|
|
23
|
+
Requires-Dist: pytest; extra == "dev"
|
|
24
|
+
Requires-Dist: ruff; extra == "dev"
|
|
25
|
+
Requires-Dist: matplotlib; extra == "dev"
|
|
26
|
+
Requires-Dist: seaborn; extra == "dev"
|
|
27
|
+
Requires-Dist: hyperopt; extra == "dev"
|
|
28
|
+
Requires-Dist: pandas; extra == "dev"
|
|
29
|
+
Requires-Dist: typer; extra == "dev"
|
|
30
|
+
Requires-Dist: optuna; extra == "dev"
|
|
31
|
+
Requires-Dist: optunahub; extra == "dev"
|
|
32
|
+
Requires-Dist: botorch; extra == "dev"
|
|
33
|
+
Requires-Dist: hebo; extra == "dev"
|
|
34
|
+
Dynamic: license-file
|
|
35
|
+
|
|
36
|
+
# heavyball
|
|
37
|
+
|
|
38
|
+
[][pypi]
|
|
39
|
+
[][license]
|
|
40
|
+
|
|
41
|
+
_High-performance, extensible, chainable optimizers for PyTorch._
|
|
42
|
+
|
|
43
|
+
## Why heavyball
|
|
44
|
+
|
|
45
|
+
- **Lightning-Fast Training**: Batched `foreach` operations deliver significant speedups on large models.
|
|
46
|
+
- **Adaptive & Extensible**: Built-in AdamW, RMSprop, Schedule-Free algorithms, and PaLM-inspired schedules.
|
|
47
|
+
- **Plug-and-Play**: Drop-in replacements for `torch.optim` with seamless integration.
|
|
48
|
+
- **Customizable**: Chainable API lets you compose optimizers and transforms (MARS correction, cautious updates, orthogonal updates).
|
|
49
|
+
- **Battle-Tested**: Extensive benchmarks and real-world examples included.
|
|
50
|
+
|
|
51
|
+
## Key Features
|
|
52
|
+
|
|
53
|
+
- Foreach-based optimizers: `ForeachAdamW`, `ForeachRMSprop`, `ForeachSFAdamW`, `Muon`, `ADOPT`, `MSAM` (Momentum SAM), …
|
|
54
|
+
- Schedule-Free optimizers with dynamic learning rate adaptation.
|
|
55
|
+
- Advanced update rules: MARS correction, cautious updates, PaLM beta2 scheduling.
|
|
56
|
+
- Chainable transforms for custom optimization recipes.
|
|
57
|
+
- Comprehensive benchmark suite (`benchmark/`).
|
|
58
|
+
- Detailed documentation and example-driven tutorials.
|
|
59
|
+
|
|
60
|
+
## Quickstart
|
|
61
|
+
|
|
62
|
+
**Install:**
|
|
63
|
+
```bash
|
|
64
|
+
pip install heavyball
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
**Basic usage:**
|
|
68
|
+
```python
|
|
69
|
+
import torch
|
|
70
|
+
from torch import nn
|
|
71
|
+
from heavyball import ForeachAdamW
|
|
72
|
+
|
|
73
|
+
model = nn.Sequential(
|
|
74
|
+
nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10)
|
|
75
|
+
)
|
|
76
|
+
optimizer = ForeachAdamW(model.parameters(), lr=1e-3)
|
|
77
|
+
|
|
78
|
+
for data, target in dataloader:
|
|
79
|
+
optimizer.zero_grad()
|
|
80
|
+
output = model(data)
|
|
81
|
+
loss = torch.nn.functional.cross_entropy(output, target)
|
|
82
|
+
loss.backward()
|
|
83
|
+
optimizer.step()
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
## Benchmarks
|
|
87
|
+
|
|
88
|
+
> Reproduce benchmarks with:
|
|
89
|
+
> ```bash
|
|
90
|
+
> python3 -m benchmark.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
|
|
91
|
+
> ```
|
|
92
|
+
|
|
93
|
+
## Migrating from HeavyBall 1.x
|
|
94
|
+
|
|
95
|
+
- Read the detailed [2.0.0 migration notes](docs/heavyball-2.0.0-migration.md) for an end-to-end checklist, including guidance for restoring legacy behaviour when needed.
|
|
96
|
+
- Use `scripts/migrate_optimizer_state.py` to rewrite pre-2.0 optimizer checkpoints:
|
|
97
|
+
```bash
|
|
98
|
+
python scripts/migrate_optimizer_state.py path/to/checkpoint.pt heavyball.ForeachAdamW --state-key optimizer
|
|
99
|
+
```
|
|
100
|
+
The utility renames legacy state entries, fans them out per parameter view, and injects the HeavyBall metadata block expected by 2.0.0.
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
## Contributing
|
|
104
|
+
|
|
105
|
+
We welcome contributions! Please check the [issue tracker][tracker] and follow these steps:
|
|
106
|
+
1. Fork the repo and create a feature branch.
|
|
107
|
+
2. Install dev dependencies: `pip install -e .[dev]`.
|
|
108
|
+
3. Run tests: `pytest`.
|
|
109
|
+
4. Submit a pull request.
|
|
110
|
+
|
|
111
|
+
## License
|
|
112
|
+
|
|
113
|
+
BSD 3-Clause — see the [LICENSE](LICENSE) file.
|
|
114
|
+
|
|
115
|
+
---
|
|
116
|
+
<p align="center">
|
|
117
|
+
Made by the HeavyBall team.
|
|
118
|
+
</p>
|
|
119
|
+
|
|
120
|
+
[pypi]: https://pypi.org/project/heavyball/
|
|
121
|
+
[license]: LICENSE
|
|
122
|
+
[tracker]: https://github.com/HomebrewML/HeavyBall/issues
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
heavyball/__init__.py,sha256=cabACszT-lZaLPkN2rANFxLxkpcDPb8q4p7iL8jfTtQ,28274
|
|
2
|
+
heavyball/chainable.py,sha256=l7uzXKMbJxKn-kgMR9In8BxYKuvIO_y8uL6P5b0LZo8,41250
|
|
3
|
+
heavyball/helpers.py,sha256=zk_S84wpGcvO9P6kn4UeaQUIDowHxcbM9qQITEm2g5I,30267
|
|
4
|
+
heavyball/utils.py,sha256=sK85OOhmPHvAbUjZhkgX5tRfE3ECai0Yx4Zvt4p8z2Q,97794
|
|
5
|
+
heavyball-2.0.0.dist-info/licenses/LICENSE,sha256=G9fFZcNIVWjU7o6Pr_4sJBRCNDU5X-zelSxIJ2D48ms,1323
|
|
6
|
+
heavyball-2.0.0.dist-info/METADATA,sha256=9JKPFmnvaMT9oigMckzaqI5DjUN0hJj3ePzc1ok3Ia8,4973
|
|
7
|
+
heavyball-2.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
|
8
|
+
heavyball-2.0.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
|
|
9
|
+
heavyball-2.0.0.dist-info/RECORD,,
|