pipeline parallelism deep learning

**Pipeline Parallelism in Deep Learning** is **the model partitioning strategy that assigns different layers (stages) of a neural network to different GPUs, flowing microbatches through the pipeline — enabling training of models too large for a single GPU's memory while achieving reasonable hardware utilization through overlapping forward and backward passes across stages**. **Pipeline Partitioning:** - **Stage Assignment**: model layers divided into K stages assigned to K GPUs; each stage holds consecutive layers; stage boundary placement balances compute time across stages to minimize pipeline bubble - **Memory Motivation**: a 175B parameter model requires ~350 GB in fp16 weights alone; pipeline parallelism distributes layers across GPUs, with each GPU holding only 1/K of the parameters plus activations for in-flight microbatches - **Communication**: only activation tensors cross stage boundaries (one tensor transfer per microbatch per stage boundary); communication volume is much smaller than all-reduce gradient synchronization in data parallelism - **Layer Balance**: unequal layer compute costs create pipeline stalls where fast stages wait for slow stages; profiling per-layer compute time and balancing memory + compute is an NP-hard partitioning problem **Pipeline Schedules:** - **GPipe (Synchronous)**: inject M microbatches forward through all stages, then all backward — results in a pipeline bubble of (K-1)/M fraction of total time; increasing microbatches M reduces bubble but increases activation memory (each stage stores all M forward activations for backward pass) - **1F1B (One-Forward-One-Backward)**: after filling the pipeline with forward passes, alternate one forward and one backward per stage — limits peak activation memory to K microbatches (vs M for GPipe); bubble fraction same as GPipe but memory is dramatically reduced - **Interleaved 1F1B (Megatron-LM)**: each GPU holds multiple non-consecutive stages (e.g., GPU 0 holds stages 0 and 4); reduces pipeline bubble by (V-1)/(V*K-1) where V is virtual stages per GPU — 2× more stage boundaries doubles communication but halves bubble - **Zero-Bubble Schedule**: advanced scheduling algorithms (Qi et al. 2023) overlap backward-weight-gradient computation with forward passes from later microbatches — theoretically eliminates bubble with careful dependency analysis **Activation Memory Management:** - **Activation Checkpointing**: discard forward activations after use, recompute during backward pass — trades 33% extra compute for ~K× activation memory reduction; essential for deep pipelines with many microbatches - **Activation Offloading**: transfer activations to CPU memory during the pipeline fill phase, fetch back during backward — overlaps CPU-GPU transfer with computation to hide latency - **Memory-Efficient Schedule**: 1F1B schedule inherently limits activation memory by starting backward passes before all forward passes complete — steady state holds only K microbatch activations simultaneously **Combining with Other Parallelism:** - **3D Parallelism**: combining pipeline parallelism (inter-layer), tensor parallelism (intra-layer), and data parallelism (across replicas) enables training models like GPT-3 (175B), PaLM (540B) on thousands of GPUs simultaneously - **Pipeline + ZeRO**: ZeRO optimizer state partitioning within each pipeline stage reduces per-GPU memory further; each stage's data-parallel workers shard optimizer states - **Pipeline + Expert Parallelism**: MoE models use expert parallelism within stages and pipeline parallelism across stage groups — Mixtral/Switch Transformer architectures leverage both Pipeline parallelism is **an essential technique for training the largest neural networks — the key engineering challenge is minimizing the pipeline bubble (idle time) through schedule optimization while managing activation memory through checkpointing, making deep pipeline training both memory-efficient and compute-efficient**.