title: A Distributed Data-Parallel PyTorch Implementation of the Distributed
    Shampoo Optimizer for Training Neural Networks At-Scale
summary: 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_authors._id: 65091d8bce81b73e63538d25
paper_authors.name: Hao-Jun Michael Shi
paper_authors._id: 65091d8bce81b73e63538d26
paper_authors.name: Tsung-Hsien Lee
paper_authors._id: 65091d8bce81b73e63538d27
paper_authors.name: Shintaro Iwasaki
paper_authors._id: 65091d8bce81b73e63538d28
paper_authors.name: Jose Gallego-Posada
paper_authors._id: 65091d8bce81b73e63538d29
paper_authors.name: Zhijing Li
paper_authors._id: 65091d8bce81b73e63538d2a
paper_authors.name: Kaushik Rangadurai
paper_authors._id: 65091d8bce81b73e63538d2b
paper_authors.name: Dheevatsa Mudigere
paper_authors._id: 65091d8bce81b73e63538d2c
paper_authors.name: Michael Rabbat
paper_title: A Distributed Data-Parallel PyTorch Implementation of the Distributed
    Shampoo Optimizer for Training Neural Networks At-Scale
paper_summary: 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_ai_summary: Shampoo, an AdaGrad-based method using block-diagonal preconditioners with Kronecker product approximations, enhances neural network training performance in PyTorch, especially distributed multi-GPU environments.
paper_ai_keywords: Shampoo AdaGrad block-diagonal preconditioner Kronecker product approximation PyTorch DTensor AllGather distributed data-parallel training ImageNet ResNet50