Skip to main content

A research project for emerging optimizers other than AdamW

Project description

Emerging Optimizers

codecov CICD NeMo Python 3.12+ GitHub Repo stars Documentation

Overview

Emerging Optimizers is a research project focused on understanding and optimizing the algorithmic behavior of emerging optimizers (including Shampoo, SOAP, Muon, and others) and their implications to performance of GPU systems in LLM training.

⚠️ Note: Emerging-Optimizers is under active development. All APIs are experimental and subject to change. New features, improvements, and documentation updates are released regularly. Your feedback and contributions are welcome, and we encourage you to follow along as new updates roll out.

Background

What are Emerging Optimizers?

Emerging optimizers represent a class of novel optimization algorithms that go beyond traditional first-order methods like Adam or SGD. These include optimizers that use matrix-based (non-diagonal) preconditioning, orthogonalization techniques, and other innovative approaches to achieve faster convergence and improved training efficiency.

Examples include Shampoo, which uses Kronecker-factored preconditioning (arXiv:1802.09568), and Muon, which uses Newton-Schulz orthogonalization (arXiv:2502.16982).

Why They Matter

Emerging optimizers have demonstrated significant practical impact in large-scale language model training. Most notably, Muon was used to train the Kimi K2 model (arXiv:2507.20534), showcasing the effectiveness of these novel approaches at scale. These optimizers can:

  • Achieve faster convergence, reducing the number of training steps required
  • Improve final model quality through better conditioning of the optimization landscape
  • Enable more efficient hyperparameter tuning due to reduced sensitivity to learning rates

Installation

Prerequisites

  • Python 3.12 (Release v0.1.0 is the last version supports Python 3.10)
  • PyTorch 2.0 or higher

Install from Source

git clone https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git
cd Emerging-Optimizers
pip install .

Usage

Example

Refer to tests for usage of different optimizers, e.g. tests/test_orthogonalized_optimizer.py::MuonTest.

Integration with Megatron Core

Integration with Megatron Core is available in dev branch, e.g. emerging_optimizers.py

Benchmarks

Coming soon.

License

This project is licensed under the Apache License 2.0 - see the LICENSE file for details.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

emerging_optimizers-0.3.0.tar.gz (101.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

emerging_optimizers-0.3.0-py3-none-any.whl (102.9 kB view details)

Uploaded Python 3

File details

Details for the file emerging_optimizers-0.3.0.tar.gz.

File metadata

  • Download URL: emerging_optimizers-0.3.0.tar.gz
  • Upload date:
  • Size: 101.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.0.1 CPython/3.12.3

File hashes

Hashes for emerging_optimizers-0.3.0.tar.gz
Algorithm Hash digest
SHA256 4eb96f38280dab17d1fa5bf46bdd7a110b1ab399e98891d0f90c467fb134e0ce
MD5 5ac3ed79f87b4561ce5b6bc8ad28f348
BLAKE2b-256 c5bb3fb2bfacc364b82356f886ce1babbc9b9041bda32bc46ae4a12f77a24d9c

See more details on using hashes here.

File details

Details for the file emerging_optimizers-0.3.0-py3-none-any.whl.

File metadata

File hashes

Hashes for emerging_optimizers-0.3.0-py3-none-any.whl
Algorithm Hash digest
SHA256 e993682a8b93b6953b3056b0d4e2714e3a3182b0c85f456bcc78d76cd7b47429
MD5 33f725c8fb71a375273fdff5f9fb0e0c
BLAKE2b-256 3afdb8aa6cad45beae4d6b585ad55dad32ed95ac74097dd7bfb8952a72421fee

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page