Source-linked AI summary
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
TL;DR
Large language models exceed individual accelerator memory, motivating model-parallel training, but existing methods face communication overhead or pipeline inefficiency. TeraPipe pipelines autoregressive Transformer computation along the token dimension and uses dynamic programming to select the execution scheme. It reports a 5.0x speedup for synchronous training of the 175-billion-parameter GPT-3 model on 48 AWS p3.16xlarge instances.
Problem
175-billion-parameter language models can exceed individual accelerator memory, while existing model-parallel methods face communication overhead or pipeline inefficiency.
Method
TeraPipe performs synchronous token-level pipeline parallelism and uses dynamic programming to compute an optimal token partition for a given model and cluster.
Results
5.0x speedup was reported for synchronous training of the 175-billion-parameter GPT-3 model on an AWS cluster with 48 p3.16xlarge instances.
Takeaways & Limitations
Token-level pipelining provides a fine-grained, model-parallel training approach that is orthogonal to other model-parallel methods and can be combined with them.
Takeaways & Limitations
The paper focuses on unidirectional autoregressive language models, not bidirectional masked language models, and evaluates optimization through per-iteration latency under the same training algorithm.
Abstract
from arXiv · showhide
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
1. Introduction
Large Transformer language models require model parallelism because their parameter memory exceeds individual accelerators, while existing approaches incur communication overhead or pipeline bubbles. TeraPipe exploits autoregressive token dependencies to pipeline within one sequence, using dynamic programming to optimize token partitioning and achieving a reported 5.0x speedup on the largest GPT-3 model.
- Motivation: 175B-parameter models can exceed accelerator memory, making model-parallel training necessary.GPT-3 can have more than 175 billion parameters, amounting to 350 GB at 16-bit precision.
- Motivation: Existing model-parallel approaches either introduce excessive inter-device communication or lose efficiency through pipeline bubbles.Operation partitioning requires synchronization, while microbatch pipelining can leave devices idle.
- Key Insight: Autoregressive Transformers allow token-level pipelining because each token depends on previous tokens but not future tokens.This enables current-token computation on one layer to overlap with previous-token computation on the next layer.
- Method: TeraPipe introduces token dimension as a new pipeline-parallel dimension for Transformer language models.The method focuses on unidirectional autoregressive models such as GPT, not bidirectional masked language models such as BERT.
- Method: A dynamic programming algorithm computes a token partition intended to maximize pipeline parallelism for a specific model and cluster configuration.The partitioning addresses under-utilization from overly fine-grained slices and uneven loads across token positions.
- Results: 5.0x higher synchronous training throughput was reported for the 175-billion-parameter GPT-3 model versus previous state-of-the-art model-parallel methods.The evaluation used an AWS cluster with 48 p3.16xlarge instances.
2. Related Work
Prior model-parallel methods partition operations or pipeline microbatches, but communication costs and pipeline bubbles limit efficiency. TeraPipe’s token-level approach is distinguished from both microbatch pipelining and wavefront methods by exploiting Transformer-specific dependencies.
- Model Parallelism: Model parallelism partitions a model across devices so models larger than one device’s memory can be trained.Existing approaches are broadly categorized as operation partitioning and pipeline parallelism.
- Operation Partitioning: Operation partitioning parallelizes matrix computations across devices but requires cross-device communication to combine partial results.Partitioning matrices along rows and columns produces partial products that must be summed.
- Microbatch Pipelining: Microbatch pipeline parallelism overlaps different microbatches across layers, but its pipeline bubbles grow as model size increases.Longer sequences reduce feasible minibatch sizes under fixed GPU memory, leaving fewer sequences available for parallel processing.
- Wavefront Parallelism: Wavefront parallelism accelerates multilayer RNNs but does not accelerate Transformers because positions within a Transformer layer lack the needed dependency.Its per-word granularity would also be too fine for efficient TeraPipe pipelining.
3. Method
TeraPipe exploits autoregressive dependencies to pipeline token slices within a single sequence, while using dynamic programming to choose slices that balance device utilization and pipeline efficiency. It is designed to complement existing model-parallel techniques.
- Pipeline parallelism within a sequence: Previous microbatch pipelines become less efficient for large models because longer sequences and smaller batch sizes create larger pipeline bubbles.Large sequence lengths reduce the number of sequences processed concurrently, while smaller batches limit bubble saturation.
- Pipeline parallelism within a sequence: Transformer autoregression lets computation at layer f_i and token t begin after earlier hidden states from f_i−1 are ready, enabling within-sequence pipelining.Self-attention uses previous positions, while the feed-forward layer uses the current position.
- Pipeline parallelism within a sequence: Token-level slicing must balance GPU underutilization from short slices against larger pipeline bubbles from long slices.On a GPT3-1B layer, sequences shorter than 256 tokens have the same forward time as a single-token sequence, while longer slices reduce the number of pipeline stages.
- Selecting optimal slicing scheme: TeraPipe uses dynamic programming to select slice lengths l_1,...,l_M summing to L and minimize total forward and backward propagation latency.The algorithm searches for an optimal token-dimension slicing scheme for a partitioned Transformer model.
- Selecting optimal slicing scheme: Its performance model separates forward time without extra context from context overhead and predicts the latter with a linear model having under 2% relative error.The model uses measured values for all choices of the slice endpoint and fits the context term from a subset of combinations.
- Combining with other parallel training methods: TeraPipe’s token pipeline is orthogonal to previous model-parallel techniques and can be combined with them to improve parallelization performance.The method preserves synchronous model-parallel training while adding token-dimension pipelining.
4. Evaluation
The evaluation measures per-iteration latency across GPT-3 model sizes and sequence lengths, comparing TeraPipe with prior model-parallel training. TeraPipe consistently accelerates larger models and increasingly long sequences, while dynamic-programming slicing outperforms uniform slicing.
- TeraPipe preserves the single-device optimization algorithm, so evaluation focuses on per-iteration wall-clock latency.
- Experiments evaluate GPT3-1B, GPT3-13B, GPT3-44B, and GPT3-175B configurations on AWS p3.16xlarge clusters.The tested models contain 1B, 13B, 44B, and 175B parameters, respectively.
- TeraPipe accelerates GPT3-13B by 1.40x and GPT3-44B by up to 2.40x across the reported settings.
- 6.75x and 5.02x speedups are achieved for GPT3-175B settings (9) and (10), respectively.
- TeraPipe provides higher speedups for larger models because smaller batch sizes weaken microbatch pipeline saturation while additional stages create more opportunities for token slicing.
- The dynamic-programming slicing scheme is 1.12x and 1.04x better than the best uniform slicing pipeline in the evaluated GPT3-44B and GPT3-175B cases.Very fine-grained slicing underutilizes GPUs, while coarse slicing creates large pipeline bubbles.
- For GPT3-13B, TeraPipe reaches 2.76x, 4.97x, and 7.83x speedups at sequence lengths 4096, 6144, and 8196, respectively.The performance gap grows with sequence length, although batch sizes are reduced because memory usage increases.
5. Conclusion
TeraPipe combines token-level pipeline parallelism with dynamic programming to optimize execution for a given model and cluster, achieving a reported 5.0x speedup for 175-billion-parameter GPT-3 models.
- TeraPipe is a token-level pipeline parallel algorithm for training large-scale Transformer language models.
- Its dynamic programming algorithm calculates the optimal pipelining execution scheme for a specific language model and cluster configuration.
- TeraPipe is orthogonal to other model parallel training methods and can be combined with them.
- 5.0x acceleration was reported for synchronous training of 175-billion-parameter GPT-3 models on 48 AWS p3.16xlarge instances versus previous methods.
A. Combine TeraPipe with Gradient Accumulation
TeraPipe can complement gradient accumulation, but memory limits can prevent gradient accumulation from filling a deep pipeline when each GPU stores few sequences.
- TeraPipe and gradient accumulation are orthogonal, and TeraPipe can provide additional speedup over gradient accumulation.
- When each GPU supports only two input sequences, sequence 3 cannot begin its forward pass until sequence 1 releases backward-pass activations.
- With 48 pipeline stages, GPU memory is insufficient to hold 48 input sequences, even on 80GB A100 GPUs.
- Consequently, TeraPipe is expected to improve training efficiency even when gradient accumulation is used.
B. Implementation
The implementation uses PyTorch and NCCL, builds on Megatron-LM for operation partitioning, and adds custom pipeline and data parallelism components.
- TeraPipe is implemented with PyTorch and NCCL.
- Megatron-LM provides operation partitioning, while microbatch-based pipeline parallelism and data parallelism are implemented separately.
- The core implementation contains 1714 lines of Python and is included in the supplementary material for open-source release.
C. Experiment Results
The supplementary experiment section provides detailed latency statistics and slicing schemes for the main experiments, dynamic-programming ablations, and longer-sequence experiments.
- Detailed latency means, standard deviations, and dynamic-programming slicing schemes are provided for all main-paper experiments.
- Table 2 documents detailed numbers and slicing schemes for the main experiments corresponding to Figure 5.
- Table 3 documents the dynamic programming algorithm ablation studies corresponding to Figure 6.
- Table 4 documents experiments with longer sequence lengths corresponding to Figure 7.