前三篇文章分别介绍了 EP、DP、TP:
接下来会尽量做到由浅入深的介绍 MP 中的 PP,既 流水线并行策略。
Naive PP
首先 PP 即 流水线并行,是将模型 按层 进行拆分,与 DP(切数据) 和 TP(层内切张量) 进行互补。
最朴素 Naive的做法是:
- 将模型按
层(layer)切分为多个阶段(stage),每个stage(可以包含多个layer)部署到不同设备

- 单个
batch的前向传播按stage顺序执行,如下表所示,GPU0执行完本设备stage的前向传播后,将中间变量激活值传给GPU1,GPU1再执行,直到GPU3完成整个batch的前向传播 - 计算得到
loss后,立即开始反向传播,GPU3执行反向传播得到梯度后,GPU2再执行,直到GPU0完成整个batch的反向传播 - 所有设备都得到各自
梯度后,再统一进行参数更新
| timeline | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 |
|---|---|---|---|---|---|---|---|---|---|
| GPU3 | FWD | BWD | UPDATE | ||||||
| GPU2 | FWD | BWD | UPDATE | ||||||
| GPU1 | FWD | BWD | UPDATE | ||||||
| GPU0 | FWD | BWD | UPDATE |
可以看到 Navie PP 有两个主要问题:
-
同一时刻,其实只有一个设备在进行计算(和单GPU没有任何区别),其余设备
均处于空闲状态,如下图所示,有大量空闲状态,也就形成了
气泡,资源利用率极低

-
数据生命周期被拉长,导致显存占用率高,如下图所示,为了在反向传播过程中完成计算,不得不将激活值等数据进行缓存;例如对于GPU0,要等待几乎整轮前向+反向传播完成,才可以释放这部分显存

GPipe
GPipe论文:https://arxiv.org/abs/1811.06965
GPipe 针对上面 Naive PP 的主要问题分别提出了两个解决方案:
(1)针对气泡问题,也就是GPU计算空闲,提出了micro-batch方案,也就是把mini-batch拆分成更小的micro-batch
下面以拆成4份的micro-batch为例,当GPU0完成第一个微批次得到激活值之后,会同时做两件事:①首先会立即把激活值传递给GPU1,这样GPU1就可以立即开始进行前向传播 ②同时GPU0也会立即进行第二个微批次的计算;以此类推,这样极大的提高了设备的并行度(例如在时刻3时,4个设备都在计算),降低了设备空闲时间,减小了气泡(反向传播也是一样)
| timeline | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| GPU3 | F1 | F2 | F3 | F4 | B4 | B3 | B2 | B1 | UPDATE | ||||||
| GPU2 | F1 | F2 | F3 | F4 | B4 | B3 | B3 | B1 | UPDATE | ||||||
| GPU1 | F1 | F2 | F3 | F4 | B4 | B3 | B2 | B1 | UPDATE | ||||||
| GPU0 | F1 | F2 | F3 | F4 | B4 | B3 | B2 | B1 | UPDATE |
相比于Naive PP的气泡,GPipe的气泡小了很多

(2)针对数据生命周期被拉长,导致显存占用率高的问题,提出了重计算方案,也就是把激活值丢掉 ,在需要用的时候重新计算(也就是在backward的时候,先计算一遍激活值)
PipeDream
PepeDream论文:https://arxiv.org/pdf/1806.03377
虽然GPipe对气泡问题进行了优化,但也只能说是在一定程度上 减小了气泡,无法进一步减少甚至消除气泡,主要是受限于F-then-B(先forward然后再backward然后聚集梯度再更新参数)这种架构。
现在看看 PipeDream 如何进一步解决 气泡 问题:
首先,如果把前向传播和反向传播想象轨道上的火车,那么不论Naive PP(一辆)还是GPipe(一批)其实都是先开过去再开回来(阶段分离);
但PipeDream提出了另一种方式,也就是 1F1B,通过调整轨道上火车 的调度策略 让前向与反向传播的小火车,可以持续的相向而行(打破阶段分离);
下面是 PipeDream 论文中的原图,可以看到,1F1B 的核心思想就是 F后紧接着执行 B 然后紧接着执行 F 然后 B 也就是 1F1B1F1B1F...,尽最大可能的减少了 气泡。

整个调度策略可以分为两个阶段:
startup阶段:和GPipe的micro-batch类似steady阶段:F后立即进行B,一段时间后就可以完全跑满GPU,几乎没有任何气泡
但是问题来了,steady 都乱成这样了,那什么时候聚集梯度、怎么更新参数呢?
PipeDream 设计了两个东西:weight stash 和 vertical sync
先看一下 weight stash,对于每个GPU来说,它独立处理一个stage,但 1F1B 的困境在与,某个时刻不同GPU在同时处理不同 mini-batch 的 前向 或 反向传播;
其中 反向传播 依赖 前向传播 的 激活值;而 激活值 又依赖于 前向传播 时的 模型weight!
但如果不管不顾更新参数的话,很可能出现 反向传播 时,其 模型参数 已经被更新过了,这样算出来的 梯度 就完全不对了,所以根本问题就是 权重 weight的版本不一致。
所以weight stash做的事情就是,在每个mini-batch开始前,对其分配一个version版本号,这样该mini-batch后续的所有操作,都建立在整个version上,同时对于stash 暂存 一份 这个版本的权重weight,当该mini-batch进行反向传播时,使用对应版本 的权重weight;
所以,weight stash 实现了 前向/反向 传播过程中中的 一致性,可以比喻成 水平方向 的 一致性。
而另一个维度 垂直方向,也就是多GPU间 参数更新的 一致性,就需要靠 vertical sync 确保 同一个 mini-batch 的数据,即使在不同设备上,也要使用 相同版本 的 参数。
其底层也通过 mini-batch 进入 stage-0 时的 版本号 实现的,该 版本号 会绑定到该 mini-batch 的 整个生命周期。
Megatron-LM
Megatron-LM的论文:https://arxiv.org/pdf/2104.04473
PipeDream 的 1FB 已经很牛了,但是大佬们依然要求不满足,Megatron-LM 又提出了更牛的 真·交错式 终极体 1F1B。
主要解决了啥问题呢,如下图所示,虽然四个GPU实现了1F1B,单其实在某些时刻,特别是startup 阶段,依然存在着 等待依赖;
相比于 PipeDream 将模型各层,连续的分给各个 stage,例如:stage1包括1/2/3/4层,stage2包括5/6/7/8层;
Megatron-LM采取交错式分层,例如:stage1包括1/2/7/8层,也就是 打破了分层的连续性

而带来的好处就是,正常情况下,7/8 需要在 GPU1 等待它的 5/6 算完后才能进行,但此时 GPU0 一直在空闲的,所以当 GPU1 执行完 5/6 后,由 GPU0 立即进行 7/8 的计算,这样就减少了2个GPU等待时间。
此外,还有很多
PP策略,例如:1F1B-Flush、DeepSeek开源周推出的DualPipe、Chimera等等等等,有时间再研究吧,学不完根本学不完...