[Paper Review] TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models

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

Topic Modeling31 references33 citations

TL;DR

TeraPipe introduces token-level pipeline parallelism for Transformer LMs and achieves up to 5.0x training speedups on GPT-3-175B over prior synchronous model-parallel methods on AWS.

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

Motivation & Objective

• Motivate the need for deeper model parallelism to train extremely large Transformer LMs beyond single-device memory limits.
• Identify a new, fine-grained pipeline dimension along the token sequence that leverages autoregressive dependencies.
• Develop a dynamic programming-based algorithm to compute optimal token-slice partitioning for maximal pipeline efficiency.

Proposed method

• Propose token-level pipeline parallelism that pipelines across the token dimension within a single input sequence.
• Model forward/backward latency as a function of token-slice sizes and cluster characteristics.
• Develop a dynamic programming algorithm to find the optimal slicing scheme over the token dimension to minimize training latency.
• Estimate forward propagation time with a simple performance model and use it to guide DP optimization.
• Show orthogonality: TeraPipe can be combined with existing data/model parallel methods (microbatching, operation partitioning, data parallelism).

Experimental results

Research questions

• RQ1 How can pipeline parallelism be extended from the layer dimension to the token dimension in autoregressive Transformers?
• RQ2 What slicing scheme over the token dimension minimizes total training latency for a given LM and cluster?
• RQ3 How does token-level pipelining interact with other model-parallel techniques and data parallelism?
• RQ4 What performance gains can be achieved on large GPT-3-scale models using token-level pipeline parallelism?
• RQ5 How does sequence length affect the effectiveness of token-level pipeline parallelism?

Key findings

• TeraPipe yields substantial speedups for large LMs, with up to 5.0x faster training for GPT-3-175B over prior synchronous model-parallel methods on 48 AWS p3.16xlarge GPUs.
• A dynamic programming approach effectively determines optimal token-slicing schemes to maximize pipeline efficiency, outperforming uniform slicing by about 1.04x–1.12x in examined cases.
• The method provides larger gains for bigger models due to memory constraints reducing batch size and increasing pipeline stages, where token-level pipelining offers more saturation opportunities.
• Longer input sequence lengths substantially boost the potential benefits of token-level pipelining, with observed speedups increasing as sequence length grows.
• TeraPipe is orthogonal to, and can be combined with, existing parallel training methods such as microbatch-based pipeline, operation partitioning, and data parallelism.