CoChat
Sign Up
PreprintCited by 70Stochastic Gradient Optimization TechniquesParallel Computing and Optimization TechniquesMachine Learning and Data Classification

The DOI matches its DataCite record.

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

Open access
Shi, Hao-Jun Michael, Lee, Tsung-Hsien, Iwasaki, Shintaro, Gallego-Posada, Jose, Li, Zhijing, Rangadurai, Kaushik, Mudigere, Dheevatsa, Rabbat, Michael·2023·arXiv
DOI 10.48550/arxiv.2309.06497arXiv 2309.06497
Open this paper in CoChat — free accountOpen PDF

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.

Overview

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2309.06497

Cites

Statistics

Verification

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

  • The approximation results in Section 3.2 hold for other algorithms like Shampoo or its versions, assuming a typical choice of momentum schemes.

    “Finally, we discuss some of the most popular optimizers in recent years, but other important algorithms like Shampoo [30, 58] or its versions are also amenable to this analysis, and the approximation results in Section 3.2 hold for them (assuming a typical choice of momentum schemes).”

    From How Memory in Optimization Algorithms Implicitly Modifies the Loss (2025) · cites this paper as [58]

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Citations

Cites

No references could be resolved yet.

Research works printed but not linked: 61

  1. Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. 2020. Disentangling Adaptive Gradient Methods from Learning Rates. arXiv:2002.11803 (2020). · arXiv
  2. Rohan Anil, Sandra Gadanho, Da Huang, Nijith Jacob, Zhuoshu Li, Dong Lin, Todd Phillips, Cristina Pop, Kevin Regan, Gil I Shamir, et al. 2022. On the Factory Floor: ML Engineering for Industrial-Scale Ads Recommendation Models. arXiv:2209.05310 (2022). · arXiv
  3. Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. 2020. Scalable Second Order Optimization for Deep Learning. arXiv:2002.09018 (2020). · arXiv
  4. Jimmy Ba, Roger Grosse, and James Martens. 2017. Distributed Second-Order Optimization using Kronecker-Factored Approximations. In ICLR.
  5. Albert S Berahas, Raghu Bollapragada, and Jorge Nocedal. 2020. An investigation of Newton-Sketch and subsampled Newton methods. Optimization Methods and Software 35, 4 (2020), 661-680.
  6. Albert S Berahas, Majid Jahani, Peter Richtárik, and Martin Takáč. 2022. Quasi-Newton methods for machine learning: forget the past, just sample. Optimization Methods and Software 37, 5 (2022), 1668-1704.
  7. Albert S Berahas, Jorge Nocedal, and Martin Takác. 2016. A Multi-Batch L-BFGS Method for Machine Learning. In NeurIPS.
  8. Raghu Bollapragada, Richard Byrd, and Jorge Nocedal. 2018a. Adaptive Sampling Strategies for Stochastic Optimization. SIAM Journal on Optimization 28, 4 (2018), 3312-3343.
  9. Raghu Bollapragada, Jorge Nocedal, Dheevatsa Mudigere, Hao-Jun Shi, and Ping Tak Peter Tang. 2018b. A Progressive Batching L-BFGS Method for Machine Learning. In ICML.
  10. Antoine Bordes, Léon Bottou, and Patrick Gallinari. 2009. SGD-QN: Careful Quasi-Newton Stochastic Gradient Descent. Journal of Machine Learning Research 10 (2009), 1737-1754.
  11. Léon Bottou et al. 1991. Stochastic Gradient Learning in Neural Networks. Proceedings of Neuro-Nîmes 91, 8 (1991), 12.
  12. Léon Bottou, Frank E Curtis, and Jorge Nocedal. 2018. Optimization Methods for Large-Scale Machine Learning. SIAM Rev. 60, 2 (2018), 223-311.
  13. NeurIPS.
  14. Wenlin Chen, James Wilson, Stephen Tyree, Kilian Weinberger, and Yixin Chen. 2015. Compressing Neural Networks with the Hashing Trick. In ICML.
  15. Aaron Defazio. 2020. Momentum via Primal Averaging: Theoretical Insights and Learning Rate Schedules for Non-Convex Optimization. arXiv:2010.00406 (2020). · arXiv
  16. Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. 2009. ImageNet: A Large-Scale Hierarchical Image Database. In CVPR.
  17. Aditya Devarakonda, Maxim Naumov, and Michael Garland. 2017. AdaBatch: Adaptive Batch Sizes for Training Deep Neural Networks. arXiv:1712.02029 (2017). · arXiv
  18. Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. 2021. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In ICLR.
  19. John Duchi, Elad Hazan, and Yoram Singer. 2011. Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. Journal of Machine Learning Research 12, 7 (2011).
  20. Massimiliano Fasi, Nicholas J Higham, and Xiaobo Liu. 2023. Computing the Square Root of a Low-Rank Perturbation of the Scaled Identity Matrix. SIAM J. Matrix Anal. Appl. 44, 1 (2023), 156-174.
  21. Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. 2018. Fast Approximate Natural Gradient Descent in a Kronecker Factored Eigenbasis. In NeurIPS.
  22. Donald Goldfarb, Yi Ren, and Achraf Bahamou. 2020. Practical Quasi-Newton Methods for Training Deep Neural Networks. NeurIPS (2020).
  23. Gene H Golub and Charles F Van Loan. 2013. Matrix Computations. JHU press.
  24. Roger Grosse and James Martens. 2016. A Kronecker-factored approximate Fisher matrix for convolution layers. In ICML.
  25. Maya R Gupta, Samy Bengio, and Jason Weston. 2014. Training Highly Multiclass Classifiers. Journal of Machine Learning Research 15, 1 (2014), 1461-1492.
  26. Vineet Gupta, Tomer Koren, and Yoram Singer. 2018. Shampoo: Preconditioned Stochastic Tensor Optimization. In ICML.
  27. Stephen J. Hanson and Lorien Y. Pratt. 1988. Comparing Biases for Minimal Network Construction with Back-Propagation. In NeurIPS.
  28. Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016. Deep Residual Learning for Image Recognition. In CVPR.
  29. Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, and Mu Li. 2019. Bag of Tricks for Image Classification with Convolutional Neural Networks. In CVPR.
  30. Dmytro Ivchenko, Dennis Van Der Staay, Colin Taylor, Xing Liu, Will Feng, Rahul Kindi, Anirudh Sudarshan, and Shahin Sefati. 2022. TorchRec: a PyTorch Domain Library for Recommendation Systems. In RecSys.
  31. Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. 2018. Averaging Weights Leads to Wider Optima and Better Generalization. UAI.
  32. Diederik P Kingma and Jimmy Ba. 2015. Adam: A Method for Stochastic Optimization. In ICLR.
  33. Frederik Kunstner, Philipp Hennig, and Lukas Balles. 2019. Limitations of the empirical Fisher approximation for natural gradient descent. NeurIPS (2019).
  34. Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. 1998. Gradient-based learning applied to document recognition. Proc. IEEE 86, 11 (1998), 2278-2324.
  35. Ilya Loshchilov and Frank Hutter. 2019. Decoupled Weight Decay Regularization. In ICLR.
  36. James Martens. 2020. New Insights and Perspectives on the Natural Gradient Method. Journal of Machine Learning Research 21, 1 (2020), 5776-5851.
  37. James Martens et al. 2010. Deep learning via Hessian-free optimization. In ICML.
  38. James Martens, Jimmy Ba, and Matt Johnson. 2018. Kronecker-factored Curvature Approximations for Recurrent Neural Networks. In ICLR.
  39. James Martens and Roger Grosse. 2015. Optimizing Neural Networks with Kronecker-factored Approximate Curvature. In ICML.
  40. James Martens and Ilya Sutskever. 2011. Learning Recurrent Neural Networks with Hessian-Free Optimization. In ICML.
  41. James Martens and Ilya Sutskever. 2012. Training Deep and Recurrent Networks with Hessian-Free Optimization. Neural Networks: Tricks of the Trade: Second Edition (2012), 479-535.
  42. Dheevatsa Mudigere, Yuchen Hao, Jianyu Huang, Zhihao Jia, Andrew Tulloch, Srinivas Sridharan, Xing Liu, Mustafa Ozdal, Jade Nie, Jongsoo Park, et al. 2022. Software-hardware co-design for fast and scalable training of deep learning recommendation models. In ISCA.
  43. Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sundaraman, Jongsoo Park, Xi- aodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G Azzolini, et al. 2019. Deep Learning Recommendation Model for Personalization and Recommendation Systems. arXiv:1906.00091 (2019). · arXiv
  44. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. 2019. PyTorch: An Imperative Style, High-Performance Deep Learning Library. NeurIPS (2019).
  45. Mert Pilanci and Martin J Wainwright. 2017. Newton Sketch: A Near Linear-Time Optimization Algorithm with Linear- Quadratic Convergence. SIAM Journal on Optimization 27, 1 (2017), 205-245.
  46. Boris T Polyak and Anatoli B Juditsky. 1992. cceleration of Stochastic Approximation by Averaging. SIAM Journal on Control and Optimization 30, 4 (1992), 838-855.
  47. Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. 2018. On the Convergence of Adam and Beyond. In ICLR.
  48. Yi Ren and Donald Goldfarb. 2021. Tensor Normal Training for Deep Learning Models. NeurIPS.
  49. Herbert Robbins and Sutton Monro. 1951. A Stochastic Approximation Method. Annals of Mathematical Statistics (1951), 400-407.
  50. David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. 1986. Learning representations by back-propagating errors. Nature 323, 6088 (1986), 533-536.
  51. Nicol N Schraudolph, Jin Yu, and Simon Günter. 2007. A Stochastic Quasi-Newton Method for Online Convex Optimization. In AISTATS.
  52. Noam Shazeer and Mitchell Stern. 2018. Adafactor: Adaptive Learning Rates with Sublinear Memory Cost. In ICML.
  53. Hao-Jun Michael Shi, Dheevatsa Mudigere, Maxim Naumov, and Jiyan Yang. 2020. Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems. In SIGKDD.
  54. Yue Song, Nicu Sebe, and Wei Wang. 2022. Fast Differentiable Matrix Square Root. ICLR.
  55. Wei Tao, Zhisong Pan, Gaowei Wu, and Qing Tao. 2018. Primal Averaging: A New Gradient Evaluation Step to Attain the Optimal Individual Convergence. IEEE Transactions on Cybernetics 50, 2 (2018), 835-845.
  56. Peng Xu, Fred Roosta, and Michael W Mahoney. 2020a. Newton-type methods for non-convex optimization under inexact Hessian information. Mathematical Programming 184, 1-2 (2020), 35-70.
  57. Peng Xu, Jiyan Yang, Fred Roosta, Christopher Ré, and Michael W Mahoney. 2016. Sub-sampled Newton Methods with Non-uniform Sampling. NeurIPS.
  58. Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. 2020. Large Batch Optimization for Deep Learning: Training BERT in 76 minutes. In ICLR.
  59. Guodong Zhang, James Martens, and Roger B Grosse. 2019. Fast Convergence of Natural Gradient Descent for Over- Parameterized Neural Networks. NeurIPS.
  60. Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. 2022. OPT: Open Pre-trained Transformer Language Models. arXiv:2205.01068 (2022). · arXiv
  61. Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al. 2023. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. arXiv:2304.11277 (2023). · arXiv

Cited by

Sorted by citations

Metadata

Authors
Shi, H.-J. M., Lee, T.-H., Iwasaki, S., Gallego-Posada, J., Li, Z., Rangadurai, K., Mudigere, D., & Rabbat, M.
Affiliations
Meta (United States) (US)
ORCID
Venue
arXiv
Publisher
arXiv
Published
2023
ISSN
2331-8422 (ISSN-L)
DOI
10.48550/arxiv.2309.06497
arXiv
2309.06497
OpenAlex
W4386793277
Cited by
70
References
0
Citations by year
2026: 30 · 2025: 30 · 2024: 10
Type
posted content
Language
en
Open access
Licence
arxiv.org/licenses/nonexclusive-distrib/1.0/
Locations
  • arXiv (Cornell University) (submittedVersion)
  • arXiv (Cornell University)
Topics
  • Physical Sciences › Computer Science › Artificial Intelligence › Stochastic Gradient Optimization Techniques
  • Physical Sciences › Computer Science › Hardware and Architecture › Parallel Computing and Optimization Techniques
  • Physical Sciences › Computer Science › Artificial Intelligence › Machine Learning and Data Classification
Keywords
PyTorch, multi-GPU training, AdaGrad, adaptive gradient methods, large-scale neural network training

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review