GPipe & PipeDream 让流水线告别卡顿 “纵享丝滑”

大模型训练横向扩展方式，是组合多个GPU设备内存，以承载训练大模型的内存需求。其中典型的代表流水线并行，将模型按层组合成stage，分配到不同的GPU设备上，通过层间交互完成前向和反向计算，以及参数更新，实现大规模模型训练。

1）流水线并行的基本方式和空泡率的推导 2）GPipe中的F-then-B策略和PipeDream优化内存的1F1B策略 3）多个stage的流水线并行导致参数失配问题和解决方案介绍

1，朴素流水线并行

流水线分布式训练的步骤：

模型参数按层组合成stage，分配到不同的计算设备，每个设备上需要计算的算子为 。

• • 前向计算：输入数据 经过 得到结果，其计算结果作为 的输入，依次进行前向计算，最终在device3设备上得到最后损失。
• • 反向传递：损失开始后向传递，并需要每个设备对应的前向结果计算梯度。在设备 Device 0上完成整个梯度计算。
• • 参数更新：同时更新每个设备上的参数，更新完成后，再进行新一轮的迭代。

图1，朴素流水线和GPipe流水线比较。

问题： 当设备Device 0 计算完成 之后，就一直空闲，直到反向传递的最后一步才又开始运行，其他时间都在空跑。同理每个设备都存在严重的资源浪费。

2，GPipe 流水线并行

优化： GPipe通过拆分batch为更小的微批次，并行微批次的前向和反向过程，进而减少空泡率。

• 并行气泡定量计算：假设气泡总时间为 ，在一个batch中微批次的数量为 m，可用的训练设备有 p个， 每一轮迭代理想的时间为 ，单个微批次的前向和反向时间分别为 。
• 一次迭代中，总的并行气泡时间，共包含 p-1 个前向和反向时间，所以总的气泡时间为 ：

一个batch理想的处理时间为：

并行气泡比率为：

理想情况： 要想减少模型气泡时间，就要使一个batch中的微批次数量 m 尽可能的大于设备数量 p。

图2，GPipe 微批次处理流程。

图，GPipe流水线流程示意图。蓝色块表示前向计算，绿色块表示反向计算，灰色块表示设备空闲。图中一个batch拆分成了8个微批次 microbatch，即m=8，p=4。

GPipe策略流程：

• • 第一个微批次，在Device1 完成前向计算后，将计算结果传递给 Device2， 同时Device1 启动第二个微批次的前向计算。
• • 最终在Device 4 完成所有微批次的前向计算，获得损失loss，并依次开始每次微批次的反向梯度计算。
• • Device 4 首先完成所有微批次的梯度计算，并取得平均梯度，更新对应stage的模型参数。同理Device 3，2，1 完成梯度计算后，更新负责stage的参数。

显著的特点：F-then-B，全部的微批次前向计算完成后，再开始反向计算，并收集所有微批次反向梯度后，进行参数更新。

问题： 但是当一个batch 中的微批次过多时，会导致额外的内存占用，因为需要保存m个前向的中间结果，用于后向计算梯度。

那怎样在使用多个微批次的同时，保持内存可控？这就是PipeDream中1F1B要优化的核心。

3，PipeDream 流水线并行

3.1，1F1B 并行策略

优化目标：将微批次梯度计算过程提前，梯度计算完成后，相应的前向激活值就可以丢弃，从而节省内存占用。

名词概念：

• • Startup State 启动状态，预热阶段。steady State 稳定状态，流水线满负荷运行。
• • Worker N 训练设备。蓝色块前向计算，绿色块反向计算，条纹块GPU空闲。反向计算消耗时间约为前向两倍。
• • input stage 包含输入层的阶段，即输入阶段，对应分配到worker 1 上的模型层。output stage 包含输出层的阶段，即输出阶段，对应分配到worker 4 上的模型层。

步骤：

• • 在启动状态，input stage先读入足够多微批次数据（数量一般为流水线深度），以保证在稳定阶段时，各个GPU上都有足够的微批次数据进行计算。图中就是 input stage 发送四个小批次传播到 output stage。
• • 一旦 output stage 完成第一个小批次的前向传播（即worker 4 第一个蓝色块 1），它就对同一微批次批次执行后向传播（worker 4 的第一个绿色 1）。
• • 然后开始交替执行后续微批次的前向传播和反向传播（即 worker 4 的 2前，2后，3前，3后.....）。
• • 当反向传播过程开始传播到管道中的早期阶段时（即work 3 ~ work 1），每个阶段开始在不同小批次的正向和反向过程之间交替进行。
• • 在稳定状态下，每台机器都忙着对一个微批次进行正向传播或反向传播。

限制：

将飞行中的微批次（ in-flight micro batches，即正在进行反向传播且需要保持激活状态的微批次）的数量限制为流水线的深度，即设备的数量，而不是一个批次中的微批次数量。

图3，PipeDream 流水线示意图。

为什么叫做 1F1B 策略？

图中，第5个微批次前向计算结束后，预热阶段结束，开始进入稳定阶段。在这个状态下，每个Worker GPU设备，依次执行一个前向计算 1F 和一个反向传播 1B。

总结：

• • 1F1B对于一个batch 的训练时间和GPipe 是相同的，没有减少并行气泡的时间，但是占用的内存相比GPlpe 减少了。
• • 1F1B要求保存的激活状态对应于 p 个，相比之下GPipe 需要保存 m 个，这样减少了内存占用。

因此，当m远大于p时，虽然PipeDream空泡率与GPipe 相同，但更加节省内存。

3.2，权重更新问题

流水线并行，会引入两种类型的参数失配问题：

问题1：同一个microbatch的前向计算和反向传播使用的参数不一致。

原生的PipeDream流水线中，每个阶段的前向传播都是使用某一个版本的参数来执行，而其反向传播则是使用不同版本的参数来执行的。如：

• • 当 microbatch 5 进入到 worker 1 时，它的前向传播逻辑在 microbatch 1 的后向传播计算之后执行，即它前向传播计算时候使用的参数是 microbatch 1 后向传播计算之后更新的参数。
• • 但是 microbatch 5 后向传播逻辑是在 "microbatch 2, microbatch 3, microbatch 4" 执行完后才开始计算，即此时使用的参数是"microbatch 1, microbatch 2, microbatch 3, microbatch 4" 后向传播计算之后更新的参数。这就导致 microbatch 5 的前向计算和后向计算时候，使用的参数不一致。
• • 理想目标是第一行worker 1 蓝色 5 号和绿色 5 号计算时候，必须都使用绿色 1 号之后更新的参数。

问题2：同一个microbatch在不同stage做同样操作（同样做前向操作，或者同样做后向传播）使用的参数版本不一致。

• • 对于 microbatch 5 在 worker 1 上的前向计算部分（蓝色5），它的前向逻辑在 microbatch 1 的后向计算以后执行。
• • 但是 microbatch 5 在 worker 2 上的前向计算部分（蓝色5），是在 "microbatch 1, microbatch 2" 的后向计算结束后才执行。
• • 这就导致了 microbatch 5 在两个stage上前向计算使用的参数版本不一致。

为解决这两个问题，PipeDream 分别采用了 weight stashing 和 Vertical Sync 两种技术

3.3，weight stashing 权重暂存

总体策略：

Weight stashing为权重维护多个版本，每个active microbatch都有一个版本。每个stage 都用最新版本的权重进行前向计算，处理输入的microbatch。计算前向传播之后，会将这份参数保存下来用于同一个microbatch的后向计算。Weight stashing确保在一个阶段内，相同版本的模型参数被用于给定小批量的向前和向后传播，但是不能保证跨阶段间，一个给定的小批次使用模型参数的一致性。

图4，权重暂存策略示意图。

符号定义： 其中w表示权重参数，i 表示微批次，n表示计算设备。比如 表示微批次2在worker 2上更新后的参数。

以microbatch 5 分别在每个worker上完成前向和后向计算时，需要使用的参数为例进行说明。

具体：

• • Worker 1 第一行的蓝色 5 依赖于它前面同一行的绿色 1。Worker 1 所在行的第一个绿色 1 结束时，代表了 microbatch 1 完成了本次流水线的 4 次前向传播，4次后向传播，所以是一个新版本的 weight 1 就是 。因此Work 1 的两个 microbatch 5（蓝色前向和绿色后向）都应该基于新版本 计算。因此需要记录下来 新版本 。
• • Worker 2 第二行的蓝色 5 依赖于它前面同一行的绿色 2。同理，Worker 1 的第一个绿色 2 结束时，代表了 microbatch 2 完成了本次流水线的 4 次前向传播，4次后向传播，所以是一个新版本的 weight 2。此时的 microbatch 5 （前向和图上未标出的绿色后向）都应该基于新版本的 weight 2 计算，因此需要记录下新版本。
• • Worker 3 第三行的蓝色 5 依赖于它前面同一行的绿色 3。当执行 microbatch 5 的前向传播时候，已经更新，所以需要记录下来 ，为了以后 microbatch 5 的后向更新使用。
• • Worker 4 第四行的蓝色 5 依赖于它前面同一行的绿色 4。当执行 microbatch 5 的前向传播时候，使用所以需要记录下来 ，供 microbatch 5 的后向更新使用。

依次类推，当进行mircobatch 6,7,8,9 ... 等以后的批次时，worker 1 需要记录 ，，， 的每一个新版本。就是 worker 1 对应microbatch 1,2,3,4 的各个权重。

3.4，vertical sync 纵向同步

总体策略：

每个microbatch进入pipeline 时都使用输入stage最新版本的参数，并且参数的版本号会伴随该microbatch数据整个生命周期，在各个阶段都是用同一个版本的参数，而不是Weight stashing那样都使用最新版本的参数，从而实现了stage间的参数一致性。

图5，vertical sync 策略介绍。

强制所有worker在计算 microbatch 5 的时候都用本worker做 microbatch 1 反向传播之后的参数，具体来说就是：

• • 对于 worker 2,3,4 都使用本阶段绿色1（1反向传播之后，更新的本阶段权重）来做 5 的前向传播。

但是，这样同步会导致很多计算浪费无用。比如5更新时用的1的权重，但2，3，4后向传播的权重都白白计算了，所以默认不使用Vertical Sync。这样虽然每层不完全一致，但是由于weight stashing的存在，所有的参数都是有效的。

3.5，交错式

优化：为了进一步的减少并行气泡时间，每个设备不仅是分配单一的stage而是多个model chunk。

图6，1F1B流水线交错式。

• • 策略：之前4个设备，每个设备有4层模型参数，即 device0：layer0-3， device1：layer4-7，等等。现在让每个设备有两个模型块 model chunks，每个chunk 有两层，即设备0：layer0-1 and layer 9-10. device1：layer 2-3 and layer 11-12，其他的设备也是如此。
• • 优化结果： 使得并行气泡降低为原来的1/N，提高了GPU的利用率，但是带来的代价是通信量的增加，需要GPU之间高带宽方式连接。

4，总结

• • GPipe 采用 “all-forward, all-backward"的策略，需要保留所有微批次的前向结果，用于反向梯度计算，造成内存消耗高。
• • 1F1B 交错式策略，通过减少保存微批次激活值优化内存占用。即在完成一个微批次 前向计算后，立即开始进行反向传播，且反向传播完成后，即可丢弃前向的中间激活值，进而节省了激活值内存。但两者完成一轮参数更新的时间是相同的，并行气泡并未较少。
• • 1F1B 交错式流水线，通过让设备负责不同的模型块，缩小的并行气泡的时间，但通信量增加了。

参考：

[1] arXiv:1811.06965 [2] arXiv:2104.04473 [3] PipeDream: Generalized Pipeline Parallelism for DNN Training