Pipeline Parallelism

A pipeline stage is a contiguous set of model layers assigned to a single GPU. The model is split vertically along its depth. For example, a 24-layer transformer might be split into 4 stages of 6 layers each. Each stage operates independently on its assigned micro-batch, passing activations to the next stage via high-speed interconnects like NVLink. The number of stages is typically equal to the number of available GPUs, defining the pipeline's length.

A micro-batch is a small, fixed-size subset of the overall training batch. The global batch is split into many micro-batches that are fed into the pipeline sequentially. This is critical for keeping the pipeline full: while one micro-batch is in the forward pass of stage 2, another can be in stage 1. Micro-batch size is a key tuning parameter, balancing GPU memory usage against pipeline bubble size. Common sizes range from 1 to 32 samples.

A pipeline bubble is a period of idle time in a GPU's schedule, representing underutilization. Bubbles occur during pipeline fill (when the first micro-batch is propagating through empty stages) and flush (when the last micro-batch completes). The bubble size is fundamentally determined by the number of pipeline stages. Techniques like 1F1B (One Forward One Backward) scheduling and interleaved scheduling are used to reduce bubble overhead and improve throughput.

1F1B (One Forward One Backward) is the standard scheduling protocol for pipeline parallelism. Each device alternates between performing a forward pass for one micro-batch and a backward pass for another. This pattern:

• Maintains a steady state where all devices are busy.
• Ensures memory usage is bounded, as activations for a micro-batch can be discarded after its backward pass.
• Creates a deterministic schedule that is easier to debug and reason about compared to more complex schemes.

Interleaved scheduling is an advanced optimization where each physical GPU is assigned multiple virtual pipeline stages. For example, on a 4-GPU system, each GPU might handle 2 virtual stages, creating an 8-stage virtual pipeline. This technique:

• Reduces the size of pipeline bubbles by increasing the number of micro-batches in flight.
• Improves GPU utilization at the cost of increased implementation complexity and memory overhead, as each GPU must hold the parameters and activations for multiple model segments.

Gradient synchronization ensures weight updates are consistent across all pipeline stages. During the backward pass, gradients flow backward through the pipeline. Once the backward pass for a micro-batch completes at the first stage, its gradients must be synchronized with the gradients from other micro-batches in the same global batch. This is typically done using an All-Reduce operation across all pipeline-parallel GPUs, often overlapped with computation to hide communication latency.

A model parallelism strategy that splits individual tensor operations (e.g., matrix multiplications within a single layer) across multiple GPUs. Unlike pipeline parallelism, which partitions layers, tensor parallelism partitions the computations within a layer.

• Key Difference: Requires extremely high-bandwidth communication (e.g., NVLink) between devices for each layer's forward pass.
• Use Case: Often combined with pipeline parallelism in very large models, where a single layer is too big for one GPU's memory.

A neural network architecture where a routing network dynamically selects a small subset of specialized 'expert' sub-networks for each input token. This enables massive model capacity with a manageable computational cost per token.

• Sparse Activation: Only the parameters of the selected experts are activated, leading to more efficient computation than a dense model of equivalent size.
• Parallelism Synergy: MoE layers are often distributed using techniques like tensor or pipeline parallelism to handle the large number of experts.

An inference optimization technique that dynamically groups incoming requests of varying sequence lengths into a single batch to maximize GPU utilization. It is highly complementary to pipeline parallelism.

• Orchestration: In a pipelined system, continuous batching can be applied to the micro-batches flowing through each stage, keeping all GPUs busy with heterogeneous workloads.
• Throughput vs. Latency: Maximizes overall system throughput, which helps amortize the pipeline 'bubble' overhead inherent in pipeline parallelism.

vLLM is a high-throughput inference engine renowned for its PagedAttention algorithm, which manages the KV cache like an operating system manages virtual memory.

• Memory Efficiency: Treats the KV cache as non-contiguous 'pages,' eliminating memory fragmentation and allowing for longer contexts and higher batch sizes.
• Integration Point: While vLLM excels at attention-layer optimization, its efficient memory management can be leveraged within the individual stages of a pipeline-parallel system to handle larger micro-batches.

A compression technique that reduces the numerical precision of a model's weights and activations (e.g., from FP16 to INT8) to decrease memory footprint and increase computational speed.

• Synergy with Parallelism: Applying post-training quantization (PTQ) or quantization-aware training (QAT) to a model before applying pipeline parallelism reduces the memory and bandwidth requirements for each GPU stage, potentially allowing for more layers per device or faster inter-stage communication.

A compiler-level optimization that combines multiple sequential neural network operations into a single, custom kernel. This reduces intermediate memory reads/writes and kernel launch overhead.

• Stage-Level Optimization: Critical for maximizing the performance of each individual 'stage' in a pipeline. Fusing operations within a stage (e.g., linear layer + activation + normalization) reduces latency, which directly improves the overall throughput of the pipeline by processing micro-batches faster.