动态调整流水线并行粒度，提升长文本大模型训练效率
InfiniPipe: Elastic Pipeline Parallelism for Efficient Variable-Length Long-Context LLM Training
- 融合令牌级与批处理级流水线，动态适应不同长度序列
- 在真实数据集上实现1.69倍于现有系统的训练加速
- 适合需要高效训练长上下文大模型的研究者使用
长上下文训练对大语言模型的上下文扩展至关重要。现有方案如序列并行会带来显著通信开销。流水线并行（PP）可降低该开销，但其效率依赖于划分粒度。批级PP结合序列打包虽内存占用高，而令牌级PP将序列切分为片段虽缓解内存压力，却可能导致硬件利用率不足。此外，真实数据集中序列长度分布不均，使固定粒度的流水线并行性能不佳。本文提出：1）弹性流水线并行（EPP），协调令牌级与批级并行以适应资源与负载异构性；2）面向阶段的分块自适应检查点技术，高效集成梯度检查点与EPP。全面实验表明，InfiniPipe相比最先进系统实现1.69倍加速。代码已开源：https://github.com/wsjdsg/InfiniPipe-code.git。
原文摘要 · Abstract (English)
Long context training is crucial for LLM's context extension. Existing schemes, such as sequence parallelism, incur substantial communication overhead. Pipeline parallelism (PP) reduces this cost, but its effectiveness hinges on partitioning granularity. Batch-level PP employing sequence packing exhibits high memory consumption in long-context scenarios, whereas token-level PP splitting sequences into slices alleviates memory overhead but may incur hardware under-utilization. Moreover, the skewed distribution of sequence length in real-world datasets renders monolithic and static granularity PP's sub-optimal performance. In this paper, we propose 1) \textit{Elastic Pipeline Parallelism} (EPP) that orchestrates token-level PP and batch-level PP to adapt to resource and workload heterogeneity, and 2) \textit{Stage-Aware Chunk-Level Adaptive Checkpointing} that efficiently integrates gradient checkpointing with EPP. Comprehensive experiments demonstrate that InfiniPipe achieves a 1.69x speedup over state-of-the-art systems. Our code is open-sourced at https://github.com/wsjdsg/InfiniPipe-code.git.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。