Advanced Optimizers (AIO)
A comprehensive, all-in-one collection of state-of-the-art optimization algorithms for deep learning. Designed for maximum efficiency, minimal memory footprint, and superior performance across diverse model architectures and training scenarios.
📦 Installation
pip install adv_optm
Requires PyTorch 2.3+ for torch.compile support.
What's New
🌟 Version 2.5.x: The Massive Refactor - Click to expand
This major update introduces a complete architectural refactor of the library:
🆕 New Optimizers & Scaling
SinkSGD_adv: Added a powerful new optimizer to the lineup.- Spectral Scaling: Now available across all optimizers, achieving width/rank invariant updates for highly stable training.
💾 Memory & State Precision Control
- Granular State Precision (
state_precision): Drastically reduce memory overhead with new optimizer state modes:factored(Rank-2 factored mode)fp32(Full precision)bf16_sr&int8_sr(BF16/Int8 with Stochastic Rounding)
- Factored Second Moment (
factored_2nd): Available for all Adam variants. Works seamlessly alongside anystate_precisionsetting to further slash memory usage.
⚙️ Advanced Dynamics & Momentum
- Variance Normalized Momentum (
normed_momentum): Applies optimizer normalization before momentum (Normalization then Momentum/NtM). Available forAdamW_adv,SignSGD_adv, andSinkSGD_adv. - Universal Nesterov Momentum: Replaced the hard-to-tune Simplified_AdEMAMix with Nesterov momentum (
nesterov) and a dedicated coefficient (nesterov_coef) across all optimizers. - Preconditioning & Signs:
- Improved CANS (
accelerated_ns): Enhanced for Muon variants by integrating a dynamic lower bound. - New OrthoGrad modes (
orthogonal_gradient): Standard OrthoGradflattenedand a new matrix-wise modeiterative.
⚓ Weight Decay Innovations
- Centered Weight Decay (
centered_wd): Pulls weights toward their pre-train state (anchor). To save memory, anchor precision (centered_wd_mode) can be set to full, float8, int8, or int4. - Fisher Weight Decay (
fisher_wd): Now available for Adam variants based on the FAdam paper. - Geometric Weight Decay: Added specifically for
SinkSGD_advandSignSGD_adv.
(Note: Lion_Prodigy_adv, Simplified_AdEMAMix, and heuristic cautious/grams modes have been deprecated in favor of these superior, theoretically-grounded features).
Click to see older release notes (v1.2.x - v2.1.x)
Version 2.1.x
- New Optimizer: Added Signum (SignSGD with momentum) to the
SignSGD_advfamily.
Version 2.0.x
- ⚡
torch.compileSupport: Fully implemented for all advanced optimizers. Enable viacompiled_optimizer=Trueto heavily fuse and optimize the optimizer step path. - 📉 1-Bit Factored Mode: Vastly improved implementation via
nnmf_factor=True. - 🛠️ Broad performance and stability improvements across all optimizers.
Version 1.2.x
- Advanced Muon Variants: Brought the groundbreaking Muon optimizer into the fold, enriched with features from recent literature.
| Optimizer | Description |
|---|---|
Muon_adv |
Advanced Muon implementation featuring CANS, NorMuon, Low-Rank Orthogonalization, and more. |
AdaMuon_adv |
Combines Muon's geometry with Adam-like adaptive scaling and sign-based orthogonalization. |
- Prodigy Speedup: Prodigy variants are now 50% faster by eliminating unnecessary CUDA syncs (Shoutout to @dxqb!).
- Stochastic Rounding for BF16: Parameter updates and weight decay now accumulate in float32 and round once at the end.
- Cautious Weight Decay: Implemented for all advanced optimizers (Paper).
- Fused Operations: Transitioned to fused and in-place operations wherever possible.
Quick Start
from adv_optm import AdamW_adv
optimizer = AdamW_adv(
model.parameters(),
lr=1e-3,
)
- For PEFT/LoRA models, tag parameters before passing them to the optimizer so spectral scaling and other features work correctly (refer to parameter tagging in docs)
🔗 Quick Links
| Section | Description |
|---|---|
| Optimizers | Full list of available optimizers with descriptions |
| Settings & Features | Common settings shared across all optimizers |
| Momentum Settings | Momentum-related configurations (Nesterov, normed momentum, etc.) |
| Adam-Related Settings | AdamW, Adopt, and Prodigy specific settings (Kourkoutas-β, FAdam, Atan2, etc.) |
| Muon-Related Settings | Muon and AdaMuon specific settings (Newton-Schulz, CANS, NorMuon, MARS-M, etc.) |
| Memory & Precision | State precision, factored modes, and memory-saving techniques |
| Weight Decay | Weight decay variants (centered, Fisher, cautious, geometric) |
| Advanced Features | OrthoGrad, spectral normalization, stochastic sign, and more |
1 maintainer
Metadata
1 maintainer
Release files for adv-optm 2.5.16
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| adv_optm-2.5.16.tar.gz | 70.0 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| adv_optm-2.5.16-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 157.9 kB
Release files / adv_optm-2.5.16.tar.gz
| Download URL | adv_optm-2.5.16.tar.gz |
|---|---|
| Size | 70.0 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
2a105ed25ca396acdc57eff682b7499045efc921b7318374c62be5fa0038f3e2
|
|
BLAKE2b-256 checksum How to use checksums |
6e61c9e3fb77d7b233d695bc97993051dcd69ac2007af78330aef1ae25c8544d
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.12.10
|
Release files / adv_optm-2.5.16-py3-none-any.whl
| Download URL | adv_optm-2.5.16-py3-none-any.whl |
|---|---|
| Size | 87.9 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
202bdd27ce2905059423fe7715dae36a3d22abf42b88ed1f2cab00ea146bfb2e
|
|
BLAKE2b-256 checksum How to use checksums |
68e9d76a6a509bd4136f97e8f1fb176ee255c4792cbb84a14487c768bd1ba359
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.12.10
|
