[论文解读] Memory-Efficient Pipeline-Parallel DNN Training

Deepak Narayanan, Amar Phanishayee+3 more|arXiv (Cornell University)|Jun 16, 2020

Topic Modeling参考文献 22被引用 62

一句话总结

PipeDream-2BW 通过带双缓冲权重更新和自动分区的内存高效流水线并行实现，在保持相似精度的前提下，加速对超大规模 transformer 模型的训练。

ABSTRACT

Many state-of-the-art ML results have been obtained by scaling up the number of parameters in existing models. However, parameters and activations for such large models often do not fit in the memory of a single accelerator device; this means that it is necessary to distribute training of large models over multiple accelerators. In this work, we propose PipeDream-2BW, a system that supports memory-efficient pipeline parallelism. PipeDream-2BW uses a novel pipelining and weight gradient coalescing strategy, combined with the double buffering of weights, to ensure high throughput, low memory footprint, and weight update semantics similar to data parallelism. In addition, PipeDream-2BW automatically partitions the model over the available hardware resources, while respecting hardware constraints such as memory capacities of accelerators and interconnect topologies. PipeDream-2BW can accelerate the training of large GPT and BERT language models by up to 20$ imes$ with similar final model accuracy.

研究动机与目标

• 在加速器内存限制下，提出训练十亿参数级DNN的挑战动机。
• 提出一个内存高效的流水线并行系统，在不产生 flush 引起的停顿的情况下提升吞吐量。
• 在异构互连中自动划分模型，同时满足设备内存约束。
• 与基线相比，展示在大型 transformer 模型（GPT/BERT）上的收敛性和吞吐量提升。

提出的方法

• 引入 2BW（双缓冲权重更新），通过维护两个权重版本并在微小批量上累积梯度来降低内存占用。
• 使用规划算法将 DNN 运算符划分为等量复制的阶段（等倍复制），并分配微小批次和梯度累积，以平衡吞吐量和内存。
• 提供一个变体 PipeDream-Flush，使用单个权重版本和流水线刷新，以在吞吐量代价下实现更低的内存占用。
• 以 1 步延迟形式形式化权重更新语义：W(t+1)=W(t)−η·∇f(W(t−1))，并讨论与 vanilla 更新的收敛性相似性。
• 引入激活重新计算以降低激活存储并增加每个 GPU 的可行微小批量大小。
• 在使用 PyTorch、AWS 上的 NVLink 对大规模 GPT/BERT 模型进行评估，比较 Megatron 风格的张量模型并行和 GPipe。

实验结果

研究问题

• RQ1 在面向十亿参数模型的流水线并行训练中，如何在内存和吞吐量之间取得平衡？
• RQ2 双缓冲权重更新方案是否能够在不发生流水线刷新且内存占用可控的情况下实现高吞吐量？
• RQ3 在保持收敛性的同时，哪些自动分区策略最适合将 transformer 风格的框架映射到异构互连上？
• RQ4 2BW 与 PipeDream-Flush 的语义是否能在大型 NLP 模型上保持与 vanilla 优化器相当的收敛性/准确性？

主要发现

• 在各种 GPT 模型上，相比优化后的模型并行基线实现端到端 1.3x 至 20x 的加速。
• 由于没有流水线刷新且内存占用更低，速度比 GPipe 快最多 3.2x。
• 可以使用 64 个八 GPU 或16-GPU 的设置训练接近 300 亿参数的大型 transformer 模型，在内存受限场景下优于 Megatron 风格的配置。
• 在 BERT 和 GPT 预训练任务中，使用 Adam 优化器配合 2BW 显示出与 vanilla Adam 相近的训练损失轨迹和下游微调准确性。
• PipeDream-Flush 在吞吐量方面有成本的情况下提供更低的内存占用，而 2BW 通过两个权重版本和激活重新计算实现更高吞吐量。
• 内存占用比较显示 GPipe 的内存挑战、PipeDream-Flush 中等占用，以及 PipeDream-2BW 的中间占用。