数据预处理流程

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

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

这个例子中，模型 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 仓库中可以查看。

注：每组参数的缩略版含义：· 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）

数据 模型并行性能对比（注：其中前 4 组的模型规模一致；后 2 组的模型规模一致。）

数据 模型 流水并行性能对比（注：第1组参数的模型比后3组都要小，因为机器内的数据并行限制了参数规模。）

OneFlow 在分布式训练领域拥有独特的设计和视角，解决了分布式训练中的各种并行难题，因此在大规模预训练模型场景下用 OneFlow 做分布式训练更易用也更高效。但相比 PyTorch 在单机单卡视角下的极致易用性，OneFlow 的前端用户接口还有明显的差距。

OneFlow 研发团队正在全力提升框架的单卡使用体验， 并从即将在 5 月发布的下个大版本 OneFlow v0.4.0 起， OneFlow 开始提供兼容 PyTorch 的全新接口以及动态图等特性。我们预计在 v0.5.0 全面兼容 PyTorch， 届时用户将 PyTorch 的模型训练脚本迁移成 OneFlow 的训练脚本几乎是一键的，除此之外， 我们还会提供 Consistent 视角的分布式 Eager，用户可以既享受动态图的易用性，又可以非常方便的进行各种分布式并行训练，欢迎前来体验。

欢迎持续关注OneFlow：https://github.com/Oneflow-Inc/oneflow

注：题图源自pixabay

花粉社群VIP加油站