这个代码库是 Megatron-LM 的一个分支。原始的 README 可以在这里找到。

零气泡流水线并行和可控内存流水线并行

零气泡流水线并行是一种新型流水线并行算法，能够将流水线并行的气泡几乎减少到零，同时保持同步语义。

可控内存流水线并行是一种新方法，用于构建具有可控激活内存的流水线并行调度。使用这种方法，我们可以显著降低流水线并行的激活内存消耗，同时保持相同的吞吐量，甚至更快。

查看我们的论文：

零气泡调度器的演示平台：

启用零气泡的快速设置：

  --zero-bubble-v-schedule
  --allow-padding-num-layers
  --enable-optimizer-post-validation

也可以尝试使用 ZERO_BUBBLE_V_SCHEDULE=1 examples/pretrain_zero_bubble.sh

或添加另一个标志来控制内存消耗或 V 调度：

  --zero-bubble-v-schedule-mem-setup half

为您自己的 Megatron 分支启用 ZB H1 调度的轻量级替代选项

# 通过 pip install zbpp_light 安装
import zbpp_light
zbpp_light.patch_megatron()

import megatron
...

推进吞吐量和内存的帕累托前沿

我们的一系列调度推进了吞吐量和内存的帕累托前沿。

调度

实现零气泡的关键是将反向传播分解为 $B$ 传递和 $W$ 传递。一个阶段的 $B$ 只依赖于其下一个阶段的 $B$，而不是像 1F1B 那样依赖于 $B$ 和 $W$。

通过控制每个构建块的生命周期，我们可以控制和降低激活内存。

1F1B
ZB1P
ZB2P
ZBV - 每个设备被分配恰好 2 个块（虚拟阶段），其中白色文本颜色表示第一个块，黑色文本颜色表示第二个块。模型块之间的依赖序列在正向和反向传递中都遵循 "V" 形模式。

	1F1B	ZB1P	ZB2P	ZBV	V-Half	V-Min
气泡率	$(p-1)/(m+p-1)=B$	$B/3$	0	0	$B/2$	$2B/3 + O(n)$ 开销
激活内存（相比1F1B）	1倍	1倍	2倍	1倍	1/2倍	1/3倍
流水线通信量（相比1F1B）	1倍	1倍	1倍	2倍	2倍	2倍

--enable-zero-bubble 启用零气泡调度。
--zero-bubble-v-schedule 启用上述推荐的V调度。隐含 --enable-zero-bubble。
--zero-bubble-v-schedule-mem-setup 设置V调度的内存限制，有效选项为 min/half/zb（默认）。
--enable-optimizer-post-validation 启用在优化器后验证中解释的优化器后验证。
--allow-padding-num-layers 允许层数不是流水线数的倍数。这允许我们在第一个和最后一个流水线阶段减少一层，以补偿嵌入层造成的气泡。
--zero-bubble-max-pending-backward 控制零气泡调度的内存限制。将其设置为1倍流水线数将得到类似ZB1P的调度，而设置为2倍流水线数将得到ZB2P。对由 --zero-bubble-v-schedule 启用的ZBV调度无效。
--zero-bubble-pipeline-timers-start-iter 和 --zero-bubble-pipeline-timers-end-iter 用于控制ZB调度器分析每个F/B/W以测量 $T_F$、$T_B$ 和 $T_W$ 的起始/结束迭代。

注意事项

V调度要求每个流水线的层数为偶数，以便每个阶段可以均匀地分成两个虚拟阶段。
为了获得更好的吞吐量，我们建议将 --num-layers 设置为 k * pipeline-model-parallel-size - 2，其中k可以是任何 $\ge1$ 的值。这用于补偿第一个/最后一个流水线阶段上的额外嵌入层，否则可能会给所有其他阶段带来气泡。