How Inner Optimizers Shape Replica Scaling
Organizations: Mila · Concordia University · Université de Montréal · CNRS
Abstract
Training large language models (LLMs) across geographically distributed and poorly interconnected compute pools or data centers has become a topic of significant interest. This has led to increased study of local optimization methods such as DiLoco for LLMs, especially those that use local (or inner) adaptive or preconditioned optimizers, which are critical for high-performance LLM training. However, a central challenge remains: it is well known that increasing the number of model replicas (replica scaling) while fixing the total training budget leads to monotonically decreasing performance. Existing work has empirically demonstrated that the choice of inner optimizer (Muon v.s. AdamW) can significantly improve this trend, but offers only limited explanation of the performance difference. In this work, we consider a range of inner-optimizers and identify two important factors that impact replica scaling: (1) the covariance structure of the pseudo-gradients and (2) the impact of batch size on the pseudogradient's noise scale. Based on these insights, we theoretically and empirically demonstrate that whitening-based pre-conditioning is a powerful inner-optimizer desideratum. Our empirical evaluation extends the results to reveal that whitening-based inner-optimizers achieve strong replica scaling (for 8-96 replicas), with performance relative to DP improving as parameters and data are scaled.
Figures & tables
| Inner optimizer | DP | ||||
|---|---|---|---|---|---|
| SOAP | 2.816 | 2.842 | 2.864 | 2.901 | 2.958 |
| SPlus | 2.853 | 2.837 | 2.853 | 2.873 | 2.912 |
| Muon | 2.858 | 2.979 | 3.073 | 3.230 | 3.417 |
| AdamW | 2.880 | 2.932 | 2.969 | 3.050 | 3.157 |
| Shampoo | 2.842 | 2.899 | 2.947 | 3.002 | 3.097 |
| Curvature ratio | ||
|---|---|---|
| Optimizer | ||
| AdamW | – | – |
| Muon | – | – |
| SOAP | – | – |
| SPlus | – | |
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
| Setting | 178M | 512M |
| Transformer layers | 9 | 12 |
| Hidden width | 1024 | 1536 |
| Attention heads / key-value heads | ||
| Feed-forward width | 2688 | 5440 |
| Activation | SwiGLU | SwiGLU |
| Sequence length | 2048 | 2048 |
LLM Training
Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.
See all 353 papers