pipeline parallelism training

**Pipeline Parallelism** is **the model parallelism technique that partitions neural network layers across multiple devices and processes multiple micro-batches concurrently in a pipeline fashion — enabling training of models too large for a single GPU by distributing consecutive layers to different devices while maintaining high GPU utilization through careful scheduling of forward and backward passes across overlapping micro-batches**. **Pipeline Parallelism Fundamentals:** - **Layer Partitioning**: divides model into stages (consecutive layer groups); stage 0 on GPU 0, stage 1 on GPU 1, etc.; each stage processes its layers then passes activations to next stage - **Sequential Dependency**: forward pass flows stage 0 → 1 → 2 → ...; backward pass flows in reverse; creates inherent sequential bottleneck - **Naive Pipeline Problem**: without micro-batching, only one GPU is active at a time; GPU utilization = 1/num_stages; completely impractical for more than 2-3 stages - **Micro-Batching Solution**: splits mini-batch into smaller micro-batches; processes multiple micro-batches in flight simultaneously; overlaps computation across stages **GPipe (Google):** - **Synchronous Pipeline**: processes all micro-batches of a mini-batch before updating weights; maintains synchronous SGD semantics; gradient accumulation across micro-batches - **Forward-Then-Backward Schedule**: completes all forward passes for all micro-batches, then all backward passes; simple but high memory usage (stores all activations) - **Pipeline Bubble**: idle time during pipeline fill (ramp-up) and drain (ramp-down); bubble_time = (num_stages - 1) × micro_batch_time; efficiency = 1 - bubble_time / total_time - **Activation Checkpointing**: recomputes activations during backward pass to reduce memory; essential for deep pipelines; trades 33% more computation for 90% less activation memory **PipeDream (Microsoft):** - **Asynchronous Pipeline**: doesn't wait for all micro-batches to complete; uses weight versioning to handle concurrent forward/backward passes with different weight versions - **1F1B Schedule (One-Forward-One-Backward)**: alternates forward and backward micro-batches after initial warm-up; reduces memory usage (stores fewer activations) compared to GPipe - **Weight Stashing**: maintains multiple weight versions for different in-flight micro-batches; ensures gradient consistency; memory overhead for storing weight versions - **Vertical Sync**: periodically synchronizes weights across all stages; balances staleness and consistency; configurable sync frequency **Pipeline Scheduling Strategies:** - **Fill-Drain (GPipe)**: fill pipeline with forward passes, drain with backward passes; high memory (stores all activations), simple implementation - **1F1B (PipeDream, Megatron)**: after warm-up, alternates 1 forward and 1 backward; steady-state memory usage (constant number of stored activations); most common in practice - **Interleaved 1F1B**: each device handles multiple non-consecutive stages; device 0: stages [0, 4, 8], device 1: stages [1, 5, 9]; reduces bubble size by increasing scheduling flexibility - **Chimera**: combines synchronous and asynchronous execution; synchronous within groups, asynchronous across groups; balances consistency and efficiency **Memory Management:** - **Activation Memory**: forward pass stores activations for backward pass; memory = num_micro_batches_in_flight × activation_size_per_micro_batch; 1F1B reduces this compared to fill-drain - **Activation Checkpointing**: stores only subset of activations (e.g., every Nth layer); recomputes others during backward; selective checkpointing balances memory and computation - **Gradient Accumulation**: accumulates gradients across micro-batches; single weight update per mini-batch; maintains effective batch size = num_micro_batches × micro_batch_size - **Weight Versioning (PipeDream)**: stores multiple weight versions for asynchronous execution; memory overhead = num_stages × weight_size; limits scalability to 10-20 stages **Micro-Batch Size Selection:** - **Trade-offs**: smaller micro-batches → more parallelism, less bubble, but more communication overhead; larger micro-batches → less overhead, but more bubble - **Optimal Size**: typically 1-4 samples per micro-batch; depends on model size, stage count, and hardware; profile to find sweet spot - **Bubble Analysis**: bubble_fraction = (num_stages - 1) / num_micro_batches; want bubble < 10-20%; requires num_micro_batches >> num_stages - **Memory Constraint**: micro_batch_size limited by per-stage memory; smaller stages can use larger micro-batches; non-uniform micro-batch sizes possible but complex **Communication Optimization:** - **Point-to-Point Communication**: stage i sends activations to stage i+1; uses NCCL send/recv or MPI; bandwidth requirements = activation_size × num_micro_batches / time - **Activation Compression**: compress activations before sending; FP16 instead of FP32 (2× reduction); lossy compression possible but affects accuracy - **Communication Overlap**: overlaps communication with computation; sends next micro-batch while computing current; requires careful scheduling and buffering - **Gradient Communication**: backward pass sends gradients to previous stage; same volume as forward activations; can overlap with computation **Combining with Other Parallelism:** - **Pipeline + Data Parallelism**: replicate entire pipeline across multiple groups; each group processes different data; scales to arbitrary GPU count - **Pipeline + Tensor Parallelism**: each pipeline stage uses tensor parallelism; enables larger models per stage; Megatron-LM uses this combination - **3D Parallelism**: data × tensor × pipeline; example: 512 GPUs = 8 DP × 8 TP × 8 PP; matches parallelism to hardware topology (TP within node, PP across nodes) - **Optimal Configuration**: depends on model size, hardware, and batch size; automated search (Alpa) or manual tuning based on profiling **Framework Implementations:** - **Megatron-LM**: 1F1B schedule with interleaving; combines with tensor parallelism; highly optimized for NVIDIA GPUs; used for GPT, BERT, T5 training - **DeepSpeed**: pipeline parallelism with ZeRO optimizer; supports various schedules; integrates with PyTorch; extensive documentation and examples - **Fairscale**: PyTorch-native pipeline parallelism; modular design; easier integration than DeepSpeed; used by Meta for large model training - **GPipe (TensorFlow/JAX)**: original implementation; synchronous pipeline with activation checkpointing; less commonly used now (Megatron/DeepSpeed preferred) **Practical Considerations:** - **Load Balancing**: stages should have similar computation time; unbalanced stages create bottlenecks; use profiling to guide layer partitioning - **Stage Granularity**: more stages → better load balance but more bubble; fewer stages → less bubble but harder to balance; 4-16 stages typical - **Batch Size Requirements**: pipeline parallelism requires large batch sizes (num_micro_batches × micro_batch_size); may need gradient accumulation to achieve effective batch size - **Debugging Complexity**: pipeline failures are hard to debug; use smaller configurations for initial debugging; comprehensive logging essential **Performance Analysis:** - **Efficiency Metric**: efficiency = ideal_time / actual_time where ideal_time assumes perfect parallelism; accounts for bubble and communication overhead - **Bubble Overhead**: bubble_time = (num_stages - 1) × (forward_time + backward_time) / num_micro_batches; minimize by increasing num_micro_batches - **Communication Overhead**: depends on activation size and bandwidth; high-bandwidth interconnect (NVLink, InfiniBand) critical; measure with profiling tools - **Memory Efficiency**: pipeline enables training models that don't fit on single GPU; memory per GPU = model_size / num_stages + activation_memory Pipeline parallelism is **the essential technique for training models that exceed single-GPU memory capacity — enabling the distribution of massive models across multiple devices while maintaining reasonable training efficiency through sophisticated scheduling and micro-batching strategies that minimize idle time and maximize hardware utilization**.