TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models

By Zhuohan Li, Siyuan Zhuang, Shiyuan Guo, Danyang Zhuo, Dawn Song, Ion Stoica. 2021

166 citations recorded by the source corpus.

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

Subfields: Artificial Intelligence, Model Architectures, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Model Architectures

In Efficiency & Compression

Citations in the corpus

Ethical and social risks of harm from Language Models (2021, background): “Net impact of efficiency gains is difficult to predict Work to reduce the wall-clock time required to train a LM (Li et al., 2021) can yield efficiency gains and reduce the environmental cost of training a model.”
GSPMD: General and Scalable Parallelization for ML Computation Graphs (2021, methodology): “However, it is limited to homogeneous pipeline stages; for heterogeneous stages we recommend integrating GSPMDwith other pipeline implementations [15, 22, 24] and sharding each stage.”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, background): “TeraPipe [20] exposes fine-grained pipeline parallelism across tokens in a single training sequence for auto-regressive models like GPT.”
Reducing Activation Recomputation in Large Transformer Models (2022, methodology): “Two forms of model parallelism are commonly used to distribute the model parameters across GPUs: parameters of each layer are distributed across many devices [18, 19, 21], and 2) pipeline parallelism where the model is split along the layer dimension of the network [7, 10, 11].”
PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, background): “Pipeline parallelism [5, 7, 10, 14, 19] partitions a model instance into stages and distributes stages across multiple devices, where activations and gradients are communicated across stage boundaries.”
Pre-Trained Models: Past, Present and Future (2021, methodology): “…pipeline methods include GPipe (Huang et al., 2019b) which can send smaller parts of samples within a mini-batch to different nodes, and Tera Pipe (Li et al., 2021) which can apply token-level pipeline mechanisms for Transformer-based models to make each token in a sequence be processed by…”
Taxonomy of Risks posed by Language Models (2022, background): “However LM fine-tuning with differential privacy2 has been limited to small models [116, 201] and it remains to be established whether it is suitable for training LMs from scratch on large datasets of web text [188].”
This paper is included in the Proceedings of the 16th USENIX Symposium on Operating Systems Design and Implementation. (unknown): “TeraPipe [33] discovers a new pipeline parallelism dimension for transformer-based LMs. Google’s Pathway system [7] is a concurrent work of Alpa.”
Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training (2021, background): “TeraPipe (Li et al., 2021b) introduced another type of pipelining speciﬁc to single-transformer architectures, where pipelining occurs across tokens instead of microbatches.”
Parallel Training of Pre-Trained Models via Chunk-Based Dynamic Memory Management (2021, background): “Pipeline Parallelism (PP) [14, 22, 26] is another approach for distributed training that divides the model’s layers into stages that the model can process in a pipeline manner.”
ERNIE 3.0 Titan: Exploring Larger-scale Knowledge Enhanced Pre-training for Language Understanding and Generation (2021, methodology): “The 4D hybrid parallelism includes data parallelism (DP) [60, 61], intra-layer tensor model parallelism (MP) [62, 57], inter-layer pipeline model parallelism (PP) [32, 63, 64, 65, 33], and the improved sharded data parallelism (Sharded) based on ZeRO [58].”
Large Language Model (LLM) for Telecommunications: A Comprehensive Survey on Principles, Key Techniques, and Opportunities (2024, methodology): “Data parallelism involves the replication of the model’s parameters and optimizer states across multiple GPUs [57], allocating to each GPU a subset of data to process and subsequently aggregate the computed gradients.”