pytorch-optimizer v4.0.0

pytorch-optimizer v4.0.0

We're releasing pytorch-optimizer 4.0.0 with more optimizers, new execution options, and extensive correctness fixes.

Since 3.0.0, we've expanded the collection to 127 registered optimizers and refactored the shared utilities, package exports, and tests. This release also raises the minimum requirements to Python 3.10 and PyTorch 2.1.

Highlights since 3.0.0

• More optimizers: Additions include the Muon family, SOAP, AdEMAMix, APOLLO, SCION, SPAM, and Magma. The latest additions are NorMuon and SaRA.
• Foreach support: Batched updates across supported optimizers to reduce Python loop overhead, especially on CUDA.
• Compilation: Optional optimizer-step compilation through create_optimizer(..., compile=True), with tested paths for Lion, native PyTorch AdamW, and StableAdamW.
• Refactoring: Centralized exports, modernized type hints, shared update utilities, and consolidated training and interface tests.
• Correctness: Fixes to update rules, weight decay, momentum, bias correction, checkpoint restoration, losses, and learning rate schedulers.

Upgrading to 4.0.0

Python 3.10+ and PyTorch 2.1+ are required. Compilation remains opt-in. Use foreach=False to disable batched updates where supported. Corrections to optimizer algorithms may change training results compared with earlier versions.

Thanks to everyone who contributed implementations, fixes, and bug reports

What's Changed

• docs: update changelogs for v3.11.0 by @kozistr in #588
• [Fix] Trigger automerge when label is added by @kozistr in #590
• [Fix] un-swap false positives and false negatives in SoftF1Loss by @plox-sumit in #593
• [Feature] Implement NorMuon optimizer by @shaneraphel in #589
• [Fix] Support strided inputs in Lovasz hinge loss by @AHMETHAKANBEZIR1 in #597
• [Fix] Prevent Tversky overflow on large half-precision masks by @AHMETHAKANBEZIR1 in #599
• [Feature] Support reduction modes in FocalCosineLoss by @AHMETHAKANBEZIR1 in #595
• Track the GRAMS first moment and debias its denominator by @shaneraphel in #537
• [Fix] put the Schedule-Free RAdam bias correction under the square root by @shaneraphel in #545
• [Fix] Correct DiffGrad, AdaFactor, SPlus, and SOAP updates by @kozistr in #601
• [Feature] Implement SaRA optimizer by @kozistr in #602
• [Build] Require Python 3.10 / PyTorch 2.1 and use lerp for moving averages by @kozistr in #603
• [Feature] Compile optimizer steps, batch AdEMAMix, and adopt Trusted Publishing by @kozistr in #604
• [Refactor] Modernize type hints and centralize package exports by @kozistr in #605
• Standardize and correct API docstrings by @kozistr in #606
• Fix optimizer updates, checkpoint precision, and memory overhead by @kozistr in #607
• Consolidate test fixtures and streamline optimizer tests by @kozistr in #608
• Fix pytest xfails and warnings and replace Pyright with ty by @kozistr in #609
• Fix optimizer updates, checkpoints, losses, and schedulers by @kozistr in #610
• Show Python 3.10+ in the README badge by @kozistr in #611
• Merge release changelog PRs using GITHUB_TOKEN by @kozistr in #612
• Fix release version preparation on protected main by @kozistr in #613
• Prepare versions and changelogs in one PR before publishing by @kozistr in #614

New Contributors

• @plox-sumit made their first contribution in #593

Full Changelog: v3.11.0...v4.0.0