# A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale

arXiv: https://arxiv.org/abs/2309.06497 · HTML: https://mlsanity.com/paper/2309.06497/

## Key facts

- Cited by 77 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 85 references that are in the map.
- First posted on arXiv 12 Sep 2023.

## Builds on (most cited first, 30 of 40)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [Bag of Tricks for Image Classification with Convolutional Neural Networks](https://mlsanity.com/paper/1812.01187/) (2018)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [Limitations of the Empirical Fisher Approximation for Natural Gradient Descent](https://mlsanity.com/paper/1905.12558/) (2019)
- [Compressing Neural Networks with the Hashing Trick](https://mlsanity.com/paper/1504.04788/) (2015)
- [Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis](https://mlsanity.com/paper/1806.03884/) (2018)
- [A Kronecker-factored approximate Fisher matrix for convolution layers](https://mlsanity.com/paper/1602.01407/) (2016)
- [Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models](https://mlsanity.com/paper/2104.05158/) (2021)
- [Fast Convergence of Natural Gradient Descent for Overparameterized Neural Networks](https://mlsanity.com/paper/1905.10961/) (2019)
- [Newton Sketch: A Linear-time Optimization Algorithm with Linear-Quadratic Convergence](https://mlsanity.com/paper/1505.02250/) (2015)
- [Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems](https://mlsanity.com/paper/1909.02107/) (2019)
- [Practical Quasi-Newton Methods for Training Deep Neural Networks](https://mlsanity.com/paper/2006.08877/) (2020)
- [Understanding AdamW through Proximal Methods and Scale-Freeness](https://mlsanity.com/paper/2202.00089/) (2022)
- [A Progressive Batching L-BFGS Method for Machine Learning](https://mlsanity.com/paper/1802.05374/) (2018)
- [Newton-Type Methods for Non-Convex Optimization Under Inexact Hessian Information](https://mlsanity.com/paper/1708.07164/) (2017)
- [Second-Order Optimization for Non-Convex Machine Learning: An Empirical Study](https://mlsanity.com/paper/1708.07827/) (2017)

## Built on by (most cited first, 30 of 77)

- [Old Optimizer, New Norm: An Anthology](https://mlsanity.com/paper/2409.20325/) (2024)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm](https://mlsanity.com/paper/2505.16932/) (2025)
- [NorMuon: Making Muon more efficient and scalable](https://mlsanity.com/paper/2510.05491/) (2025)
- [Modular Duality in Deep Learning](https://mlsanity.com/paper/2410.21265/) (2024)
- [Dion: Distributed Orthonormalized Updates](https://mlsanity.com/paper/2504.05295/) (2025)
- [MARS: Unleashing the Power of Variance Reduction for Training Large Models](https://mlsanity.com/paper/2411.10438/) (2024)
- [A New Perspective on Shampoo's Preconditioner](https://mlsanity.com/paper/2406.17748/) (2024)
- [ASGO: Adaptive Structured Gradient Optimization](https://mlsanity.com/paper/2503.20762/) (2025)
- [Implicit Bias of Spectral Descent and Muon on Multiclass Separable Data](https://mlsanity.com/paper/2502.04664/) (2025)
- [MuonBP: Faster Muon via Block-Periodic Orthogonalization](https://mlsanity.com/paper/2510.16981/) (2025)
- [LoRA Done RITE: Robust Invariant Transformation Equilibration for LoRA Optimization](https://mlsanity.com/paper/2410.20625/) (2024)
- [Accelerating Neural Network Training: An Analysis of the AlgoPerf Competition](https://mlsanity.com/paper/2502.15015/) (2025)
- [COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs](https://mlsanity.com/paper/2502.17410/) (2025)
- [Can We Remove the Square-Root in Adaptive Gradient Methods? A Second-Order Perspective](https://mlsanity.com/paper/2402.03496/) (2024)
- [Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner](https://mlsanity.com/paper/2506.03595/) (2025)
- [Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales](https://mlsanity.com/paper/2512.05620/) (2025)
- [SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training](https://mlsanity.com/paper/2412.13148/) (2024)
- [How Muon's Spectral Design Benefits Generalization: A Study on Imbalanced Data](https://mlsanity.com/paper/2510.22980/) (2025)
- [A Stable Whitening Optimizer for Efficient Neural Network Training](https://mlsanity.com/paper/2506.07254/) (2025)
- [Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning](https://mlsanity.com/paper/2603.09697/) (2026)
- [Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization](https://mlsanity.com/paper/2509.03378/) (2025)
- [What Really Matters in Matrix-Whitening Optimizers?](https://mlsanity.com/paper/2510.25000/) (2025)
- [RMNP: Row-Momentum Normalized Preconditioning for Scalable Matrix-Based Optimization](https://mlsanity.com/paper/2603.20527/) (2026)
- [Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers](https://mlsanity.com/paper/2605.18106/) (2026)
- [When, Where and Why to Average Weights?](https://mlsanity.com/paper/2502.06761/) (2025)
- [Structure-Preserving Digital Twins via Conditional Neural Whitney Forms](https://mlsanity.com/paper/2508.06981/) (2025)
- [FedMuon: Federated Learning with Bias-corrected LMO-based Optimization](https://mlsanity.com/paper/2509.26337/) (2025)
- [SNOO: Step-K Nesterov Outer Optimizer - The Surprising Effectiveness of Nesterov Momentum Applied to Pseudo-Gradients](https://mlsanity.com/paper/2510.15830/) (2025)
- [Clarifying Shampoo: Adapting Spectral Descent to Stochasticity and the Parameter Trajectory](https://mlsanity.com/paper/2602.09314/) (2026)

## Abstract (arXiv)

Shampoo is an online and stochastic optimization algorithm belonging to the AdaGrad family of methods for training neural networks. It constructs a block-diagonal preconditioner where each block consists of a coarse Kronecker product approximation to full-matrix AdaGrad for each parameter of the neural network. In this work, we provide a complete description of the algorithm as well as the performance optimizations that our implementation leverages to train deep networks at-scale in PyTorch. Our implementation enables fast multi-GPU distributed data-parallel training by distributing the memory and computation associated with blocks of each parameter via PyTorch's DTensor data structure and performing an AllGather primitive on the computed search directions at each iteration. This major performance enhancement enables us to achieve at most a 10% performance reduction in per-step wall-clock time compared against standard diagonal-scaling-based adaptive gradient methods. We validate our implementation by performing an ablation study on training ImageNet ResNet50, demonstrating Shampoo's superiority over standard training recipes with minimal hyperparameter tuning.