Pipeline Parallelism

The core mechanism of pipeline parallelism is the sequential partitioning of a model's computational graph. Layers are grouped into contiguous microbatches and assigned to different devices (e.g., GPUs). For a model with N layers and P devices, each device is responsible for approximately N/P consecutive layers. This creates a deterministic processing order where the output of one device's stage becomes the input for the next, forming a compute pipeline. This is distinct from tensor model parallelism, which splits individual layers.

A fundamental challenge is the pipeline bubble—idle time where devices wait for data from preceding stages. To minimize this, the technique of microbatch interleaving is used. Instead of processing one full batch at a time, the batch is split into many smaller microbatches. These are fed into the pipeline in a staggered fashion, keeping all devices busy concurrently. Common scheduling patterns include:

• GPipe (1F1B): A steady-state schedule of 1 Forward pass followed by 1 Backward pass per device.
• Interleaved 1F1B: Further splits stages, allowing more frequent interleaving and better bubble reduction at the cost of more communication.

Pipeline parallelism excels at maximizing throughput (samples processed per second) for batch inference workloads. Once the pipeline is full (the warm-up phase), the system can process multiple microbatches simultaneously at different pipeline stages. This makes it highly effective for offline processing, large-scale batch scoring, and serving scenarios where request queues can be grouped. Its efficiency scales with batch size, unlike data parallelism which is limited by per-device memory for a single input.

This technique directly addresses GPU memory constraints. By distributing layers, each device only needs to store the parameters, gradients, and activations for its assigned segment of the model. This allows for serving models that are far larger than the memory of any single accelerator. The memory footprint per device is roughly Total Model Size / Number of Pipeline Stages. However, it requires storing activation checkpoints for the backward pass, trading some compute for memory savings.

Communication is a defining characteristic. Devices only communicate with their immediate pipeline neighbors (previous and next stage). This creates a predictable, point-to-point communication pattern, typically using high-bandwidth links like NVLink. The primary overhead is the latency of sending activations and gradients between stages. The efficiency of the pipeline is highly sensitive to this communication latency and the balance of compute time across stages (load balancing). An imbalanced pipeline will be dominated by the slowest stage.

Pipeline parallelism is one of the three core distributed training strategies, each with distinct characteristics:

• Data Parallelism: Replicates the entire model on each device and splits the data batch. Excellent for small-to-medium models but hits memory limits for large ones.
• Tensor/Model Parallelism: Splits individual layers (e.g., matrix multiplications) across devices. Minimizes pipeline bubbles but incurs very high communication costs for every layer.
• Pipeline Parallelism: Splits layers between devices. Has lower communication volume than tensor parallelism but introduces pipeline bubbles. In practice, modern systems like Megatron-LM use 3D parallelism, combining all three techniques to train trillion-parameter models.

PyTorch provides native pipeline parallelism support through its torch.distributed.pipelining module (and the legacy torch.distributed.pipeline.sync.Pipe). It allows model partitioning using torch.fx symbolic tracing and handles automatic splitting of micro-batches, gradient accumulation, and backward pass scheduling.

• Core API: Pipe class wraps a torch.nn.Sequential module split across devices.
• Scheduling: Implements the GPipe schedule (1F1B for interleaved).
• Use Case: Primarily for training very large models (e.g., >100B parameters) where a single device lacks sufficient memory.

Microsoft's DeepSpeed library includes a sophisticated pipeline parallelism engine as part of its ZeRO-3 optimization suite. It is designed for extreme-scale model training and supports 3D parallelism (combining data, tensor, and pipeline parallelism).

• Pipeline Engine: Manages the 1F1B (One-Forward-One-Backward) schedule with interleaved stages for better GPU utilization.
• Zero Bubble: Aims to minimize the pipeline "bubble" where devices are idle.
• Integration: Works seamlessly with Hugging Face Transformers and Megatron-LM for training models like GPT-3 and BLOOM.

NVIDIA's Megatron-LM framework is a cornerstone for training large transformer models. It combines pipeline parallelism with tensor parallelism (intra-layer model parallelism) and data parallelism for optimal scaling.

• Hybrid Parallelism: A single model layer may be split across 8 GPUs using tensor parallelism, while the sequence of layers is pipelined across many more devices.
• Communication Optimization: Uses efficient NCCL primitives for gradient synchronization across pipeline stages.
• Production Use: The standard for training models like MT-NLG (530B parameters) and custom enterprise LLMs.

The Alpa project (and underlying JAX ecosystem) automates the parallelization of large models. It treats pipeline parallelism as one strategy within a unified compiler pass that can also decide on data and operator parallelism.

• Automated Orchestration: Users provide a single-device model; Alpa's compiler generates an execution plan and necessary communication code.
• Inter-Operator Parallelism: Alpa's pipeline parallel strategy can partition at the granularity of individual operators, not just layers.
• Future Direction: Represents the move towards compiler-managed parallelism, reducing manual partitioning effort.

For batch inference scenarios, specialized servers use pipeline parallelism to maximize throughput, not for training. They treat the model pipeline as a producer-consumer system.

• Triton Inference Server: Supports ensemble models, where different stages can be placed on different GPU/CPU devices, forming an inference pipeline.
• Custom Orchestration: High-throughput serving systems for recommendation models or large vision transformers often implement a pipeline where preprocessing, model execution, and postprocessing are distinct, parallelizable stages.
• Key Difference: Inference pipelines often avoid backward passes and complex gradient synchronization, simplifying scheduling.

Implementing pipeline parallelism introduces distinct systems challenges that frameworks must address:

• Pipeline Bubble: The idle time at the beginning and end of a batch as the pipeline fills and drains. Schedules like 1F1B with interleaving are designed to reduce this.
• Checkpointing: For fault tolerance, frameworks must implement activation recomputation (selective re-forwarding) or activation checkpointing to manage memory without sacrificing throughput.
• Load Balancing: Achieving balanced computation time across stages is critical. Imbalanced partitions create bottlenecks, limiting the speedup of the entire pipeline.
• Communication Overhead: The bandwidth and latency between devices (e.g., NVLink vs. PCIe) directly limit the practical number of pipeline stages and micro-batch size.

The overarching family of techniques for partitioning a single model across multiple devices. Pipeline parallelism is a specific type of model parallelism where layers are split sequentially. Other types include:

• Tensor Parallelism: Splits individual layers (e.g., the weight matrices of an attention head) across devices, requiring high-bandwidth communication.
• Expert Parallelism: Used in Mixture-of-Experts models, where different expert sub-networks are placed on different devices. The choice depends on the model architecture and the cluster's communication topology.

A complementary technique where the same model is replicated across multiple devices, and each device processes a different subset of the input data (a mini-batch). Gradients are synchronized across devices after each forward/backward pass. While pipeline parallelism splits the model, data parallelism replicates it. They are often combined in large-scale training as 3D parallelism (data, pipeline, and tensor). For inference, data parallelism is used for horizontal scaling to increase request throughput.

A critical inference optimization for improving GPU utilization when serving LLMs. Unlike static batching, which waits for a full batch, continuous batching dynamically groups incoming requests of varying sequence lengths and schedules them for execution as soon as resources are free. It is highly synergistic with pipeline parallelism:

• The pipeline's micro-batches are perfect units for continuous batching.
• Together, they maximize throughput by keeping all pipeline stages consistently busy, even with irregular request arrival times.

A model architecture where only a sparse subset of neural network components (the 'experts') are activated for a given input. Inference for MoE models involves two parallelisms:

1. Expert Parallelism: Distributing different experts across devices.
2. Pipeline Parallelism: Often used to handle the sequential non-expert layers (like attention blocks) within the model. Serving MoE models at scale requires sophisticated routing and scheduling to manage the conditional execution flow across the distributed hardware.

The overarching process of deploying and executing models in production. Pipeline parallelism is a serving architecture decision made to enable the serving of models too large for a single device. It is implemented within inference servers like NVIDIA Triton or vLLM, which manage the lifecycle, scheduling, and API exposure for pipelined models. The serving layer handles client requests, forms batches, and manages the data movement between the pipeline stages.

The performance of pipeline parallelism is bounded by the communication links between devices. Key technologies include:

• NVLink/NVSwitch: High-bandwidth, low-latency interconnects between GPUs in the same node (e.g., up to 900 GB/s with NVLink 4).
• GPUDirect RDMA: Allows direct memory access between GPUs in different servers over InfiniBand or Ethernet, bypassing the host CPU. The pipeline bubble is directly influenced by the latency of transferring activations and gradients between stages.