[论文解读] Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines

Shigang Li, Torsten Hoefler|arXiv (Cornell University)|Jul 14, 2021

Parallel Computing and Optimization Techniques参考文献 53被引用 8

一句话总结

Chimera 提出了一种双向流水线并行方案，通过同步训练大规模神经网络，相比最先进方法，吞吐量最高提升 2.34 倍。通过智能调度多个流水线中的前向和反向传播，它将流水线空洞减少了高达 50%，并均衡了激活内存使用，从而在 2,048 块 GPU 上实现高效训练，且精度损失最小。

ABSTRACT

Training large deep learning models at scale is very challenging. This paper proposes Chimera, a novel pipeline parallelism scheme which combines bidirectional pipelines for efficiently training large-scale models. Chimera is a synchronous approach and therefore no loss of accuracy, which is more convergence-friendly than asynchronous approaches. Compared with the latest synchronous pipeline approach, Chimera reduces the number of bubbles by up to 50%; benefiting from the sophisticated scheduling of bidirectional pipelines, Chimera has a more balanced activation memory consumption. Evaluations are conducted on Transformer based language models. For a GPT-2 model with 1.3 billion parameters running on 2,048 GPU nodes of the Piz Daint supercomputer, Chimera improves the training throughput by 1.16x-2.34x over the state-of-the-art synchronous and asynchronous pipeline approaches.

研究动机与目标

• 解决现有大规模深度学习训练中流水线并行方案存在的低效与内存不平衡问题。
• 减少流水线空洞并提升系统利用率，同时不损害训练收敛性或精度。
• 在同步训练框架中均衡加速器之间的激活内存消耗。
• 实现在分布式超算系统上对 GPT-2 和 BERT 等大模型的高吞吐、可扩展训练。
• 提供灵活可配置的流水线调度策略，可适配多种模型与系统架构。

提出的方法

• Chimera 采用双向流水线，在多个流水线阶段同时执行前向和反向传播，减少空闲时间和空洞。
• 它使用一种新型调度策略——前向加倍、反向减半或直接拼接，以管理前向与反向传播之间的负载不平衡。
• 通过避免权重过时，保持同步训练，确保收敛稳定性和精度。
• 利用流水线阶段之间的点对点通信，并通过优化微批次大小和流水线深度配置，最小化 allreduce 开销。
• Chimera 动态配置流水线深度（D）和微批次大小（B），以平衡通信、计算与内存使用。
• 一个精确的性能模型指导系统配置，以在不同模型与硬件配置下实现最优吞吐量。

实验结果

研究问题

• RQ1 如何优化流水线并行以减少大规模深度学习训练中的空洞并提升利用率？
• RQ2 双向流水线是否能在保持同步训练与收敛性的同时，实现比单向流水线更高的吞吐量？
• RQ3 与现有流水线方法相比，Chimera 如何在加速器之间均衡激活内存消耗？
• RQ4 不同的调度策略（前向加倍、反向减半、直接拼接）对训练效率与可扩展性有何影响？
• RQ5 Chimera 在多大程度上可扩展至大批次大小与深层模型，同时保持最小的重计算开销？

主要发现

• 在 2,048 块 GPU 节点上，Chimera 对 13 亿参数的 GPT-2 模型，相比最先进同步与异步流水线方法，训练吞吐量提升 1.16 倍至 2.34 倍。
• 与现有同步流水线方法相比，Chimera 将流水线空洞减少了高达 50%，显著提升系统利用率。
• 采用直接拼接调度时，Chimera 在大批次下分别相较 GPipe、GEMS 和 DAPPLE 实现 1.13 倍、2.07 倍和 1.06 倍的加速。
• 在 512 块 GPU 上训练 GPT-2 时，采用前向加倍的 Chimera 分别相较 PipeDream-2BW、GPipe、GEMS 和 DAPPLE 实现 1.13 倍、1.18 倍、2.60 倍和 1.34 倍的加速。
• 四个流水线且 D=32 时，在 32 层 GPT-2 模型上达到峰值性能，但通常两个流水线能实现最高吞吐量，因在空洞与 allreduce 开销之间达到最佳平衡。
• Chimera 在保持与最先进方法相同峰值激活内存成本的同时，实现了更均衡的加速器间内存消耗。