LionMuon optimizer speeds up language model training with less compute
LionMuon: Alternating Spectral and Sign Descent for Efficient Training
Machine Learning
Summary
Training large language models requires a lot of computer power, and the choice of optimizer can save significant resources. The authors introduce LionMuon, a new optimizer that mixes two types of update steps to balance speed and accuracy. It runs an expensive but strong step occasionally and cheaper steps in between, reducing overall computation and communication costs. Tests show LionMuon reaches better training results faster than several popular optimizers.
What this means in practice
- •For deep learning engineers: Train language models more efficiently with fewer compute resources by using LionMuon to reduce expensive optimization steps and communication.
- •For cloud platform operators: Lower GPU usage and training time for large models by deploying LionMuon optimizer in multi-GPU distributed training setups.
Authors
Arman Bolatov, Artem Riabinin, Nikita Kornilov, Andrey Veprikov, Samuel Horváth, Martin Takáč, Aleksandr Beznosikov
Abstract
Pretraining a language model takes enormous compute, and the right optimizer can save a good part of it. Muon's spectral step gives a stronger direction than a sign step, but it is expensive. Every step runs Newton-Schulz iterations on the full matrix and, in distributed training, an extra all-reduce. Sign steps, as in Lion and Signum, are cheap and stay local to each device. We propose LionMuon, which takes one Muon step every $P$ iterations and Lion steps in between, with a single dual-EMA momentum buffer shared by both. Muon's compute and communication are paid once per $P$ steps, and the optimizer state is half of AdamW's. A single-EMA variant, SignMuon, already improves on Muon. We prove complexity bounds under heavy-tailed noise in which the period sets an interpolation between Muon's and Lion's smoothness and noise constants, and which say when LionMuon is faster than both. On 124M and 355M models trained on FineWeb, LionMuon with $P=2$ and $P=5$ reaches a lower loss than Muon, AdamW, Lion and Signum at the same number of tokens. Under 4-GPU data-parallel training it reaches Muon's final loss with a third less wall-clock on PCIe, and it beats the communication-efficient Muon variants Dion and MuonBP on loss at no more exposed communication, while keeping the exact gradient. Code: https://github.com/brain-lab-research/lion-muon