[论文解读] PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers

Chaoyang He, Shen Li+2 more|arXiv (Cornell University)|Feb 5, 2021

Neural Networks and Applications参考文献 22被引用 16

一句话总结

PipeTransformer 提出了一种用于分布式训练 Transformer 模型的自动化、弹性流水线框架，该框架在训练过程中动态冻结层，将活跃层打包到更少的 GPU 上，并通过启动额外的流水线副本增加数据并行宽度。该方法在 ViT 和 BERT 模型上实现了高达 2.83 倍的速度提升，且未造成精度损失。

ABSTRACT

The size of Transformer models is growing at an unprecedented pace. It has only taken less than one year to reach trillion-level parameters after the release of GPT-3 (175B). Training such models requires both substantial engineering efforts and enormous computing resources, which are luxuries most research teams cannot afford. In this paper, we propose PipeTransformer, which leverages automated and elastic pipelining and data parallelism for efficient distributed training of Transformer models. PipeTransformer automatically adjusts the pipelining and data parallelism by identifying and freezing some layers during the training, and instead allocates resources for training of the remaining active layers. More specifically, PipeTransformer dynamically excludes converged layers from the pipeline, packs active layers into fewer GPUs, and forks more replicas to increase data-parallel width. We evaluate PipeTransformer using Vision Transformer (ViT) on ImageNet and BERT on GLUE and SQuAD datasets. Our results show that PipeTransformer attains a 2.4 fold speedup compared to the state-of-the-art baseline. We also provide various performance analyses for a more comprehensive understanding of our algorithmic and system-wise design. We also develop open-sourced flexible APIs for PipeTransformer, which offer a clean separation among the freeze algorithm, model definitions, and training accelerations, hence allowing it to be applied to other algorithms that require similar freezing strategies.

研究动机与目标

• 为解决大规模 Transformer 模型训练带来的高计算与工程成本问题。
• 通过利用深层网络层自底向上收敛的观察，减少训练时间和资源消耗。
• 设计一种能够根据冻结层动态调整流水线配置和数据并行复制的系统。
• 克服在运行时冻结、动态流水线、变化进程组中的集体通信以及跨副本激活缓存等挑战。

提出的方法

• 提出一种自适应冻结算法，利用可调阈值 α 来根据训练过程中激活值的稳定性识别并冻结层。
• 采用 AutoPipe 重新划分模型，将活跃层打包到更少的 GPU 上，以最小化通信开销和流水线空洞。
• 使用 AutoDP 在释放的 GPU 上动态启动新的流水线副本，同时保持分层通信组以实现高效的集体操作。
• 实现 AutoCache，使现有和新增的流水线副本能够共享缓存的激活值，其性能分析器用于确定最优的激活缓存启动时间。
• 将这些组件整合为一个整体系统，根据实时层冻结情况自动调整模型划分与复制策略。
• 利用 PyTorch 的流水线和数据并行 API 确保兼容性与高效执行。

实验结果

研究问题

• RQ1 在训练过程中是否可以有效利用动态层冻结来加速 Transformer 的分布式训练？
• RQ2 如何在层冻结的背景下联合优化流水线与数据并行性？
• RQ3 在弹性流水线中，冻结阈值 α 和微批次分块的性能权衡如何？
• RQ4 在一个流水线副本数量动态变化的系统中，如何高效管理集体通信？
• RQ5 在动态训练设置中，激活缓存的最佳启用时机是什么？

主要发现

• PipeTransformer 在 ViT（ImageNet）和 BERT（SQuAD/GLUE）上相较于最先进基线实现了高达 2.83 倍的速度提升，且未造成精度下降。
• 较大的冻结阈值 α（如 α = 1/5）可带来更大的加速效果，但会造成轻微精度损失，表明速度与性能之间存在可调和的权衡。
• 最优微批次数量 M 随流水线长度 K 显著变化，凸显了在避免性能下降前进行预分析的重要性。
• 仅在足够多的层被冻结后，激活值缓存才变得有益——AutoCache 在 ViT 上于 3 个层冻结后开始缓存，在 BERT 上于 5 个层冻结后开始缓存。
• 在 ImageNet 上通信开销占总训练时间的 5.9%，在 SQuAD 上占 8.8%，表明进一步减小模型尺寸可显著提升效率。
• 减少流水线空洞与增加数据并行宽度的乘法效应带来了显著的、非叠加性的加速效果。