[论文解读] Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training

Yuanzhong Xu, HyoukJoong Lee+4 more|arXiv (Cornell University)|Apr 28, 2020

Advanced Neural Network Applications参考文献 19被引用 19

一句话总结

本文提出在数据并行训练中自动对权重更新计算进行跨副本分片，以减少因所有副本上重复且未分区的更新所导致的性能瓶颈。通过在 XLA 中使用静态分析和图变换，该方法高效地在副本间分片权重和优化器辅助变量，并优化通信，无需代码更改或额外硬件，在如 Transformer 等大规模模型上实现了高达 45% 的加速。

ABSTRACT

In data-parallel synchronous training of deep neural networks, different devices (replicas) run the same program with different partitions of the training batch, but weight update computation is repeated on all replicas, because the weights do not have a batch dimension to partition. This can be a bottleneck for performance and scalability in typical language models with large weights, and models with small per-replica batch size which is typical in large-scale training. This paper presents an approach to automatically shard the weight update computation across replicas with efficient communication primitives and data formatting, using static analysis and transformations on the training computation graph. We show this technique achieves substantial speedups on typical image and language models on Cloud TPUs, requiring no change to model code. This technique helps close the gap between traditionally expensive (ADAM) and cheap (SGD) optimizers, as they will only take a small part of training step time and have similar peak memory usage. It helped us to achieve state-of-the-art training performance in Google's MLPerf 0.6 submission.

研究动机与目标

• 解决数据并行训练中因重复且未分区的权重更新计算导致的性能瓶颈。
• 减少 ADAM 等昂贵优化器的运行时间，这些优化器原本因在所有副本上执行完整权重更新而主导训练时间。
• 在不增加额外设备的前提下，实现对权重和辅助变量（如动量、方差）在现有副本间的高效分片。
• 通过智能分片策略和通信模式选择，最小化通信和数据格式化开销。
• 在保持与现有模型代码兼容的同时，实现在大规模模型上的显著加速和内存节省。

提出的方法

• 对 XLA 计算图进行静态分析，识别适合分片的重复操作（如权重更新）。
• 使用控制流分析确定最优通信点，并估算分片候选操作的性能增益。
• 转换计算图，在关键位置插入分片操作和通信原语（如 all-gather、all-reduce）。
• 设计与加速器内存布局（如分块内存）对齐的数据分片格式，以最小化通信和内存访问成本。
• 针对小规模（最小化分片大小）和大规模（最小化通信延迟）训练场景，采用不同的分片策略。
• 利用 XLA 的函数式中间表示（IR），副作用最小化，以简化分析并支持高级优化（如减少权重张量的存活范围）。

实验结果

研究问题

• RQ1 是否可以在不修改模型代码的前提下，自动对数据并行训练中的权重更新计算在副本间进行分片？
• RQ2 在副本间分片权重和辅助变量的性能与通信开销之间存在怎样的权衡？
• RQ3 自动分片在多大程度上通过减少权重和辅助张量的存活范围，影响内存使用情况，特别是峰值内存？
• RQ4 在大规模训练中，分片在多大程度上能减少 ADAM 等昂贵优化器相对于 SGD 的运行时间？
• RQ5 在权重大小和批量大小各不相同的模型中，分片策略的有效性有何差异？

主要发现

• 在 1024 个 TPUv3 芯片上，该技术使 Transformer 模型的单步时间减少了 45%，从 46.5ms 降低至 25.6ms。
• 即使在较小的模型（如 ResNet-50）上，也观察到在大规模（1024 个芯片）下 6% 的加速，表明其具有广泛的适用性。
• 语言模型（如 Transformer）在小规模（16 个芯片）下实现了 9% 的性能提升，凸显了其在大权重模型中的显著影响。
• 该优化通过重用辅助变量的缓冲区，减少了峰值内存使用，尤其使 NCF 等辅助变量开销较大的模型受益。
• 该方法使 ADAM 和 SGD 的峰值内存使用量相当，有效弥合了昂贵优化器与廉价优化器之间的内存差距。
• 该方法无需修改模型代码，仅使用现有副本，避免了额外的硬件或基础设施需求。