Pipeline Parallelism

The fundamental mechanism that enables high device utilization. The pipeline is filled with multiple micro-batches of data simultaneously, each at a different stage. This creates a steady state where all devices are active concurrently, overlapping computation across the pipeline. For example, while Device 2 processes the forward pass for micro-batch N, Device 1 can process the backward pass for micro-batch N-1.

• Bubble Overhead: The initial filling and final draining of the pipeline create idle periods known as pipeline bubbles, which reduce efficiency.
• Goal: Maximize the duration of the steady state relative to bubble phases to approach ideal speedup.

The model is split into sequential stages, where each stage is a contiguous set of layers assigned to a single accelerator (e.g., GPU or NPU). This is a form of model parallelism.

• Granularity Trade-off: Fewer, larger stages minimize communication but increase pipeline bubbles and memory pressure per device. More, smaller stages reduce per-stage memory but increase communication overhead and bubble frequency.
• Partitioning Strategy: Optimal partitioning balances computational load, memory footprint, and the communication volume at stage boundaries to minimize the pipeline's critical path.

Data must be transferred between devices at each stage boundary. This communication is a primary bottleneck and must be optimized.

• Activation Passing: During the forward pass, each stage sends its output activations to the next stage.
• Gradient Passing: During the backward pass, gradients are passed backward through the pipeline.
• Overlap Potential: Modern implementations use techniques like double buffering to overlap this communication with computation, hiding some latency. The bandwidth of the inter-device interconnect (e.g., NVLink, InfiniBand) is crucial.

A key advantage is the reduction of peak memory consumption on each device compared to other parallelism strategies.

• Only Stage Parameters: Each device stores only the parameters and optimizer states for its assigned stage, not the entire model.
• Multiple Activations in Flight: A device must store the activations for all micro-batches currently in its segment of the pipeline (e.g., for the backward pass). This is managed via activation recomputation (checkpointing), which trades compute for memory by selectively re-computing activations during the backward pass.
• Enables Larger Models: This per-device memory reduction is what makes training models with hundreds of billions of parameters feasible.

The order in which forward and backward passes for different micro-batches are processed. The schedule determines efficiency and affects weight update semantics.

• Gpipe (Synchronous): Uses a simple, synchronous schedule (1F1B). Straightforward but suffers from large pipeline bubbles.
• Interleaved Schedules: More complex schedules like 1F1B (One Forward, One Backward) or PipeDream-Flush are used to reduce bubbles and improve throughput.
• Asynchronous vs. Synchronous: Some early schemes (PipeDream) used asynchronous updates, which could lead to weight staleness. Modern approaches like PipeDream-2BW and others in DeepSpeed and Megatron-LM maintain synchronous, consistent weight updates.

Pipeline Parallelism is rarely used alone. It is almost always combined with Data Parallelism and/or Tensor Parallelism to scale to thousands of accelerators and achieve optimal performance.

• Pipeline + Data Parallelism: Different pipeline instances (called model replicas) train on different subsets of data. This requires gradient synchronization across data-parallel groups after the backward pass.
• Pipeline + Tensor Parallelism: Within a single stage, the computation for individual layers (e.g., the linear layers in a Transformer block) can be further split via Tensor Parallelism. This is a hallmark of frameworks like Megatron-LM.
• 3D Parallelism: The combined use of Pipeline, Tensor, and Data Parallelism represents the state-of-the-art approach for training the largest foundation models.

Pipeline Parallelism splits a neural network's layers into sequential stages, each assigned to a different accelerator (e.g., GPU). Data flows through these stages as a sequence of micro-batches. The key innovation is the interleaving of forward and backward passes for different micro-batches across stages, creating a compute bubble at the start and end of the pipeline but achieving high device utilization in steady state. This is distinct from Model Parallelism, which splits individual layers, and Data Parallelism, which replicates the entire model.

• Stages: A contiguous block of model layers assigned to one device.
• Micro-batches: Small subdivisions of a training batch processed sequentially.
• Bubble: Idle time in the pipeline during ramp-up and ramp-down phases.

Introduced by Google in 2019, GPipe established the modern pipeline parallelism paradigm for training giant models. It uses a simple synchronous gradient descent approach with gradient accumulation across micro-batches. The primary challenge GPipe addresses is memory: by keeping only one micro-batch's activations per device at a time, it can train models whose layers exceed a single device's memory. However, its naive schedule leads to significant pipeline bubbles.

• Key Mechanism: Forward passes for all micro-batches, then backward passes in reverse order.
• Memory Efficiency: Only one set of intermediate activations per device is stored, enabled by gradient checkpointing.
• Limitation: The simple FIFO schedule results in low utilization during bubble phases.

Microsoft's PipeDream introduced optimized scheduling to minimize pipeline bubbles. Its key innovation is the 1F1B (One-Forward-One-Backward) schedule. Instead of completing all forward passes first (like GPipe), 1F1B interleaves forward and backward passes for different micro-batches as soon as possible, drastically reducing idle time and improving throughput.

• 1F1B Schedule: A device performs a forward pass for micro-batch i, then immediately a backward pass for an earlier micro-batch, keeping the pipeline full.
• Weight Stashing: PipeDream maintains multiple weight versions to ensure correct gradient updates despite asynchronous forward/backward passes.
• Throughput Gain: Can be several times faster than GPipe's naive schedule for deep pipelines.

Implementing pipeline parallelism involves navigating inherent trade-offs and technical challenges.

• Pipeline Bubble: Idle time is unavoidable. Its size is proportional to the number of pipeline stages. Optimized schedules (1F1B) minimize but cannot eliminate it.
• Activation Memory: Storing activations for the backward pass is memory-intensive. Gradient checkpointing (or activation recomputation) is essential, trading 33% more compute for a ~√N reduction in memory.
• Weight Synchronization: In schedules like 1F1B, ensuring consistent weight versions for forward and backward passes requires careful orchestration (weight stashing or periodic synchronization).
• Load Balancing: Uneven partitioning of layers across stages creates a straggler stage, limiting overall throughput. Optimal partitioning is non-trivial.

Communication between pipeline stages is a critical performance factor. The pattern is point-to-point and sequential, unlike the all-reduce operations used in data parallelism.

• Communication Pattern: Each device sends activations forward to the next stage and receives gradients backward from it.
• Overlap Strategy: High-performance implementations overlap communication with computation. For example, while computing a layer's forward pass, the output activation from the previous micro-batch can be asynchronously sent to the next stage.
• Hardware Topology: Performance is optimal when consecutive pipeline stages are placed on devices with high-bandwidth interconnects (e.g., NVLink). Network hops between stages add significant latency.

A broader category of distributed training where a model's parameters are partitioned across devices. Pipeline Parallelism is a specific form of Model Parallelism that organizes partitions into a sequential pipeline. Other forms include:

• Tensor Parallelism: Splits individual tensor operations (e.g., matrix multiplications within a layer) across devices.
• Spatial Partitioning: Distributes different channels or filters of a layer across devices. The key distinction is that Pipeline Parallelism partitions the model by layer sequence, while other methods partition within layers.

The most common distributed training strategy, where identical copies of the entire model are placed on each device, and the training dataset is split into shards (mini-batches). Each device processes its shard independently in the forward and backward pass, and gradients are averaged across all devices via an All-Reduce operation before updating weights. It contrasts with Pipeline Parallelism, which splits the model itself rather than the data. The two are often combined in large-scale training as hybrid parallelism.

A fine-grained model parallelism technique where individual matrix multiplications within a layer (e.g., the linear projections in a Transformer's attention block) are split across devices. For example, a large weight matrix is partitioned column-wise or row-wise. Devices must communicate (e.g., via All-Gather operations) after the distributed computation to combine results. It is often used within a single pipeline stage in Pipeline Parallelism to handle layers too large for one device, creating a two-level hierarchical parallel strategy.

A memory optimization paradigm for Data Parallelism that partitions the optimizer states, gradients, and model parameters across devices instead of replicating them. ZeRO-Offload and ZeRO-Infinity further extend this by offloading parts to CPU or NVMe memory. While ZeRO optimizes memory within Data Parallelism, it is frequently combined with Pipeline and Tensor Parallelism to enable the training of trillion-parameter models by eliminating redundant memory copies across all parallel dimensions.

A memory optimization technique critical for enabling deep pipelines. It trades compute for memory by selectively saving only a subset of layer activations during the forward pass. Non-checkpointed activations are recomputed during the backward pass as needed. In Pipeline Parallelism, this drastically reduces the memory footprint of each pipeline stage, allowing for more layers per device or the use of larger micro-batches. It is essential for training very deep models where activation memory is the primary bottleneck.

The orchestration algorithm that determines the order in which micro-batches are injected into and flow through the pipeline to maximize hardware utilization. Key schedules include:

• GPipe (Synchronous): Uses a simple 1F1B (One Forward pass followed by One Backward pass) schedule with periodic pipeline flushes, leading to 'bubbles' of idle time.
• Interleaved 1F1B: A more advanced schedule that splits stages further (virtual pipelining) to reduce the bubble size.
• Asynchronous Schedules: Allow forward and backward passes of different micro-batches to overlap more aggressively, at the potential cost of training stability. The schedule directly impacts pipeline efficiency and throughput.