TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models
Zhuohan Li Siyuan Zhuang Shiyuan Guo Danyang Zhuo Hao Zhang Dawn Song Ion Stoica
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
中文速览
训练GPT-3这类拥有数千亿参数的超大语言模型时，单张GPU根本装不下模型，必须把模型拆分到多台设备上并行计算，但现有方案要么设备间通信开销极大，要么因流水线"空泡"（bubble，即设备等待造成的闲置）而效率低下。TeraPipe发现了一个全新的并行维度：由于自回归语言模型（autoregressive language model）中每个词（token）的计算只依赖它之前的词，可以在单条训练序列内部按词切片、构建细粒度的流水线（token-level pipeline parallelism），让不同设备同时处理同一序列的不同词段与不同层，从而大幅减少空泡、提升设备利用率。为了找到最优的词段划分方式，研究者设计了一套基于动态规划（dynamic programming）的调度算法，自动根据模型结构和集群配置计算最佳切分方案。实验表明，在48台AWS p3.16xlarge实例组成的集群上训练1750亿参数的GPT-3最大版本时，TeraPipe比当前最优的模型并行方法快5倍，为超大规模语言模型的高效同步训练提供了一条新路径。
原文 arXiv:2102.07988;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2102.07988v2