Distributed Shampoo: Making Second-Order
Optimization Practical at Scale

Shi, Lee, Iwasaki, Gallego-Posada, Li, Rangadurai, Mudigere, Rabbat — Meta, Mila/Université de Montréal, NVIDIA — Sep 2023
arXiv:2309.06497 ↗ ResNet50 · ImageNet-1k ≤10% step overhead ZeRO-style sharding

From Diagonal to Distributed: the Shampoo lineage

Two independent theoretical roads — online convex optimization and natural-gradient descent — converge on the same Kronecker-factorization trick. This paper is the systems engineering that finally makes it cheap enough to run at scale.

Preconditioner memory vs. layer size (4096×4096 linear layer, log scale)

Full-matrix AdaGrad's preconditioner for one layer this size would be 16M×16M — the bar below is capped and annotated, not to scale.

Kronecker factorization: approximating the impossible

Instead of one dense preconditioner per layer, Shampoo factors it into two much smaller matrices — one per input dimension, one per output dimension — and inverts each separately.

low correlation medium high correlation hover a cell for its value

Refusing to replicate: ZeRO-style preconditioner sharding

Diagonal optimizer state is cheap enough to copy on every worker. Shampoo's factor matrices are not — so each worker owns and computes only a shard, greedily load-balanced by parameter size.

The AllGather round

Each worker computes its shard of search directions, then a single AllGather (via PyTorch DTensor) reunites the full set before the parameter update.

ResNet50 / ImageNet-1k: the receipts

Per-step wall-clock overhead

Optimizer state size (relative to model size)

References

  1. [1]A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale — Shi, Lee, Iwasaki, Gallego-Posada, Li, Rangadurai, Mudigere, Rabbat, 2023 — arXiv:2309.06497
  2. [2]Adaptive Subgradient Methods for Online Learning and Stochastic Optimization — Duchi, Hazan, Singer, 2011
  3. [3]Shampoo: Preconditioned Stochastic Tensor Optimization — Gupta, Koren, Singer, 2018
  4. [4]Scalable Second Order Optimization for Deep Learning — Anil, Gupta, Koren, Regan, Singer, 2020
  5. [5]Optimizing Neural Networks with Kronecker-factored Approximate Curvature — Martens, Grosse, 2015
  6. [6]Optimizing Neural Networks with Kronecker-factored Approximate Curvature (K-FAC) — Martens, Grosse, 2015
  7. [7]On the Factory Floor: ML Engineering for Industrial-Scale Ads Recommendation Models — Anil, Gadanho, Huang, et al., 2022
  8. [8]Adafactor: Adaptive Learning Rates with Sublinear Memory Cost — Shazeer, Stern, 2018
  9. [9]SOAP: Improving and Stabilizing Shampoo using Adam — Vyas, Morwani, Zhao, et al., 2024