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

📄 arXiv: 2504.14519v1 📥 PDF

作者: Zhouyang Li, Yuliang Liu, Wei Zhang, Tailing Yuan, Bin Chen, Chengru Song, Di Zhang

分类: cs.LG, cs.AI

发布日期: 2025-04-20

💡 一句话要点

SlimPipe：面向长文本LLM训练的内存高效流水线并行方法

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 流水线并行 长文本LLM 内存优化 负载均衡 1F1B调度

📋 核心要点

  1. 现有流水线并行方法在长文本LLM训练中面临激活内存压力大和流水线气泡多的挑战，导致效率降低。
  2. SlimPipe通过均匀序列切片和1F1B调度，将累积激活减少到单个切片，并采用负载重分配来平衡计算成本。
  3. 实验表明，SlimPipe在长文本场景下显著提升了模型FLOPs利用率，并在大模型上保持了较高的GPU利用率。

📝 摘要（中文）

流水线并行（PP）是训练大型语言模型（LLM）的关键技术，因为它能够以相对较低的通信开销缓解模型状态带来的内存压力。然而，在长文本场景中，现有的流水线并行方法无法解决显著的激活内存压力，这主要是由于多个微批次中激活的累积导致峰值内存消耗。此外，这些方法不可避免地会引入大量的流水线气泡，进一步阻碍效率。为了解决这些挑战，我们提出SlimPipe，一种细粒度的流水线并行新方法，它采用均匀序列切片和one-forward-one-backward（1F1B）调度。它将累积的激活从多个微批次减少到仅一个，该激活被分成多个切片。虽然切片被均匀划分，但由于因果注意力，切片的计算成本并不相等。我们开发了一种复杂的负载重新分配技术来解决这种负载不平衡。SlimPipe实现了（1）接近零的内存开销和（2）最小的流水线气泡。SlimPipe的有效性已通过对各种模型架构、上下文窗口大小和SlimPipe特定配置的全面测试得到证明。例如，在Llama 70B模型上，与最先进的方法相比，对于512K的上下文长度，SlimPipe显著提高了模型FLOPs利用率（MFU）高达1.57倍。更值得注意的是，对于2048K的上下文长度，它在256个NVIDIA Hopper 80GB GPU上保持超过45%的利用率，而其他方法要么遭受显著的性能下降，要么由于内存限制而完全失败。

🔬 方法详解

问题定义：现有流水线并行方法在训练长文本LLM时，由于激活值在多个微批次间的累积，导致峰值内存占用过高，容易超出GPU的内存限制。此外，传统流水线并行方法还会引入显著的流水线气泡，降低训练效率。因此，如何降低长文本LLM训练中的内存占用，并减少流水线气泡，是本文要解决的关键问题。

核心思路：SlimPipe的核心思路是通过细粒度的流水线并行，结合均匀序列切片和1F1B调度，将多个微批次的激活值压缩到一个切片中，从而显著降低内存占用。同时，通过负载重分配策略，解决因果注意力机制带来的计算不平衡问题，减少流水线气泡。

技术框架：SlimPipe的整体框架包括以下几个主要阶段：1) 输入序列被均匀切分成多个切片；2) 这些切片在流水线中的不同阶段并行处理，采用1F1B调度；3) 针对因果注意力机制导致的计算不平衡，进行动态的负载重分配，确保每个GPU的计算负载相对均衡。

关键创新：SlimPipe的关键创新在于：1) 细粒度的流水线并行，通过均匀序列切片和1F1B调度，有效降低了激活内存占用；2) 负载重分配策略，解决了因果注意力机制带来的计算不平衡问题，减少了流水线气泡。与现有方法相比，SlimPipe能够在长文本场景下实现更高的内存效率和计算效率。

关键设计：SlimPipe的关键设计包括：1) 均匀序列切片的大小，需要根据GPU内存和计算能力进行调整；2) 负载重分配策略的具体实现，需要考虑不同切片的计算复杂度，并动态调整任务分配；3) 1F1B调度的具体实现，需要保证前向和反向计算的正确依赖关系。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SlimPipe在Llama 70B模型上，对于512K的上下文长度，模型FLOPs利用率（MFU）提升高达1.57倍。更重要的是，对于2048K的上下文长度，SlimPipe在256个NVIDIA Hopper 80GB GPU上保持超过45%的利用率，而其他方法要么性能显著下降，要么因内存限制而完全失败。这些结果验证了SlimPipe在长文本LLM训练中的有效性。

🎯 应用场景

SlimPipe可广泛应用于长文本LLM的训练，例如用于训练具有超长上下文窗口的语言模型，从而提升模型在长文档理解、对话生成、代码生成等任务上的性能。该方法能够降低训练成本，加速模型迭代，并使得在资源受限的硬件上训练大型模型成为可能，具有重要的实际应用价值。

📄 摘要（原文）

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.