Pipeline Parallelism

The model's layers are divided into a sequence of stages, with each stage assigned to a different GPU. During training, a micro-batch of data passes through these stages sequentially, analogous to an assembly line. This is distinct from tensor parallelism, which splits individual layers. The number of layers per stage is a critical hyperparameter that balances memory savings against communication overhead and pipeline bubbles.

A fundamental inefficiency where GPUs are idle while the pipeline fills and drains. For a pipeline with p stages, a significant fraction of time is spent not performing useful computation. This bubble overhead is inversely proportional to the number of micro-batches in flight. Techniques like 1F1B (One Forward pass followed by One Backward pass) scheduling are designed to minimize these bubbles and improve GPU utilization.

To amortize pipeline bubble overhead, the training batch is split into smaller micro-batches. These micro-batches are injected into the pipeline in a staggered fashion, keeping all stages busy concurrently. The micro-batch size is a key tuning parameter:

• Smaller micro-batches reduce per-GPU memory pressure but may underutilize compute.
• Larger micro-batches improve compute efficiency but increase memory consumption and pipeline flush time.

Activations (forward pass) and gradients (backward pass) must be communicated between GPUs hosting adjacent stages. This point-to-point communication occurs over high-speed interconnects like NVLink or InfiniBand. The volume of data transferred is proportional to the activation size and micro-batch size, making it a potential bottleneck. This overhead contrasts with the all-to-all communication patterns often found in data parallelism.

The primary advantage is memory distribution. Each GPU only needs to store the parameters, gradients, and optimizer states for its assigned stage, not the entire model. This enables training of extremely large models (e.g., models with hundreds of billions of parameters) that would be impossible to fit on a single device, even with techniques like gradient checkpointing.

Pipeline parallelism is rarely used alone. It is typically combined with data parallelism across multiple pipeline instances to scale further. This hybrid approach, sometimes called pipeline-model-data parallelism, uses pipeline parallelism to fit the model and data parallelism to increase the effective batch size. It can also be combined with tensor parallelism within a single stage for even larger models, a strategy used in frameworks like Megatron-LM.

Pipeline parallelism is the enabling technology for training the world's largest foundation models, such as GPT-3, PaLM, and Llama 2, which can have hundreds of billions or trillions of parameters. It allows researchers to scale model size far beyond the memory limits of individual accelerators by distributing sequential layers across a pipeline of GPUs. Key implementations include:

• Megatron-LM (NVIDIA): A pioneering framework that combines pipeline parallelism with tensor parallelism for efficient large-model training.
• DeepSpeed (Microsoft): Its PipelineEngine manages complex schedules and integrates with ZeRO memory optimization.
• FairScale (Meta): Provides a Pipe module for implementing pipeline parallelism within PyTorch.

A critical implementation detail is the micro-batch scheduling algorithm that determines GPU utilization. The 1F1B schedule is the industry standard for optimal efficiency. It works by:

• Interleaving forward and backward passes for different micro-batches on each GPU.
• Ensuring that once the pipeline is 'warmed up,' every GPU is continuously busy with either a forward or backward pass.
• Minimizing the pipeline bubble—the idle time where GPUs wait for data from other stages—compared to a naive sequential schedule. This schedule is essential for achieving high hardware utilization and making pipeline-parallel training economically viable.

The performance bottleneck in pipeline parallelism is often the communication of activations and gradients between stages. Implementations must carefully manage:

• Activation Memory: Storing activations for the backward pass consumes significant GPU memory. Techniques like gradient checkpointing (recomputing activations) are often combined with pipeline parallelism to trade compute for memory.
• Communication Overlap: Advanced frameworks overlap the communication of tensors between GPUs with computation on other micro-batches to hide latency.
• Network Topology: Placing consecutive pipeline stages on GPUs connected by high-bandwidth links (e.g., NVLink) is critical for performance.

Pipeline parallelism is rarely used alone. Real-world systems employ 3D parallelism, a combination of techniques:

• Data Parallelism: Splits the training batch across multiple replicas of the entire model pipeline.
• Tensor Parallelism: Splits individual layers (e.g., attention heads, MLP blocks) across GPUs within a single pipeline stage for intra-layer parallelism.
• Pipeline Parallelism: Splits the model layers across stages. Frameworks like DeepSpeed and Megatron-DeepSpeed automate the orchestration of this 3D parallelism, allowing engineers to specify a total number of GPUs and automatically configure the optimal combination of parallel strategies for a given model size.

Production deployments must address inherent challenges:

• Fault Tolerance: A failure in one GPU halts the entire pipeline. Checkpointing the state of all pipeline stages simultaneously is complex but necessary for resuming long-running jobs.
• Load Balancing: The computational load must be evenly distributed across pipeline stages. An imbalance creates a bottleneck at the slowest stage. Model partitioning algorithms analyze layer FLOPs and memory usage to create balanced stages.
• Heterogeneous Hardware: In cloud environments, pipeline stages might run on different GPU instance types, requiring dynamic scheduling to accommodate variable performance.

While primarily a training technique, pipeline parallelism concepts are applied to inference for ultra-large models. Inference-serving systems can partition a model across multiple GPUs or even nodes, streaming tokens through the pipeline. This is distinct from and often combined with continuous batching for request-level parallelism. The key difference is the absence of a backward pass, allowing for different optimization strategies to minimize inference latency while handling models that cannot fit on a single server.

Tensor parallelism is a model parallelism technique that splits individual layers of a neural network (e.g., the weight matrices within a transformer's attention or feed-forward blocks) across multiple GPUs. Unlike pipeline parallelism, which partitions the model by layers (vertical split), tensor parallelism partitions operations within a single layer (horizontal split).

• Key Mechanism: For a linear layer Y = XW, the input X and weight matrix W are split along specific dimensions. The computation is performed in parallel across devices, requiring all-reduce communication after the operation to combine results.
• Use Case: Essential for layers too large to fit on a single GPU's memory. Often combined with pipeline parallelism in systems like Megatron-LM to train models with trillions of parameters.
• Trade-off: Introduces very frequent, fine-grained communication, making it efficient only within high-bandwidth interconnects like NVLink.

Data parallelism is the most common distributed training strategy, where identical copies of the entire model are placed on multiple GPUs, and each processes a different subset (mini-batch) of the training data.

• Key Mechanism: Gradients are computed independently on each device, then averaged via an all-reduce operation before updating the model weights synchronously (Synchronous SGD) or asynchronously.
• Contrast with Pipeline Parallelism: Data parallelism replicates the entire model, scaling with the dataset size. Pipeline parallelism splits the model itself, scaling with model size. They are orthogonal and often combined: pipeline parallelism splits the model across a pipeline parallel group, while data parallelism replicates each pipeline stage across a data parallel group*.
• Limitation: Requires the full model to fit on a single GPU, which is impossible for modern foundation models, necessitating model parallelism techniques.

Gradient checkpointing (or activation checkpointing) is a memory optimization technique that trades compute for memory, enabling the training of larger models or larger batch sizes within limited GPU memory—a critical complement to pipeline parallelism.

• Key Mechanism: Instead of storing all intermediate activations (needed for the backward pass), the system only saves select 'checkpoint' activations. Non-checkpointed activations are recomputed on-demand during the backward pass.
• Synergy with Pipeline Parallelism: Pipeline parallelism reduces the memory footprint per GPU by storing only a subset of layers. Gradient checkpointing further reduces the peak memory within each stage by minimizing activation storage. This combination is standard for training giant models.
• Impact: Can reduce activation memory by up to 75%, but typically increases total computation time by 20-30% due to recomputation.

A Mixture of Experts (MoE) model is a sparsely activated architecture where each input token is routed to only a small subset of specialized sub-networks ('experts'), dramatically increasing parameter count without a proportional increase in computation.

• Key Mechanism: A gating network selects 1-2 experts per token. Only the weights of the activated experts are used for computation, making the model conditionally sparse.
• Relationship to Parallelism: MoE models are exceptionally well-suited for pipeline and tensor parallelism. Experts can be distributed across devices, and the sparse activation pattern creates natural parallelism. Training systems like GShard and Switch Transformers use sophisticated parallelism strategies to handle the routing and communication overhead.
• Benefit: Enables models with over a trillion parameters while keeping the computational cost per token similar to a dense model of 10-100B parameters.

ZeRO is a memory optimization paradigm for data parallelism that partitions the three main model states—optimizer states, gradients, and parameters—across devices to eliminate memory redundancy.

• Key Stages: ZeRO-1 partitions optimizer states, ZeRO-2 partitions gradients, and ZeRO-3 partitions all model parameters. Each stage requires progressively more communication to gather partitions when needed.
• Contrast with Pipeline Parallelism: ZeRO is a form of optimized data parallelism, not model parallelism. However, ZeRO-3 can be seen as a form of tensor parallelism on a massive scale. It is often used instead of or alongside pipeline parallelism in frameworks like DeepSpeed to train massive models. Pipeline parallelism handles layer-to-layer distribution, while ZeRO handles memory-efficient state management within stages.
• Outcome: Can train models 10x larger on the same hardware compared to standard data parallelism.

Model sharding is a general term for any technique that partitions a model's parameters across multiple devices. It is the overarching category that includes both pipeline and tensor parallelism.

• Key Concept: The goal is to overcome the memory limitations of a single accelerator. Sharding can be applied to different dimensions of the model: by layer (pipeline), by operation within a layer (tensor), or by parameter groups (ZeRO).
• Implementation: Requires a parallelism strategy (how to split) and a execution schedule (when to compute and communicate). Pipeline parallelism uses a 1F1B (One Forward pass followed by One Backward pass) or similar schedule to minimize the 'bubble' of idle time.
• System Design: Modern training frameworks (e.g., PyTorch Fully Sharded Data Parallel, DeepSpeed, Megatron) provide automated or semi-automated sharding strategies, allowing engineers to compose different parallelism forms (3D Parallelism: Data + Tensor + Pipeline) for optimal efficiency.