Sep, 2023
用于大规模训练神经网络的分布式数据并行 PyTorch 实现的分布式 Shampoo 优化器
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...
TL;DRShampoo 是一种在线和随机优化算法，属于 AdaGrad 方法家族，用于训练神经网络。它构建了一个块对角预处理器，其中每个块都由神经网络的每个参数的粗糙 Kronecker 乘积逼近所组成。我们提供了该算法的完整描述，以及我们实现中利用 PyTorch 进行规模化训练的性能优化。我们的实现通过在每次迭代中通过 PyTorch 的 DTensor 数据结构分配与每个参数块相关的内存和计算，并在计算的搜索方向上执行 AllGather 原语，从而实现快速的多 GPU 分布式数据并行训练。这种主要的性能提升使我们能够实现与标准的对角缩放自适应梯度方法相比每步墙钟时间最多降低 10% 的性能。我们通过在训练 ImageNet ResNet50 上进行消融研究来验证我们的实现，展示了 Shampoo 在最小超参数调整下对标准训练配方的优越性。