Pipeline Parallelism

The model is partitioned into sequential stages, each assigned to a different GPU. During the forward pass, activations flow from the first stage to the last. In the backward pass, gradients flow in reverse. This creates a pipeline where different GPUs work on different micro-batches simultaneously, but for a single input, computation is sequential across devices.

To keep all GPUs busy and reduce bubble time (idle periods), the input batch is split into smaller micro-batches. These micro-batches are fed into the pipeline in a staggered fashion. For example, while GPU 2 processes the forward pass of micro-batch 1, GPU 1 can start the forward pass of micro-batch 2. This technique, central to frameworks like GPipe, is crucial for achieving high GPU utilization and throughput.

Performance is heavily dependent on the communication overhead between stages. Each stage must send its output activations to the next GPU and receive gradients during the backward pass. This makes the technique sensitive to interconnect bandwidth (e.g., NVLink vs. PCIe). The pipeline depth (number of stages) introduces a fundamental latency cost proportional to the number of devices.

Its primary advantage is memory distribution. By splitting the model across devices, the memory footprint for parameters, gradients, and optimizer states is divided, enabling the training of models that are too large to fit on a single GPU. This is often combined with other parallelism strategies like data parallelism or ZeRO for training foundation models with hundreds of billions of parameters.

1F1B (One Forward pass followed by One Backward pass) is a popular pipeline scheduling algorithm that optimizes memory usage. Instead of running all forward passes first, it interleaves forward and backward passes for different micro-batches. This reduces the amount of activation memory that needs to be stored in the pipeline, as activations can be freed sooner after their corresponding backward pass is completed.

Pipeline parallelism is rarely used alone. It is typically combined with:

• Data Parallelism: Different pipeline instances handle different data batches.
• Tensor Parallelism: Within a single stage, individual layers (like the attention or MLP blocks) are further split across GPUs. Frameworks like Megatron-LM and DeepSpeed implement hybrid 3D parallelism, integrating pipeline, tensor, and data parallelism to train models at extreme scale.

Tensor parallelism is a model parallelism technique that splits individual layers (operations) of a neural network across multiple GPUs. Unlike pipeline parallelism, which assigns whole layers to different devices, tensor parallelism divides the computations within a single layer.

• Key Mechanism: For a linear layer Y = XW, the weight matrix W is partitioned column-wise or row-wise across devices. Each GPU computes a portion of the output, requiring all-to-all communication (e.g., via all-reduce) to combine results.
• Use Case: Essential for layers that are too large to fit on one GPU's memory, such as the massive feed-forward layers in transformer models.
• Communication Overhead: High, as synchronization is required after every parallelized operation.

Data parallelism is a distributed training strategy where the same model is replicated across multiple GPUs, and each replica processes a different subset (mini-batch) of the training data.

• Key Mechanism: Gradients are computed independently on each device, then averaged across all devices via a collective operation like all-reduce before updating the model weights.
• Primary Goal: To scale training by processing more data simultaneously, reducing the time per epoch.
• Contrast with Model Parallelism: Data parallelism splits the data; model parallelism (like pipeline or tensor) splits the model. They are often combined in training massive models.

ZeRO is a memory optimization paradigm for data parallel training that partitions the optimizer states, gradients, and model parameters across processes to eliminate memory redundancy.

• Core Innovation: Instead of each GPU holding a full copy of the model (~16GB for a 1B parameter model in FP16), ZeRO stages (1, 2, and 3) progressively partition these components, reducing per-device memory footprint.
• Impact: Enables the training of models far larger than the memory of any single GPU by leveraging the aggregate memory of a cluster.
• Relation to Pipeline Parallelism: ZeRO is complementary; it optimizes data parallelism, which is frequently used alongside pipeline parallelism in frameworks like DeepSpeed.

Continuous batching (or iterative batching) is an inference-time optimization that dynamically groups incoming requests into a single batch to maximize GPU utilization, even if the requests start at different times.

• Key Mechanism: As some requests finish generation, new requests are seamlessly slotted into the vacated space within the ongoing batch. This contrasts with static batching, which waits for a fixed batch size to form.
• Primary Benefit: Dramatically improves throughput (tokens/second) and GPU utilization for LLM serving, as the GPU is kept constantly busy.
• Deployment Context: A critical feature in high-performance inference servers like vLLM and Text Generation Inference (TGI).

vLLM is an open-source, high-throughput LLM inference and serving engine. Its key innovation is PagedAttention, an algorithm that manages the KV cache similarly to how an operating system manages virtual memory.

• Problem Solved: The KV cache for long sequences is non-contiguous and variable-sized, leading to significant memory fragmentation and waste (up to 80%).
• PagedAttention Solution: It divides the KV cache into fixed-size blocks. These blocks can be non-contiguously stored in GPU memory and are dynamically allocated as needed.
• Result: Enables near-zero waste in KV cache memory, allowing for larger batch sizes and longer context lengths, which directly improves serving throughput and cost-efficiency.

Model serving is the process of deploying a trained machine learning model into a production environment where it can receive input data, perform inference, and return predictions via an API.

• Core Components: Includes the inference endpoint (the API), the serving runtime (software that loads the model and executes it), and the surrounding infrastructure for scaling, monitoring, and updating.
• Key Challenges: Low latency, high throughput, efficient resource utilization, and seamless model updates.
• Serving Frameworks: Specialized tools like Triton Inference Server, Text Generation Inference (TGI), and vLLM are used to address the unique demands of serving large, computationally intensive models like LLMs.