通过内存卸载提升流水线并行扩展性，显著降低显存占用。
PipeOffload: Improving Scalability of Pipeline Parallelism with Memory Optimization
- 提出选择性内存卸载策略，动态释放冗余激活数据。
- 在多阶段配置下，单设备显存随阶段数增长而下降，最高降19%。
- 适合大规模语言模型训练，尤其适用于显存受限场景。
流水线并行（PP）广泛用于训练大语言模型（LLMs），但其可扩展性常受高激活内存消耗制约，尤其当流水线阶段数增加时。本文通过实证研究发现，在多数标准配置下，至少一半甚至全部的激活数据可无额外开销地卸载至外部存储。当无法完全卸载时，提出一种新型选择性卸载策略，使峰值激活内存下降幅度优于线性关系。此外，将内存卸载与其他技术协同优化，兼顾整体吞吐与显存限制。实验表明，随着总阶段数增加，单设备激活内存显著降低，使得PP相比张量并行更具优势，最多实现19%加速且显存更低。代码已开源。
原文摘要 · Abstract (English)
Pipeline parallelism (PP) is widely used for training large language models (LLMs), yet its scalability is often constrained by high activation memory consumption as the number of in-flight microbatches grows with the degree of PP. In this paper, we focus on addressing this challenge by leveraging the under-explored memory offload strategy in PP. With empirical study, we discover that in the majority of standard configurations, at least half, and potentially all, of the activations can be offloaded with negligible overhead. In the cases where full overload is not possible, we introduce a novel selective offload strategy that decreases peak activation memory in a better-than-linear manner. Furthermore, we integrate memory offload with other techniques to jointly consider overall throughput and memory limitation. Our experiments proves that the per-device activation memory effectively reduces with the total number of stages, making PP a stronger alternative than TP, offering up to a 19\% acceleration with even lower memory consumption. The implementation is open-sourced at \href{https://github.com/sail-sg/zero-bubble-pipeline-parallelism}{this url}.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。