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

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

arXiv commentICML 2021

References in corpus (16)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Language Models are Few-Shot Learners
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• One weird trick for parallelizing convolutional neural networks
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Training Deep Nets with Sublinear Memory Cost
• Big Bird: Transformers for Longer Sequences
• Long Range Arena: A Benchmark for Efficient Transformers
• Supporting Very Large Models using Automatic Dataflow Graph Partitioning
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• Optimizing Performance of Recurrent Neural Networks on GPUs
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• ZeRO-Offload: Democratizing Billion-Scale Model Training
• Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks
• Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
• DAPPLE: A Pipelined Data Parallel Approach for Training Large Models

Cited by in corpus (6)

• GSPMD: General and Scalable Parallelization for ML Computation Graphs
• PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
• Pre-Trained Models: Past, Present and Future
• Pipeline Parallelism for Inference on Heterogeneous Edge Computing
• AutoDDL: Automatic Distributed Deep Learning with Near-Optimal Bandwidth Cost
• Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training