Pipeline parallelism is a distributed training strategy for deep neural networks that partitions a model's layers into sequential stages, each assigned to a different device (such as GPUs or TPUs). It combines model parallelism with data parallelism by processing micro-batches through the stages in a software-pipelined fashion, aiming to increase training throughput and reduce the idle time (bubbles) that naive layer-wise parallelism would incur.
Background and motivation
Training large neural networks, especially large language models and other deep learning models, requires enormous computational resources. As model sizes grow into the billions or trillions of parameters, single-device training becomes infeasible due to memory and compute limits. Data parallelism replicates the entire model on each device and splits the data, but it fails when the model exceeds a single device's memory. Model parallelism splits the model itself across devices, but naive layer-wise placement leads to severe underutilization: only one device computes at a time while others wait, resulting in low throughput and high idle time.
Pipeline parallelism addresses this by dividing the model into stages (typically contiguous groups of layers) and processing multiple micro-batches concurrently. The key insight is to overlap computation across stages: while stage 1 computes on micro-batch 2, stage 2 can compute on micro-batch 1, and so on. This resembles an assembly line, where each stage continuously processes incoming data and passes results to the next stage.
Core concepts
The fundamental unit of work in pipeline parallelism is the micro-batch. A large training batch is split into smaller micro-batches that flow through the pipeline stages sequentially. Each stage performs forward and backward passes on its assigned layers. The pipeline is "filled" by feeding micro-batches one after another, and after an initial warm-up period, all stages can be busy simultaneously, achieving high throughput.
The efficiency of a pipeline is often measured by the bubble ratio, which is the fraction of time that stages are idle. For a pipeline with p stages and m micro-batches, the ideal throughput is achieved when m is much larger than p, minimizing the relative bubble overhead. However, increasing m also increases memory usage for activations, creating a trade-off.
Major schemes
Several pipeline parallelism schemes have been proposed, differing in how they schedule forward and backward passes and how they handle gradient updates.
GPipe
Google introduced GPipe in 2019. It partitions the model into stages and uses a simple schedule: each micro-batch performs forward passes through all stages, then backward passes through all stages in reverse order. GPipe uses synchronous gradient updates, meaning all stages wait for all micro-batches to complete before updating weights. This ensures consistent gradients but introduces a bubble that grows with the number of stages. GPipe also requires either re-computation of activations or storing them, trading compute for memory.
PipeDream
PipeDream, from Microsoft Research (2019), uses an asynchronous schedule where each stage processes micro-batches in a round-robin fashion, performing forward and backward passes as soon as possible. This reduces the bubble but introduces weight staleness: different stages may use different versions of the weights, which can hurt convergence. PipeDream also uses a technique called "weight stashing" to maintain multiple versions of weights for different micro-batches, increasing memory overhead.
PipeDream-2BW and variants
PipeDream-2BW (2020) improves on PipeDream by using two weight buffers (one for forward, one for backward) to reduce memory and improve convergence. Variants like V-Pipe and PipeMare explore different synchronization strategies.
Megatron-LM and interleaved scheduling
NVIDIA's Megatron-LM (2019) combined tensor parallelism (splitting individual layers) with pipeline parallelism. Later, Megatron-2 (2020) introduced an interleaved schedule that divides the model into more stages than devices, allowing each device to handle multiple stages. This reduces the bubble size by increasing the number of micro-batches that can be in flight, at the cost of more communication.
Chimera and other recent schemes
Chimera (2021) uses a bidirectional pipeline to further reduce bubbles. Other schemes like PTD-P (2021) and ZeroBubble (2023) explore more sophisticated schedules to eliminate or reduce pipeline bubbles.
Mathematical and practical considerations
Pipeline parallelism is often combined with other parallelism strategies. Tensor parallelism splits individual layers across devices, while pipeline parallelism splits the model vertically. Data parallelism replicates the pipeline across groups of devices. This hybrid approach is standard in training large models, as seen in systems like Megatron-Turing NLG and OpenAI's GPT-4.
The choice of pipeline depth (number of stages) and micro-batch size affects performance. Deeper pipelines reduce per-stage memory but increase communication and bubble overhead. Larger micro-batches improve utilization but consume more memory. The optimal configuration depends on the model size, device memory, and interconnect bandwidth.
Applications and impact
Pipeline parallelism has been instrumental in training some of the largest AI models. For example, GPT-3 (175B parameters) was trained using a combination of model and pipeline parallelism. Anthropic's Claude models and Google's PaLM also rely on similar techniques. The approach is supported by major frameworks, including PyTorch (through the torch.distributed.pipeline module), TensorFlow (via tf.distribute), and JAX (with pjit).
Challenges and limitations
Pipeline parallelism introduces several challenges. Communication overhead arises from sending activations and gradients between stages, which can become a bottleneck on slow interconnects. Load imbalance occurs when stages have unequal compute times, reducing efficiency. Memory pressure from storing activations for backward passes can be mitigated by activation re-computation but at a compute cost. Additionally, asynchronous schemes can suffer from convergence issues due to stale weights, while synchronous schemes incur bubbles.
Future directions
Research continues on improving pipeline parallelism, focusing on reducing bubbles, minimizing memory, and adapting to heterogeneous hardware. Techniques like automatic stage partitioning, dynamic load balancing, and integration with mixture-of-experts models are active areas. The rise of extremely large models and the need for efficient training on cloud clusters will likely keep pipeline parallelism a key component of distributed training systems.