pytorch-optimizer

The pytorch_optimizer library supports foreach implementations for its optimizers. Using foreach=True typically provides a speedup in execution time on CUDA devices by utilizing vectorized operations, though it may increase peak memory usage.

Based on benchmarks using an MLP model (29M parameters) on an NVIDIA GeForce GTX 1060, the average speedup for foreach implementations is approximately 1.13x.

Key Trade-offs:

• Speed: Most optimizers (e.g., AdaFactor, Lion, SGDW) show significant speedups (ranging from 1.1x to 1.3x).
• Memory: foreach often results in higher peak memory consumption (e.g., AdaFactor saw a +75.3% increase). Some optimizers like LARS and SGDW show 0.0% memory difference.
• Exceptions: Some optimizers like Adan and ADOPT may perform slightly slower or show increased memory overhead when using foreach.

Many optimizers in pytorch_optimizer (including AdaFactor, StableAdamW, Lion, AdaBelief, and Amos) now support the foreach implementation. This can improve training speed by 1.1x to 1.5x at the cost of moderate memory usage.

• Default behavior: Like official PyTorch optimizers, the default is foreach=None. When set to None, CUDA paths will automatically prefer the foreach implementation over the standard for-loop implementation.
• Disabling foreach: If you require the legacy for-loop behavior, explicitly set foreach=False.

To use the 8-bit bitsandbytes optimizers, you must install the package with the bitsandbytes extra. This enables support for the following 8 optimizers:

• bnb_adagrad8bit
• bnb_adam8bit
• bnb_adamw8bit
• bnb_lion8bit
• bnb_lamb8bit
• bnb_lars8bit
• bnb_rmsprop8bit
• bnb_sgd8bit

pip3 install pytorch-optimizer[bitsandbytes]

To run the built-in visualizations for the optimizers, execute one of the following commands from the project root:

1. Using just: just visualize
2. Using python module execution: python3 -m examples.visualize_optimizers

# Option 1
just visualize

# Option 2
python3 -m examples.visualize_optimizers

Install the toolkit using pip.

Requirements:

• Python >=3.8
• PyTorch >=1.10

Optional Integrations: Some features require additional packages which are not installed by default:

• bitsandbytes
• q-galore-torch
• torchao

pip install pytorch-optimizer

In version 2.7.0, rectified versions of certain optimizers were merged into their original classes. Instead of using separate classes, you can now use the original optimizer and enable the rectified behavior by setting the rectify=True parameter.

Specifically:

• diffGrad now replaces the previous diffRGrad + diffGrad distinction.
• Lamb now replaces the previous RaLamb + Lamb distinction.

Example usage:

# To use the rectified version of Lamb
optimizer = Lamb(params, lr=1e-3, rectify=True)

# To use the rectified version of diffGrad
optimizer = diffGrad(params, lr=1e-3, rectify=True)

You can use the Refined Schedule-Free AdamW variant by using the ScheduleFreeAdamW optimizer and setting the decoupling_c parameter.

# Example usage of the refined variant
optimizer = ScheduleFreeAdamW(model.parameters(), decoupling_c=...) 

The Muon, DistributedMuon, AdaMuon, and AdaGO optimizers now support various coefficient presets for the zero_power_via_newton_schulz_5 method. You can use these presets or provide custom schedules via the ns_coeffs parameter.

Available Presets:

• original
• quintic
• polar_express
• polar_express_safer