GPT-3模型为何难以复现？这也许是分布式AI框架的最优设计（4）

发布人：AI科技大本营时间：2021-05-15 来源：工程师发布文章

为什么分布式深度学习框架要像 OneFlow 这样设计？

上一个章节，我们从用户角度分析和比较了 OneFlow 和 PyTorch（Megatron）的分布式易用性。这一章节我们会从框架设计和开发者的角度解释，为什么我们这一套 Placement + SBP 的设计是分布式训练更本质的设计。其他框架和高级定制用户的在所有分布式并行上的努力，其实都是 SBP 下的一个特例而已。

1. OneFlow 如何实现流水并行？

OneFlow 的运行时 Actor 机制有以下几个特点：

天然支持流水线， Actor 通过内部的状态机和产出的 Regst 个数以及上下游的 Regst 消息机制解决了流控问题（Control Flow）。Actor 的状态机如下图所示：

Actor 状态机

Actor 组成的计算图运行时调度是去中心化的，每个 Actor 当前是否可以执行都仅与自己的状态、空闲 Regst 数量以及收到的消息有关。

所以使用 Actor 做流水并行，本身就不需要自己定制复杂的调度逻辑。我们可以先举一个数据加载的 Pipeline 示例，当一个由 Actor 组成的数据预处理流程如下图所示时（我们可以将各个阶段约减为一个 Actor）：

数据预处理流程

当这4个Actor之间的 RegstNum 均为2时，如果训练时间比较长（训练是整个网络的瓶颈），我们就会得到下面这种流水线的时间线：

数据预处理 pipeline 时间线

在执行几个 Batch 之后， 4 个阶段的执行节奏完全被最长的那个阶段所控制。这就是 OneFlow 使用背压机制（Back Pressure）解决流控问题。（比如 Preprocessing 阶段，该 Actor 是否执行不仅仅跟有没有数据相关，也要考虑自己有没有“空闲”的 Regst 块可写）

所以流水并行问题，在 OneFlow 中就是 Regst 数量的问题。在实际实现中， OneFlow 采用了一个更通用的算法实现了 Megatron 的流水并行：插入 Buffer Op。在逻辑计算图上，我们会给后向消费前向的边插入一个 Buffer Op， Buffer 的 Regst 数量和 Stage 相关。由于后向对前向的消费经过 Checkpointing 优化后，每个 Placement Group 下只会有非常少的几条消费边。整体的算法实现可以通过下面这个示意图来解释：

OneFlow 通过插入 Buffer Op 实现流水并行

假设整个网络分为 4 个 stage，共有 8 个 Transformer Layer，则我们需要在前 3 个（stage_num - 1）stage 的前后向计算图中插入 Buffer Op。最后一个 stage 由于每做完一个 micro-batch 的前向，立马做该 micro-batch 的反向，则不需要插入 Buffer。buffer 的 regst_num 跟 stage_num 相关。（图中是理想情况下，假设 stage 之间的传输开销可以忽略不计，则至少需要 stage_num - 1 的 buffer_size）由于我们对每一个 Transformer Layer 做了 Checkpointing，则每个 Layer 仅有一条前向到后向的数据边，则只需要插入一个 Buffer Op。

跟 Megatron 复杂的手写调度器和手写通信原语相比， OneFlow 系统层面只需要插入 Buffer 就可以实现流水并行。

2. OneFlow 如何实现数据+模型的混合并行？

我们以 Linear Layer 的数据 + 模型并行为例，来解释所有的数据并行和模型并行的组合，本质上都是被 SBP 所描述的 Signature 而已。任何并行方式的设备间通信操作，该在整个网络的哪里插入、该插入什么通信操作、每个设备该和谁通信，完全都是 SBP 自动推导得到的，而且还保证数学上的一致性。有了 OneFlow，算法工程师就告别了分布式并行中的通信原语了。不仅如此，OneFlow 的框架开发者绝大多数时候也不需要关心分布式里的通信原语，SBP 这层抽象使得算子/网络跟分布式通信解耦。

我们先以 1-D SBP 为例，之后再扩展到 2-D SBP。1-D SBP 下的数据并行，对于一个 Linear Layer 而言，主要是其中的 MatMul（矩阵乘法）计算。我们假设矩阵乘法计算在逻辑视角上是一个 (m, k) x (k, n) = (m, n) 的计算，m 表示一共有多少个样例， k 和 n 分别是 Linear Layer 中的隐藏层神经元数量以及输出神经元数量。

数据并行的逻辑计算图 -> 物理计算图的映射关系如下图所示：

数据并行下逻辑计算图转物理计算图

数据并行下，每个设备上都有全部的模型（Tensor b， Shape = (k, n)），假设共有两张卡，则 GPU 0 上有前一半的数据（Tensor a，Shape = (m/2, k)），GPU 1 上有后一半的数据，则我们说 Tensor a 的 SBP Parallel = Split(0)。同时我们可以看到矩阵乘的输出 Tensor out，也是按照第 0 维切分的。

模型并行对于 Linear Layer 而言，有两种，分别是切模型 Tensor 的第0维（行切分，对应 Megatron 里的 RowParallelLinear）和第1维（列切分，对应 Megatron 里的 ColumnParallelLinear）。

第一种行切分（RowParallelLinear）模型并行的逻辑计算图 -> 物理计算图的映射关系如下图所示：

模型并行（行切分）逻辑图转物理图

模型并行下，每个设备都只有一部分的模型，在这个例子中， GPU 0 上有前一半的模型， GPU 1上有后一半的模型，每个设备上的模型大小 Tensor b 的 Shape = (k/2, n)。在这种情况下，每个设备输出的 Tensor out 都是完整的数据大小， Shape = (m, n)，但每个位置上的元素的值，都是逻辑上的输出 out 对应位置的值的一部分，即 out 的 SBP Parallel = PartialSum 。

第二种列切分（ColumnParallelLinear）模型并行的逻辑计算图 -> 物理计算图的映射关系如下图所示：

模型并行（列切分）逻辑图转物理图

这个例子中，模型 Tensor b 是按照 Split(1) 切分的，输出 Tensor out 也是按照 Split(1) 切分的，每个设备都需要全部的数据。

在 GPT 网络中，实际上的模型并行是组合使用 RowParallelLinear 和 ColumnParallelLinear 实现的（ColumnParallelLinear 后面接了 RowParallelLinear）。

注：这里我没有列出整个 Transformer Layer 的 SBP 推导信息，只说了“交替使用 Row 和 Linear，插入 AllReduce ”。后续我会补充实际 GPT 的网络结构 + 模型并行 SBP 信息的图和说明。

因为 Column 的输出 Tensor SBP 是 Split(1)， Row 的输入数据 Tensor SBP 也是 Split(1)，所以当 Column 后接 Row 时，两者之间是不需要插入任何通信的。但由于 Row 的输出是 PartialSum，当后面消费该 Tensor （在网络中是 Add 操作）的 Op 需要全部的数据时（Broadcast），此处就需要插入 AllReduce 实现通信了。

这个在 OneFlow 中称之为 Boxing。当两个逻辑上的 Op 对于同一个逻辑上的 Tensor 看待的 SBP Parallel 不一致时， OneFlow 系统会自动插入通信节点以完成数据的切分/传输/拼接等操作，使得下游 Op 总能拿到按照自己期望 SBP 切分的 Tensor。

下面展示两个 Boxing 的示例。第一个是 PartialSum -> Broadcast ，为了得到完整的数据我们需要将 PartialSum 的 Tensor 对应位置加起来，这时候 OneFlow 会自动插入 AllReduce 操作（这里类比 Megatron 在模型脚本里手写 AllReduce 通信原语）。

Boxing：通过 AllReduce 实现 PartialSum 转 Broadcast

第二个是 Split(1) -> Broadcast，此时我们需要将按照第1维切分的 Tensor 拼接起来，这时候 OneFlow 会自动插入 AllGather 操作：

Boxing：通过 AllGather 实现 Split(1) 转 Broadcast

在 OneFlow 中，所有的分布式通信操作都是基于 SBP 的推导结果，按照需要插入。OneFlow 通过 Boxing 机制，就实现了任意的数据并行和模型并行。

2-D SBP 其实就是将两组 1-D SBP 按照设备拓扑的维度拼起来就可以得到。

对于数据并行的 MatMul （a x b = out）操作， SBP Signature是：{a : Split(0), b : Broadcast, out : Split(0)}，模型并行（行切分， RowParallelLinear）的 SBP Signature 是：{a : Split(1), b : Split(0), out : PartialSum}，那如果逻辑上的一个 MatMul Op 同时做数据并行和模型并行，其 2-D SBP Signature 就是：{a : [Split(0), Split(1)], b : [Broadcast, Split(0)], out : [Split(0), PartialSum] } （其实这个 out 的 [Split(0), PartialSum] 就是 GPT 中 RowParallelLinear 输出 Tensor 的 SBP），当下游消费这个 out Tensor 的期望 SBP 是 [Split(0), Broadcast]时， OneFlow 会自动在这两组 Op 之间插入一组 AllReduce 通信 Op，且这组 AllReduce 通信 Op 是按照设备拓扑的第 0 维分组进行的。（如设备拓扑是 (4, 8) 时，所有的 AllReduce 会分为 4 组，每组内的 8 个设备会互相 AllReduce，组间没有通信。这就实现了组内（机器内）模型并行，组间（机器间）数据并行）

其实 GPT 中用到的 2-D SBP 只是最简单情形的特例，分布式下的并行经过 2-D SBP 可以拓展出非常多复杂、灵活多边的组合出来。而针对复杂的组合，再想用 Megatron 那套就非常难做了，但是对于 OneFlow 而言，二者的难度是一样的，因为本质上是用 Boxing 完成一组 2-D SBP 的变换而已。

GPT 分布式训练性能对比：OneFlow vs Megatron

OneFlow 跟 Megatron 相比，除了用户接口（分布式易用性）和框架设计上更简洁、更易用，我们在已有的测试规模上性能也略微领先 Megatron（其实经过 NVIDIA 的深度优化， Megatron 在 GPU 上的分布式训练性能已经接近极致了，DeepSpeed 也比不上，而我们在 Megatron 的基础上易用性、效率都更进一步）。

注：由于我们自己拥有的集群规模限制，我们只测试了 4 机 32卡 16GB V100 规模内的性能结果。我们非常欢迎有大规模集群硬件的伙伴一起通过 OneFlow 合作研究、训练大规模预训练模型。未来我们会公布更大规模集群上 OneFlow 的优异表现。

以下的所有实验数据均在相同的硬件环境、相同的第三方依赖（CUDA、 cuDNN等）、使用相同的参数和网络结构下，对比了 OneFlow 和 Megatron 在 GPT 模型下的性能表现。所有的性能结果均保证公开且可复现。我们的 GPT 模型脚本在：Oneflow-Inc/OneFlow-Benchmark 仓库，公开的评测报告、复现方式稍后在：Oneflow-Inc/DLPerf 仓库中可以查看。

1.数据并行性能对比

注：每组参数的缩略版含义：

· DP 数据并行；MP 模型并行；2D 数据 & 模型的混合并行；PP 流水并行

· dxmxp_B_hxl 其中：

· d = 数据并行度（data-parallel-size）

· m = 模型并行度（tensor-model-parallel-size）

· p = 流水并行度（pipeline-model-parallel-size）

· B = 总的BatchSize（global-batch-size）

· h = 隐藏层大小（hidden-size）影响每层 Transformer Layer 的模型大小

· l = Transformer Layer 层数（num-layers）

数据并行性能对比

2.模型并行性能对比

注：由于单卡 GPU 显存限制，我们对于各组参数里的模型大小是不同的，所以整体不是像数据并行那样呈一个线性增加的关系。如第 4 组参数（MP_1x32x1_16_3072x32）的模型大小是第 2 组参数（MP_1x8x1_16_1536x16）的 8 倍以上。NVIDIA 论文中有模型规模跟各个参数的计算公式：图片

其中 l 表示 num-layers ，h 表示 hidden-size， V 表示词表大小（vocabulary size = 51200）， S 表示句子长度（a sequence length = 2048）， P 表示参数规模。