1. 概述

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

本文主要介绍：

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

2. 朴素流水线并行

2.1. 流水线分布式训练的步骤

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

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

图 1：朴素流水线和 GPipe 流水线比较

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

3. GPipe 流水线并行

3.1. 优化策略

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

3.2. 并行气泡定量计算

假设气泡总时间为，在一个 batch 中微批次的数量为 m，可用的训练设备有 p 个，每一轮迭代理想的时间为，单个微批次的前向和反向时间分别为。

一次迭代中，总的并行气泡时间，共包含 p-1 个前向和反向时间，所以总的气泡时间为：

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

并行气泡比率为：

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

图 2：GPipe 微批次处理流程

图中一个 batch 拆分成了 8 个微批次 microbatch，即 m=8，p=4。

3.3. GPipe 策略流程

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

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

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

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

4. PipeDream 流水线并行

4.1. 1F1B 并行策略

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

4.1.1. 名词概念

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

4.1.2. 步骤

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

4.1.3. 限制

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

图 3：PipeDream 流水线示意图

4.1.4. 为什么叫做 1F1B 策略？

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

4.1.5. 总结

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

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

4.2. 权重更新问题

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

4.2.1. 问题 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 号之后更新的参数

4.2.2. 问题 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 两种技术

4.3. weight stashing 权重暂存

4.3.1. 总体策略

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 就是。因此 Worker 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 的后向更新使用

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

4.4. vertical sync 纵向同步

4.4.1. 总体策略

每个 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 的存在，所有的参数都是有效的。

4.5. 交错式

4.5.1. 优化策略

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

图 6：1F1B 流水线交错式

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

5. 总结

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

方法	气泡时间	内存占用	通信开销
GPipe (F-then-B)	较高	极高 m 个	低
1F1B- 非交错式	相同	较低 p 个	低
1F1B- 交错式	极低	中等	高

6. 参考文献

[1] arXiv:1811.06965

[2] arXiv:2104.04473

[3] PipeDream: Generalized Pipeline Parallelism for DNN Training