用于大规模神经网络训练的分布式数据并行 PyTorch 版 Distributed Shampoo 优化器实现
机器学习
2023-09-14 v1 分布式、并行与集群计算
数学软件
最优化与控制
摘要
Shampoo 是一种在线随机优化算法，属于用于训练神经网络的 AdaGrad 方法族。它构造了一个块对角预条件子，其中每个块由对神经网络每个参数的全矩阵 AdaGrad 的粗粒度 Kronecker 积近似组成。在这项工作中，我们提供了该算法的完整描述以及我们的实现在 PyTorch 中用于大规模训练深度网络所利用的性能优化。我们的实现通过 PyTorch 的 DTensor 数据结构分发每个参数块相关的内存和计算，并在每次迭代对计算出的搜索方向执行 AllGather 原语，从而实现快速多 GPU 分布式数据并行训练。这一主要的性能增强使我们能够实现与标准基于对角缩放的自适应梯度方法相比最多 10% 的每步挂钟时间性能降低。我们通过在对 ImageNet ResNet50 的训练中进行消融研究来验证我们的实现，展示了 Shampoo 在最小超参数调优下相对于标准训练方案的优势。
引用
@article{arxiv.2309.06497,
title = {A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale},
author = {Hao-Jun Michael Shi and Tsung-Hsien Lee and Shintaro Iwasaki and Jose Gallego-Posada and Zhijing Li and Kaushik Rangadurai and Dheevatsa Mudigere and Michael Rabbat},
journal= {arXiv preprint arXiv:2309.06497},
year = {2023}
}
备注
38 pages, 8 figures, 5 tables