Pipeline Parallelism

The neural network's computational graph is partitioned layer-wise across multiple devices. Each device is assigned a distinct subset of consecutive layers (a stage). Data flows sequentially from one stage to the next, forming a processing pipeline. This is distinct from data parallelism, where the entire model is replicated, and from tensor parallelism, where individual layers are split.

• Example: In a 12-layer transformer, Device 1 handles layers 1-4, Device 2 handles layers 5-8, and Device 3 handles layers 9-12.

To keep the pipeline full and maximize hardware utilization, the training batch is split into smaller microbatches. These microbatches are fed into the pipeline in a staggered fashion. Different devices work on different microbatches concurrently, a technique known as interleaved scheduling.

• Key Benefit: Hides the communication latency between stages by ensuring computation is almost always occurring somewhere in the pipeline.
• Challenge: Requires careful scheduling to avoid pipeline bubbles (idle time) during the initial fill and final drain phases.

Pipeline bubbles are periods of idle time within a stage, representing the fundamental inefficiency of the paradigm. They occur during the pipeline fill (startup) and pipeline drain (wind-down) phases, and whenever stages have imbalanced computational loads.

• Bubble Time: For a pipeline with p stages and a microbatch count m, the fraction of time spent in bubbles is approximately (p-1) / (m + p - 1).
• Mitigation: Increasing the number of microbatches (m) relative to stages (p) reduces the bubble fraction, improving pipeline utilization.

Communication occurs at the boundaries between pipeline stages. After a device finishes processing its assigned layers for a microbatch, it must send the activations (forward pass) or gradients (backward pass) to the next device. This communication is a primary bottleneck.

• Overlap: Efficient implementations overlap this communication with computation for other microbatches to hide latency.
• Topology: Performance is highly sensitive to the interconnect bandwidth (e.g., NVLink, InfiniBand) between devices hosting adjacent stages.

Each device only stores the parameters and optimizer states for its assigned subset of layers. This provides a linear reduction in per-device memory footprint compared to data parallelism, enabling the training of models far larger than the memory of any single device.

• Memory Advantage: For a model with M parameters split across p devices, each device holds roughly M/p parameters.
• Trade-off: This comes at the cost of increased communication and the pipeline bubble overhead.

Different scheduling algorithms manage the flow of forward and backward passes to optimize memory or performance.

• GPipe Schedule: Processes all microbatches in the forward pass first, then all in the backward pass. Simple but requires storing activations for all microbatches, leading to high activation memory.
• 1F1B (One Forward, One Backward) Schedule: Interleaves forward and backward passes for each microbatch. A device performs a forward pass for microbatch i, then later a backward pass for microbatch i-k (where k is a pipeline-dependent constant). This significantly reduces the peak activation memory required.
• Interleaved 1F1B: A more advanced variant that further splits stages into more, smaller virtual stages to improve load balancing and reduce bubble size.

The fundamental pattern involves splitting a neural network's computational graph into sequential stages, each assigned to a different device (e.g., GPU, NPU). Data is processed as a stream of microbatches. While Device 1 processes microbatch N, Device 2 processes microbatch N-1, creating an assembly-line effect. This requires careful management of activation memory (the outputs of each layer passed between devices) and gradient synchronization during backward passes.

• Stage Partitioning: The model is split at layer boundaries. The goal is to balance computational load and communication overhead between stages.
• Microbatch Streaming: The training batch is divided into smaller microbatches that are fed into the pipeline sequentially to keep all devices busy.
• Bubble Management: Pipeline bubbles (idle time) occur at the start (warm-up) and end (drain) of processing a batch. Techniques like 1F1B (One Forward pass followed by One Backward pass) scheduling minimize these bubbles.

Different algorithms schedule the forward and backward passes of microbatches to optimize for memory or throughput.

• GPipe (Google): Uses a simple F-then-B schedule. All microbatches complete their forward passes across all stages before any backward pass begins. This is simple but requires storing all intermediate activations in memory, leading to high memory pressure.
• 1F1B (One-Forward-One-Backward): A more memory-efficient schedule. Once the pipeline is warm, each device alternates between a forward pass for one microbatch and a backward pass for another. This allows activations to be freed sooner, reducing peak memory consumption.
• Interleaved 1F1B: A variant where each physical device hosts multiple virtual stages. This further improves pipeline utilization and balance by allowing finer-grained partitioning of the model.

PyTorch provides native support for pipeline parallelism via the torch.distributed.pipeline.sync.Pipe module (and experimental async versions).

Key Components:

• Pipe Class: Wraps a torch.nn.Sequential module split across devices.
• Automatic Splitting: The Pipe class can automatically partition the model or accept a user-defined partition.
• Chunks Argument: This specifies the number of microbatches, directly controlling the granularity of the pipeline stream.

Example Workflow:

1. Move model segments to different devices.
2. Wrap the distributed module in Pipe.
3. The pipeline automatically handles forward/backward pass scheduling, communication, and gradient synchronization.

Microsoft's DeepSpeed library offers advanced pipeline parallelism as part of its 3D parallelism strategy (combining Data, Model, and Pipeline parallelism).

Key Features:

• PipelineEngine: A core class that manages the pipeline execution, supporting the 1F1B schedule and its interleaved variant.
• Integration with ZeRO: Can be combined with ZeRO memory optimization stages to reduce the memory footprint of optimizer states, gradients, and parameters within each pipeline stage.
• Flexible Config: Pipeline parallelism is configured via a JSON file, specifying the number of pipeline stages and micro-batch size.
• Improved Fault Tolerance: Includes features for saving and restoring pipeline engine state during training.

Effective pipeline parallelism is a balance between computation and communication.

• Communication Primitives: Relies heavily on point-to-point operations like send/recv or P2P operations in NCCL (for GPUs) or vendor-specific collectives for NPUs to transfer activations and gradients between stages.
• Activation Checkpointing (Gradient Checkpointing): Critical for memory management. Instead of storing all activations for the backward pass, only a subset (e.g., at stage boundaries) is stored. The others are recomputed during the backward pass, trading compute for memory.
• Bandwidth vs. Latency: The throughput of the pipeline is often limited by the slowest stage (pipeline stall) and the communication bandwidth between devices. Models with large activation sizes (e.g., in generative models) are particularly sensitive to inter-device link speed.

Implementing pipeline parallelism on Neural Processing Units requires adapting to specific architectural constraints.

• Vendor SDK Integration: Leveraging proprietary communication libraries (e.g., HCCL for Ascend, Habana Collective Communications Library) for efficient inter-device data transfer.
• Memory Hierarchy Alignment: Staging activation data in the optimal level of the NPU's memory hierarchy (e.g., on-chip buffer vs. HBM) before sending to the next device to minimize transfer latency.
• Computation/Communication Overlap: Using hardware-specific asynchronous copy engines or DMA controllers to overlap the communication of activations/gradients with the computation of the next microbatch.
• Partitioning for Balanced Load: Profiling layer execution time on the target NPU is essential, as the cost of operations can differ significantly from GPUs, affecting optimal stage boundaries.

Data parallelism is a parallel computing paradigm where the same operation (e.g., a forward pass through a neural network) is applied concurrently to different subsets (batches) of a dataset across multiple processing units (e.g., GPUs or NPU cores).

• Key Mechanism: Each device holds a complete copy of the model. The global batch is split into smaller microbatches, which are processed in parallel. Gradients are then synchronized across devices (e.g., via All-Reduce) to update the model.
• Primary Goal: Scale training by processing more data simultaneously, reducing the time per epoch.
• Contrast with Pipeline Parallelism: While data parallelism replicates the model, pipeline parallelism partitions the model. They are often combined: a model is split across devices (pipeline parallelism), and each model partition is replicated (data parallelism) for further scaling.

Model parallelism is a broad technique for distributing the computational graph or parameters of a neural network across multiple processors or devices, primarily to handle models whose memory footprint exceeds the capacity of a single unit.

• Key Mechanism: Different layers or subsets of layers are placed on different devices. During forward/backward passes, activations and gradients are communicated between devices.
• Primary Goal: Enable the training of models that are too large to fit on one accelerator.
• Relationship to Pipeline Parallelism: Pipeline parallelism is a specific, optimized form of model parallelism. Traditional model parallelism may split a single batch across layers, leading to device idle time. Pipeline parallelism introduces the microbatch concept to keep all devices busy simultaneously, transforming the execution into a staged pipeline.

Tensor parallelism is a fine-grained form of model parallelism that splits individual tensor operations (e.g., large matrix multiplications within a layer) across multiple devices.

• Key Mechanism: For a linear layer Y = XA + B, the weight matrix A is partitioned along its rows or columns. The input X is broadcast, partial results are computed on each device, and then synchronized via a collective operation (e.g., All-Gather).
• Primary Goal: Distribute the computational load of massive individual layers that are bottlenecks even within a pipeline stage.
• Practical Use: Often used within a single device or combined with pipeline parallelism. For example, in the Megatron-LM architecture, tensor parallelism splits attention and MLP layers, while pipeline parallelism splits sequences of layers.

Task parallelism (or functional parallelism) is a parallel computing model where different, independent tasks or functions are executed concurrently on multiple processing units.

• Key Mechanism: The program is decomposed into distinct tasks that can run in parallel. These tasks may operate on the same or different data and communicate as needed.
• Primary Goal: Exploit concurrency in applications with heterogeneous or independent workloads.
• Contrast with Pipeline Parallelism: Pipeline parallelism is a specific, structured form of task parallelism where the "tasks" are sequential stages of a single, data-dependent computational graph. In a general task-parallel system, tasks might have complex, irregular dependencies, not the linear, producer-consumer relationship of a pipeline.

Synchronization primitives are low-level programming constructs that coordinate the execution and memory access of concurrent threads or processes, which are critical for implementing correct parallel schemes like pipeline parallelism.

• Key Primitives:
  • Barriers: Force all threads/processes to wait until every one reaches a specific point (e.g., between pipeline stages in a synchronous schedule).
  • Semaphores/Mutexes: Control access to shared resources (e.g., a shared queue holding microbatches between stages).
  • Atomic Operations: Ensure indivisible updates to shared variables (e.g., counters for tracking microbatch IDs).
  • Memory Fences: Enforce ordering of memory operations, crucial for correctness in weak memory models.
• Relevance: Efficient pipeline execution requires careful synchronization to manage the flow of microbatches, maintain data dependencies, and implement schedules like 1F1B (One-Forward-One-Backward).

Amdahl's Law is a fundamental principle that models the theoretical speedup of a parallel program, providing critical context for evaluating parallelism strategies like pipeline parallelism.

• Formula: Speedup = 1 / (S + P/N), where S is the fraction of serial work, P is the parallelizable fraction (S+P=1), and N is the number of processors.
• Implication: The serial portion S becomes the ultimate bottleneck. Even infinite processors cannot achieve a speedup greater than 1/S.
• Connection to Pipeline Parallelism:
  • Strong Scaling (Fixed Problem): Pipeline parallelism aims to reduce time for a fixed model by adding devices. Its speedup is limited by pipeline bubbles (idle time during fill/flush) and any serial operations (e.g., input/output).
  • Weak Scaling (Fixed Problem per Device): Pipeline parallelism excels here, as adding more devices (pipeline stages) allows for proportionally larger models to be trained without increasing time per iteration, assuming communication overhead is managed.