[论文解读] TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models

Zhuohan Li, Siyuan Zhuang+5 more|arXiv (Cornell University)|Feb 16, 2021

Topic Modeling参考文献 31被引用 33

一句话总结

TeraPipe 引入面向令牌的流水线并行化，用于 Transformer 语言模型，在 AWS 上的 GPT-3-175B 相对于先前的同步模型并行方法实现最高 5.0x 的训练加速。

ABSTRACT

Model parallelism has become a necessity for training modern large-scale deep language models. In this work, we identify a new and orthogonal dimension from existing model parallel approaches: it is possible to perform pipeline parallelism within a single training sequence for Transformer-based language models thanks to its autoregressive property. This enables a more fine-grained pipeline compared with previous work. With this key idea, we design TeraPipe, a high-performance token-level pipeline parallel algorithm for synchronous model-parallel training of Transformer-based language models. We develop a novel dynamic programming-based algorithm to calculate the optimal pipelining execution scheme given a specific model and cluster configuration. We show that TeraPipe can speed up the training by 5.0x for the largest GPT-3 model with 175 billion parameters on an AWS cluster with 48 p3.16xlarge instances compared with state-of-the-art model-parallel methods. The code for reproduction can be found at https://github.com/zhuohan123/terapipe

研究动机与目标

• 推动需要更深层的模型并行性，以在单设备内存限制之外训练极大型的 Transformer 语言模型。
• 确定沿令牌序列的新颖、细粒度的流水线维度，利用自回归依赖关系。
• 开发基于动态规划的算法，以计算最大化流水线效率的最优令牌切片分区。

提出的方法

• 提出在单个输入序列的令牌维度上进行流水线的令牌级流水线并行化。
• 将前向/后向延迟建模为令牌切片大小和集群特征的函数。
• 开发一个动态规划算法，在令牌维度上找到最优切片方案以最小化训练延迟。
• 使用一个简单的性能模型估算前向传播时间，并用它来指导 DP 优化。
• 展示正交性：TeraPipe 可以与现有的数据/模型并行方法（基于微批的流水线、操作切分、数据并行等）结合使用。

实验结果

研究问题

• RQ1 如何将流水线并行性从层维度扩展到自回归 Transformer 的令牌维度？
• RQ2 对于给定的语言模型和集群，令牌维度上的哪种切片方案能最小化总训练延迟？
• RQ3 令牌级流水线与其他模型并行技术和数据并行如何交互？
• RQ4 在大型 GPT-3 规模的模型上使用令牌级流水线并行化可以实现哪些性能增益？
• RQ5 序列长度如何影响令牌级流水线并行性的有效性？

主要发现

• TeraPipe 在大型语言模型上带来显著的加速，对于 GPT-3-175B，在 48 块 AWS p3.16xlarge GPU 上，相较于此前的同步模型并行方法，训练时间快多达 5.0x。
• 动态规划方法有效确定最大化流水线效率的最优令牌切片方案，在所研究的案例中比均匀切片高出约 1.04x–1.12x。
• 该方法在更大模型上带来更大收益，因为内存限制会降低批量大小并增加流水线阶段，令牌级流水线提供了更多的饱和机会。
• 更长的输入序列长度显著提升令牌级流水线化的潜在收益，观测到的加速随着序列长度增长而增加。
• TeraPipe 与现有的并行训练方法（如基于微批的流水线、操作切分和数据并行）正交且可以与之结合。