Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, Michael Rabbat

TL;DR

The paper introduces a distributed PyTorch implementation of the Shampoo optimizer that uses block-diagonal Kronecker-factor preconditioners to bridge the gap between full-matrix and diagonal adaptive methods. By distributing preconditioner computation across workers and employing DTensor with AllGather, the approach achieves near-diagonal overhead while delivering stronger per-step conditioning for training deep networks at scale. Comprehensive ablations and ImageNet/ResNet50 experiments demonstrate faster convergence and reduced training volatility with modest overhead, supporting Shampoo as a practical alternative for large-scale training. The work also details numerous implementation and numerical strategies, including layer-wise grafting, memory- and compute-balancing, and robust numerical solvers, and provides an open-source PyTorch implementation for the community.

Abstract

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.

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

TL;DR

The paper introduces a distributed PyTorch implementation of the Shampoo optimizer that uses block-diagonal Kronecker-factor preconditioners to bridge the gap between full-matrix and diagonal adaptive methods. By distributing preconditioner computation across workers and employing DTensor with AllGather, the approach achieves near-diagonal overhead while delivering stronger per-step conditioning for training deep networks at scale. Comprehensive ablations and ImageNet/ResNet50 experiments demonstrate faster convergence and reduced training volatility with modest overhead, supporting Shampoo as a practical alternative for large-scale training. The work also details numerous implementation and numerical strategies, including layer-wise grafting, memory- and compute-balancing, and robust numerical solvers, and provides an open-source PyTorch implementation for the community.

Abstract

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.
Paper Structure (51 sections, 2 theorems, 58 equations, 8 figures, 5 tables, 3 algorithms)

This paper contains 51 sections, 2 theorems, 58 equations, 8 figures, 5 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Main Contributions
  3. Terminology and Notation
  4. Problem Statement and Shampoo Algorithm
  5. Neural Network Training
  6. Diagonal Adaptive Gradient Methods
  7. The Shampoo Algorithm
  8. Layer-wise Learning Rate Grafting
  9. Implementation Details
  10. Training Heuristics
  11. First and Second Moment Estimation
  12. $\ell_2$-Regularization and (Decoupled) Weight Decay
  13. Momentum and Nesterov Acceleration
  14. Exponent Override and Exponent Multiplier
  15. Numerical Considerations
  16. ...and 36 more sections

Key Result

theorem 1

The momentum method defined in eq:full-momentum1-eq:full-momentum2 is equivalent to the heavy ball iteration and the stochastic primal iterate averaging iteration with $c_t = \frac{\gamma_t}{\gamma_t + 1} \in [0, 1)$ for some $\gamma_t > 0$, $\delta_t = \frac{\alpha_{t - 1}}{\alpha_t} \mu_t = \frac{\gamma_{t - 1}}{\gamma_t + 1}$, and $\alpha_t = \frac{\eta_t}{\gamma_t + 1} = (1 - c_t) \eta_t$. I

Figures (8)

  • Figure 1: Outline of each distributed data-parallel iteration with the Distributed Shampoo optimizer.
  • Figure 2: Picture of block-diagonal and Kronecker product approximations used in Shampoo.
  • Figure 3: A single optimizer step of Distributed Shampoo.
  • Figure 4: Maximum buffer size allocation for AllGather primitive. In practice, the preconditioned tensors is a collection of view() of the 1D communication buffer tensor to avoid extra copy at ($\ast$) in the figure.
  • Figure 5: Picture of merging (left) and blocking (right).
  • ...and 3 more figures

Theorems & Definitions (2)

  • theorem 1
  • theorem 2