PipeDream: generalized pipeline parallelism for DNN training

By D. Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, G. R. Ganger, Phillip B. Gibbons, Matei A. Zaharia. 2019

1,228 citations recorded by the source corpus.

DNN training is extremely time-consuming, necessitating efficient multi-accelerator parallelization. Current approaches to parallelizing training primarily use intra-batch parallelization, where a single iteration of training is split over the available workers, but suffer from diminishing returns at higher worker counts. We present PipeDream, a system that adds inter-batch pipelining to intra-batch parallelism to further improve parallel training throughput, helping to better overlap computation with communication and reduce the amount of communication when possible. Unlike traditional pipelining, DNN training is bi-directional, where a forward pass through the computation graph is followed by a backward pass that uses state and intermediate data computed during the forward pass. Naïve pipelining can thus result in mismatches in state versions used in the forward and backward passes, or excessive pipeline flushes and lower hardware efficiency. To address these challenges, PipeDream versions model parameters for numerically correct gradient computations, and schedules forward and backward passes of different minibatches concurrently on different workers with minimal pipeline stalls. PipeDream also automatically partitions DNN layers among workers to balance work and minimize communication. Extensive experimentation with a range of DNN tasks, models, and hardware configurations shows that PipeDream trains models to high accuracy up to 5.3X faster than commonly used intra-batch parallelism techniques.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Electrical and Electronic Engineering, Efficiency & Compression, 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

In Computer Vision and Pattern Recognition

In Electrical and Electronic Engineering

In Efficiency & Compression

In Optimization & Training

Citations in the corpus

Unified Scaling Laws for Routed Language Models (2022, methodology): “There are still more variations of model parallelism, notably Pipeline Parallelism [Narayanan et al., 2019, Huang et al., 2019], which we do not use.”
Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model (2022, methodology): “We use a 1F1B pipeline schedule [42] that alternates forward and backward propagations.”
ZeRO: Memory optimizations Toward Training Trillion Parameter Models (2019, methodology): “A different implementation of PP in PipeDream [14] keeps multiple copies of stale parameters”
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.”
PanGu-α: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation (2021, methodology): “The second regime is pipeline model parallelism [24, 25, 26, 27, 28], which partitions the total layers to stages, and then places stages to different devices, as shown in Figure 4(c).”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, methodology): “Asynchronous and bounded-staleness approaches such as PipeMare, PipeDream, and PipeDream-2BW [23, 29, 30, 45] do away with flushes completely, but relax weight update semantics.”
Memory-Efficient Pipeline-Parallel DNN Training (2020, methodology): “The number of byes communicated during an all-reduce can either be explicitly measured, or estimated using a closed-form expression [16].”
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.”
Mixture-of-Experts with Expert Choice Routing (2022, background): “Model parallelism [30] splits weights and tensors across multiple cores while pipeline parallelism [18, 24] splits different layers across devices with micro-batches pipelined to the different layers.”
Pre-Trained Models: Past, Present and Future (2021, background): “One has to carefully manage the data movement between devices, which may otherwise become the performance bottleneck (Narayanan et al., 2019; Jiang et al., 2020a).”
DAPPLE: a pipelined data parallel approach for training large models (2020, methodology): “Recently, pipeline parallelism[24, 36, 52] has been proposed as a promising approach for training large DNN models.”
Wavelet: Efficient DNN Training with Tick-Tock Scheduling (2021, methodology): “Recent literature (Huang et al., 2019; Narayanan et al., 2019) leverages inter-batch pipelining to improve GPU utilization in model parallel training.”
LongVILA: Scaling Long-Context Visual Language Models for Long Videos (2024, methodology): “The most straight-forward and user-friendly solution within HuggingFace Transformers is to use its pipeline parallelism inference feature, which shards a single model across multiple devices on a layer-by-layer basis (Huang et al., 2019; Narayanan et al., 2019).”
Understanding Capacity-Driven Scale-Out Neural Recommendation Inference (2020, background): “…for model-parallelism in TensorFlow–and in turn huge models–but have focused on language models and training which don’t have the variety and number of features shown in Figure 5 and emphasize training throughput compared to inference’s latency-bounded throughput [15], [36], [37], [40], [41].”
DeepSeek-Coder: When the Large Language Model Meets Programming - The Rise of Code Intelligence (2024, methodology): “These include tensor parallelism (Korthikanti et al., 2023), alongside ZeRO data parallelism (Rajbhandari et al., 2020) and PipeDream pipeline parallelism (Narayanan et al., 2019 utilize clusters outfitted with NVIDIA A100 and H800 GPUs.”