Penghui Qi*, Xinyi Wan*, Guangxing Huang, Min Lin Sea AI Lab
* Equal Contributions 2025年2月27日
See English Version
Deepseek 在他们的开源周第四天开源了 DualPipe。这是一种为提升训练性能而设计的流水线并行(PP)与专家并行(EP)的协同设计。
在本文中,我们展示了 DualPipe 的 Dual 部分实际上会导致 2 倍的参数冗余,这是不必要的,并且几乎可以无代价移除,对调度的其他属性影响极小。其关键在于通过简单的“对半裁剪”操作将其转换为 V-Shape 调度。进一步地,我们展示了当不需要专家并行(EP)时,可以进一步提高效率,最终得到无流水线气泡的 ZBV 调度。
需要注意的是,DualPipe 调度可以分为两个镜像对称的部分,如下图所示。例如,设备 和 拥有相同的pipeline stage,并且调度方式完全一致。
图 1. DualPipe 可以分为两个镜像对称的部分。
通过仅保留前半部分的设备(并将从下至向上的microbatch的后半段接到从上至下的microbatch的前半段),可以获得一个不含“双向”部分的调度方案,同时保持与原始方案相同的气泡率、内存占用等属性。
图 2. 设备对半裁剪及阶段拼接示意图。
图 3. 对半裁剪调度方案。
我们称该调度方案为对半裁剪调度(Cut-in-half schedule)。可以看到,它不再有参数重复。然而,由于设备数量减少了一半,每个设备的参数内存占用仍然保持不变。根据相同的理念,我们可以通过增加设备数量并减少每个阶段的层数,设计一个与DualPipe设备数量一致的对半裁剪调度方案。由于层数减少,每个设备的参数量降至原来的 50%。
我们在下表中展示了详细的比较,所有方案均基于相同数量的设备(记为 )。
表 1. 各种流水线调度的对比
方法 | 流水线气泡 | 参数 | Activation Memory | PP 通信 |
---|---|---|---|---|
1F1B | (-1)(𝐹+𝐵) | 1× | 1x | |
ZB1P | (-1)(𝐹+𝐵-2𝑊) | 1× | 1x | |
DualPipe | (/2-1)(𝐹&𝐵+𝐵-3𝑊) | 2× | +1 | 1x |
Cut-in-half* | ((-1)/2)(𝐹&𝐵+𝐵-3𝑊) | 1× | +1/2 | 2x |
* 注意: 表示设备数量。对于对半裁剪方案,它指的是裁剪后的设备数量。
对半裁剪方案的 PP 通信量是其他方法的两倍。然而,由于参数内存减少了一半,这种优势弥补了通信开销的增加,因为相较于 EP 通信,PP 通信的开销较小。
对半裁剪调度类似于 V-Shape 调度,我们强调它与以往研究的联系。V-Shape 调度最早由 [3] 提出,随后在 ZB-V 调度 [1] 中优化至零气泡,并在 [2] 中进一步修改以减少内存占用。
对半裁剪(Cut-in-half)中的流水线气泡部分源自对专家并行(EP)的优化。前向和后向计算在稳定阶段重叠,以减少 EP 通信的开销。当不考虑 EP 时,这些气泡可以进一步减少,从而将对半裁剪方案转换为 ZB-V 调度。
解耦F/B并squeeze
如下图所示,将前向传播和反向传播解耦成独立的过程可以带来更灵活的依赖关系,使得可以进一步压缩调度。现在的调度更类似于 ZB-V 调度。
通过在冷却阶段进一步解耦 B/W 并绕过优化器同步(在 [1] 的第 4 节中介绍),我们可以获得 ZB-V 调度,并实现零气泡(zero bubble)。
如果您觉得这篇博客有帮助,请引用:
[1] https://arxiv.org/pdf/2401.10241
[2] https://arxiv.org/abs/2405.15362
[3] https://dl.acm.org/doi/pdf/10.1145/3581784.3607073