Pipeline Parallelism with Controllable Memory
作者: Penghui Qi, Xinyi Wan, Nyamdavaa Amar, Min Lin
分类: cs.LG, cs.CL, cs.DC
发布日期: 2024-05-24 (更新: 2024-11-04)
🔗 代码/项目: GITHUB
💡 一句话要点
提出可控内存的流水线并行框架，显著提升大模型训练吞吐量。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 流水线并行 深度学习 大模型训练 内存优化 吞吐量提升
📋 核心要点
- 现有流水线并行调度方法缺乏系统性设计，导致内存利用率不高，成为制约训练效率的关键因素。
- 论文提出一种基于构建块分解的流水线调度框架，通过控制构建块的生命周期来优化激活内存。
- 实验表明，该方法在纯流水线并行和混合并行场景下均能显著提升吞吐量，最高提升达55%。
📝 摘要（中文）
本文提出了一种流水线并行框架，该框架将流水线调度分解为重复的构建块。研究表明，构建块的生命周期决定了流水线调度的峰值激活内存。通过分析发现，现有的大部分流水线调度方法在内存利用率方面效率低下。为了解决这个问题，本文引入了一系列具有可控激活内存的高效构建块，可以在不牺牲效率的情况下将峰值激活内存降低到1F1B的1/2，甚至在吞吐量相当的情况下降低到1/3。同时，可以在保持与1F1B相同激活内存的情况下，实现几乎零流水线气泡。评估结果表明，在纯流水线并行设置中，本文方法在吞吐量方面优于1F1B，提升幅度为7%到55%。在实际场景中，通过对混合并行超参数进行网格搜索，本文方法在大型语言模型上的吞吐量比1F1B基线提高了16%。该实现已开源。
🔬 方法详解
问题定义：现有的流水线并行调度方法，如1F1B，在内存效率方面存在不足。激活内存占用高，限制了模型规模和训练效率。缺乏系统性的调度方法，难以根据硬件资源和模型特性进行优化。
核心思路：论文的核心思路是将流水线调度分解为重复的构建块，并发现构建块的生命周期与峰值激活内存之间存在直接关系。通过设计具有可控激活内存的高效构建块，可以在保证吞吐量的前提下，显著降低内存占用。
技术框架：该框架首先将流水线调度分解为一系列重复的构建块。然后，分析构建块的生命周期，确定其对峰值激活内存的影响。接着，设计新的构建块，通过控制其生命周期来降低激活内存。最后，将这些构建块应用于实际的流水线并行训练中。
关键创新：最重要的技术创新在于提出了一种新的流水线并行调度方法，该方法通过控制构建块的生命周期来实现可控的激活内存。与传统的1F1B方法相比，该方法可以在不牺牲吞吐量的情况下，显著降低内存占用，或者在相同内存占用下，实现更高的吞吐量。
关键设计：论文设计了一系列具有不同生命周期的构建块，允许用户根据实际需求选择合适的构建块。例如，设计了可以降低激活内存到1F1B的1/2甚至1/3的构建块，以及可以实现几乎零流水线气泡的构建块。具体的参数设置和损失函数等细节未在摘要中提及，属于未知信息。
🖼️ 关键图片
📊 实验亮点
实验结果表明，在纯流水线并行设置中，该方法在吞吐量方面优于1F1B，提升幅度为7%到55%。在实际场景中，通过对混合并行超参数进行网格搜索，该方法在大型语言模型上的吞吐量比1F1B基线提高了16%。这些结果表明，该方法在实际应用中具有显著的优势。
🎯 应用场景
该研究成果可广泛应用于大规模深度学习模型的训练，尤其是在资源受限的环境下。通过降低内存占用，可以支持更大规模的模型训练，或者在相同的硬件资源上实现更高的训练效率。该方法对于自然语言处理、计算机视觉等领域的预训练模型和微调任务具有重要意义。
📄 摘要（原文）
Pipeline parallelism has been widely explored, but most existing schedules lack a systematic methodology. In this paper, we propose a framework to decompose pipeline schedules as repeating a building block, and show that the lifespan of the building block decides the peak activation memory of the pipeline schedule. Guided by the observations, we find that almost all existing pipeline schedules, to the best of our knowledge, are memory inefficient. To address this, we introduce a family of memory efficient building blocks with controllable activation memory, which can reduce the peak activation memory to 1/2 of 1F1B without sacrificing efficiency, and even to 1/3 with comparable throughput. We can also achieve almost zero pipeline bubbles while maintaining the same activation memory as 1F1B. Our evaluations demonstrate that in pure pipeline parallelism settings, our methods outperform 1F1B by from 7% to 55% in terms of throughput. When employing a grid search over hybrid parallelism hyperparameters in practical scenarios, our methods demonstrate a 16% throughput improvement over the 1F1B baseline for large language models. The implementation is open-sourced at https://github.com/sail-sg/zero-bubble-pipeline-parallelism.