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

TL;DR

This work tackles the memory and throughput bottlenecks of training long-context LLMs with pipeline parallelism. It introduces SlimPipe, a fine-grained scheme that uses uniform sequence slicing coupled with a 1F1B schedule and attention-context exchange to drastically reduce activation memory and pipeline bubbles, while distributing the workload across devices. The approach yields near-zero memory overhead and maintains high MFU, achieving up to $1.57\times$ MFU improvements at 512K context and preserving substantial throughput at $2048K$ contexts across hundreds of GPUs, with offloading extending the practical context length further. These results demonstrate SlimPipe’s ability to enable efficient long-context training without full checkpointing or prohibitive cross-node communications, significantly advancing scalable LLM training for ultra-long contexts.

Abstract

Pipeline Parallelism (PP) serves as a crucial technique for training Large Language Models (LLMs), owing to its capability to alleviate 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 one-forward-one-backward (1F1B) schedule. It reduces the accumulated activations from several microbatches to just one, which is split into several slices. Although the slices are evenly partitioned, the computation cost is not equal across slices due to causal attention. We develop a sophisticated workload redistribution technique to address this load imbalance. SlimPipe achieves (1) near-zero memory overhead and (2) minimal pipeline bubbles simultaneously. The effectiveness of SlimPipe has been proven by thorough testing with diverse model architectures, context window sizes, and SlimPipe-specific configurations. For example, on the Llama 70B model, compared to state-of-the-art methods, SlimPipe significantly boosts the Model FLOPs Utilization (MFU) to up to $1.57\times$ for a context length of 512K. More notably, for a context length of 2048K, it maintains over 45% utilization on 256 NVIDIA Hopper 80GB GPUs, while other approaches either suffer significant performance drops or fail entirely due to memory constraints.

Figures (14)

• Figure 1: Comparison of GPU memory footprint between Classic PP huang2019gpipenarayanan2021memorynarayanan2021efficientqi2024zeroli2021terapipe and SlimPipe across various pipeline parallelism sizes. While both approaches consume identical GPU memory for model parameters, SlimPipe's activation memory decreases proportionally as pipeline parallelism size increases, in contrast to Classic PP's constant activation memory requirements.
• Figure 2: Maximum context lengths supported by different PP schemes in training Llama 7B with 8-way TP and 8-way PP.
• Figure 3: Theoretical bubble fractions of different PP schemes (PP size 8) in training Llama 13B with 4 microbatches and a 256K context length.
• Figure 4: The top figure shows the default 1F1B schedule. The bottom figure shows the SlimPipe schedule, where each microbatch is broken into multiple (in this case, 8) slices. $M_a$ represents the total activation size of one microbatch. $n$ and $p$ indicate the number of slices per sequence and the PP size, respectively. The activation memory accumulated during the warm-up phase on device 1 reduces from $M_a$ to $(1+\delta)\frac{M_a}{p}$, where $\delta=2(p-1)/n$. In the mean time, the warm-up bubble shrinks by about $n$ times.
• Figure 5: SlimPipe in its interleaving form. Each device is assigned 2 stages. Dark colors show the first stage and light colors show the second stage. For simplicity, the microbatch number is not labeled and the number in the box is the slice number. Note that there are only two microbatches in this pipeline, each split into 8 slices. Whereas in the classic interleaved 1F1B pipeline, at least 4 (the PP size) microbatches are required. More details are depicted in Figure \ref{['fig:imba_bubble']} and Figure \ref{['fig:post_bubble']}.