SlimPipe: Memory-Thrifty and Efficient Pipeline Parallelism for Long-Context LLM Training

Pipeline Parallelism serves as a crucial technique for training Large Language Models, as it alleviates memory pressure from model states with relatively low communication overhead. However, in long-context scenarios, existing pipeline parallelism methods fail to address the substantial activation memory pressure, primarily due to the peak memory consumption resulting from the accumulation of activations across multiple microbatches. Moreover, these approaches inevitably introduce considerable pipeline bubbles, further hindering efficiency. To tackle these challenges, we propose SlimPipe, a novel approach to fine-grained pipeline parallelism that employs uniform sequence slicing coupled with 1F1B schedule. This method constrains the activation memory accumulation to a single microbatch at any stage, by partitioning the sequence into smaller slices. Although the slices are evenly partitioned, the computation cost is not equal across slices due to causal attention. To counter this, we introduce a sophisticated workload redistribution technique to balance the computation. Consequently, SlimPipe achieves two key benefits simultaneously: (1) near-zero memory overhead and (2) minimal pipeline bubbles. For a 512K context length, SlimPipe boosts Model FLOPS Utilization (MFU) by up to 1.57 × compared to state-of-the-art methods. More notably, it sustains over 45% MFU on 256 GPUs for an extreme context length of 2048K.