PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers

By Chaoyang He, Shen Li, M. Soltanolkotabi, S. Avestimehr. 2021

38 citations recorded by the source corpus.

The size of Transformer models is growing at an unprecedented pace. It has only taken less than one year to reach trillion-level parameters after the release of GPT-3 (175B). Training such models requires both substantial engineering efforts and enormous computing resources, which are luxuries most research teams cannot afford. In this paper, we propose PipeTransformer, which leverages automated and elastic pipelining and data parallelism for efficient distributed training of Transformer models. PipeTransformer automatically adjusts the pipelining and data parallelism by identifying and freezing some layers during the training, and instead allocates resources for training of the remaining active layers. More specifically, PipeTransformer dynamically excludes converged layers from the pipeline, packs active layers into fewer GPUs, and forks more replicas to increase data-parallel width. We evaluate PipeTransformer using Vision Transformer (ViT) on ImageNet and BERT on GLUE and SQuAD datasets. Our results show that PipeTransformer attains a 2.4 fold speedup compared to the state-of-the-art baseline. We also provide various performance analyses for a more comprehensive understanding of our algorithmic and system-wise design. We also develop open-sourced flexible APIs for PipeTransformer, which offer a clean separation among the freeze algorithm, model definitions, and training accelerations, hence allowing it to be applied to other algorithms that require similar freezing strategies.

Subfields: Artificial Intelligence, Control and Systems Engineering, Mechanical Engineering, Optimization & Training

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

1. Bryan Catanzaro
2. Mohammad Shoeybi
3. Matei A. Zaharia
4. M. Patwary
5. D. Narayanan
6. Amar Phanishayee
7. Chien-chin Huang
8. Sam Shleifer

In Control and Systems Engineering

1. Bryan Catanzaro

In Mechanical Engineering

1. Myle Ott

In Optimization & Training

1. Mohammad Shoeybi
2. Bryan Catanzaro
3. M. Patwary
4. Matei A. Zaharia
5. Myle Ott
6. Chien-chin Huang
7. Min Xu
8. Xiaoxiao Li

Citations in the corpus

• Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, background): “PipeTransformer [19] elastically adjusts the degree of pipelining and data parallelism by freezing layers with “stable” weights, and instead dedicates resources to train the remaining “active” layers.”
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, background): “For example, sequence parallelism [12] reduces activation memory in conjunction with tensor parallelism; Pipetransformer [5] designed a dynamic 2D parallelism that allows changing the dimensions of pipeline and data parallelism on the fly, depending on learning signals.”
• Subgraph Federated Learning with Missing Neighbor Generation (2021, background): “(He et al., 2021a; Xie et al., 2021) design graph level FL schemes with graph datasets dispersed over multiple data owners, which are inapplicable to our distributed subgraph system construction.”