撰文 | 成诚

2020 年，最轰动的 AI 新闻莫过于 OpenAI 发布的 GPT-3 了。它的1750亿参数量及其在众多NLP任务上超过人类的出众表现让大家坚信：大模型才是未来。但与之带来的问题是，训练超大模型所需的算力、存储已不再是单机就能搞定的了（之前的 BERT 还是可以用 DGX-1/2 这样的超级服务器训练）。

NVIDIA 估算过，如果要训练GPT-3 ，即使单个机器的显存/内存能装得下，用 8 张 V100 的显卡（一台 DGX-1 的配置），训练时长预计要 36 年；即使用 512 张 V100 ，训练也需要将近 7 个月；如果你拥有 1024 张 80GB A100， 那么完整训练 GPT-3 的时长可以缩减到 1 个月。

这意味着训练大模型一定是一个分布式问题。对算力的需求还是一个相对容易解决的问题，因为拥有大集群的组织并不只 OpenAI 一家，而如何解决上千块 GPU 的分布式训练问题才是关键。

即使你是一位非常优秀的数据科学家，知晓并能解决 Transformer 相关的所有算法问题，但如果你不是分布式专家，不知道如何解决分布式训练时上百台服务器之间的通信、拓扑、模型并行、流水并行等问题，你甚至都无法启动这次训练。这也解释了为什么时隔一年，只有 NVIDIA 、微软等大企业可以复现 GPT-3 。

目前开源的 GPT 模型库主要是 NVIDIA 的 Megatron-LM 和微软的 DeepSpeed。其中，微软的 DeepSpeed 的模型并行等内核取自 Megatron，且 DeepSpeed 主打的是，在数据并行下如何以更少的机器去跑更大的模型 （ ZeRO 、 ZeRO-Offload 等都是用梯度切片、计算、内存/硬盘换入换出来省显存），所以我们本文主要介绍和对比 Megatron 。

这里先简单比较一下 NVIDIA 的 Megatron 和 微软的 DeepSpeed：DeepSpeed 本质上是一种“节省显存”的数据并行，即：数据并行的优化版。DeepSpeed 假设了单层参数量可以在单张显卡上放得下，如果不满足这个假设，那么仍然需要使用模型并行，而且 DeepSpeed 的模型并行是通过调用 Megatron 来实现的。根据 NVIDIA 最新的那篇论文（链接：https://arxiv.org/abs/2104.04473，也是下面本文重点要介绍的），Megatron 在大规模训练的效率是超过 DeepSpeed 不少的。而 DeepSpeed 的论文一直强调：可以用更少机器训练更大的模型，但没有突出过在效率上的优势。DeepSpeed 后来又出了一篇论文：ZeRO-Infinity（链接：https://arxiv.org/abs/2104.07857），当单层参数量在单张显卡上放不下的时候，它通过对这一层算子切片，一片一片来执行，使得单卡也能跑起来一个巨大的层，可以理解成一种 “时间”轴上展开的模型并行。

Megatron 和 DeepSpeed 都是基于 PyTorch ，分别由 NVIDIA 和微软经过深度定制开发，专门为支持 PyTorch 分布式训练 GPT 而设计的。我们会简单介绍一下 NVIDIA 如何使用 PyTorch 搞分布式训练 GPT ，然后重点介绍 OneFlow 如何用一套通用设计非常简单清晰地解决了这个难题，同时我们还在已有的测试规模上性能超过 NVIDIA。

相信读完此文，你就会发现 PyTorch 、 Megatron ( NVIDIA ) 、DeepSpeed ( Microsoft ) 都走了一个非常长的弯路，这条弯路从大方向上就走错了，不仅是弯路，你还会发现 Megatron 的代码只能被 NVIDIA 的分布式训练专家所复用，它对于 PyTorch 的算法工程师而言门槛极高，是非常难用的，以至于任何想要用 PyTorch 复现一个分布式大模型的算法工程师，都得先等 NVIDIA 开发完才能再使用 Megatron 提供的模型。同时，我们也通过这样一个例子来证明：为什么一个分布式深度学习框架要像 OneFlow 这样设计。

本文内容较多，先列出主要目录：

1.分布式训练 GPT 的必要并行技术

流水并行梯度累加后向重计算1F1B 策略

2.Megatron：PyTorch 分布式训练的极限，痛点在哪儿？

流水并行 PyTorch 需要手写专家级复杂调度器模型并行 PyTorch 需要手工排线，在 kernel 里手写特定的、经过复杂推导的通信原语

3.OneFlow 用一致性视角（Consistent View）轻松填平分布式训练难的鸿沟

流水并行 ，只需要配置 Placement 就够了数据 模型的混合并行，只需要配置 Variable 的 SBP 就够了OneFlow：让每一位算法工程师都有能力训练 GPT

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

OneFlow 系统层面如何实现流水并行OneFlow 系统层面如何实现混合（数据 & 模型）并行

5.GPT 训练性能对比: OneFlow vs Megatron

纯数据并行性能对比纯模型并行性能对比混合并行性能对比流水并行 混合并行 性能对比

6.小结

一、分布式训练 GPT 的必要并行技术

最近，NVIDIA 放出了一篇重量级的论文：Efficient Large-Scale Language Model Training on GPU Clusters ，用了 3072 张 80GB A100 训练 GPT（ NVIDIA 也确实够壕，这个集群的成本就不止 5 亿了），最大规模的模型参数量达到了 1T（是 GPT-3 原版的 5 倍）。

NVIDIA 训练 GPT-3 最大到 1T 参数规模

论文里 NVIDIA 介绍了分布式训练超大规模模型的三种必须的并行技术：

数据并行（Data Parallelism）模型并行（Tensor Model Parallelism）流水并行（Pipeline Model Parallelism）

其中数据并行是大家都熟知的最常见的并行方式，而模型并行（NVIDIA 论文里叫做 "Tensor 级别的模型并行" ）是对某一层（如 Linear/Dense Layer 里的 Variable ）的模型 Tensor 切分，从而将大的模型 Tensor 分成多个相对较小的 Tensor 进行并行计算；流水并行（NVIDIA 论文里叫做流水线级别的模型并行），是将整个网络分段（stage），不同段在不同的设备上，前后阶段流水分批工作，通过一种“接力”的方式并行。

对于最大的 1T 规模的模型，NVIDIA 一共使用了 384 台 DGX-A100 机器（每台装有 8 张 80GB A100 GPU），机器内部各 GPU 间使用超高速 NVLink 和 NVSwitch 互联，每台机器装有 8个 200Gbps 的 InfiniBand (IB) 网卡，可以说是硬件集群顶配中的顶配了。

那么，这些机器是如何协同工作的？GPT 网络是由很多层 Transformer Layer 组成的，每一层内部是一个由多层 MLP 和 attention 机制组成的子图，对于参数规模 1T 的 GPT 而言就有 128 层的 Transformer Layer，这个超大超深的网络被分割成了 64 个 stage ，每个 stage 跑在 6 台 DGX-A100 上，其中 6 台机器之间进行数据并行，每台机器内部的 8 张卡之间做模型并行，整个集群的 3072 张 A100 按照机器拓扑被划分成了 [6 x 8 x 64] 的矩阵，同时使用数据并行 & 模型并行 & 流水并行 进行训练。

BSP 各个阶段串行执行

BSP 且没有 Gradient Accumulation 下的流水并行。假设整个网络被等分切成 4 个 stage，每个 stage 使用一个 device ，则在BSP下，各个设备串行执行，中间有大段的气泡。一般后向计算时间是前向计算的两倍，如果算上 Checkpointing 的重计算部分，是前向计算的三倍。我们可以从上图中看到，这种情况下有 70% 的时间设备是空闲的。

而 SSP 就是异步模型更新，允许前向计算时可以不使用最新的模型，而使用落后几个版本之内的模型。SSP 在 GPT-3 的训练中并没有被 NVIDIA 采用，其主要原因有以下几点：

SSP 的模型收敛性并没有被严格论证， 且有论文 GeePS 指出 SSP 的收敛效果不如 BSP ；SSP 会在 GPU 上同时出现几个不同版本的模型，而 GPT-3 又是一个非常大的模型网络，多份模型所带来的显存开销不可接受；BSP 通过 Gradient Accumulation Checkpointing 就可以很好的解决 GPT-3 中的流水并行问题。另外， NVIDIA 的论文的分析前提就是 BSP 情况下， 根据严格的参数优化器更新方式， 流水并行气泡的占比是 Bubble time fraction = (p - 1) / m，其中 p 是 stage 数， m 是梯度累加的 micro-batch 数。如果采用 SSP，则 NVIDIA 整篇文章的理论基础就没有了。

"Pipeline parallelism comes in a few flavors: the mode discussed in this paper uses flushes to ensure exact strict optimizer semantics."

2.梯度累加

Gradient Accumulation 就是把一个大 Batch 拆分成多个 micro-batch ， 每个 micro-batch 前后向计算后的梯度累加，在最后一个micro-batch累加结束后，统一更新模型。

micro-batch 跟数据并行有高度的相似性：数据并行是空间上的， 数据被拆分成多个 tensor，同时喂给多个设备并行计算，然后将梯度累加在一起更新；而 micro-batch 是时间上的数据并行， 数据被拆分成多个 tensor， 按照时序依次进入同一个设备串行计算，然后将梯度累加在一起更新。当总的 batch size 一致，且数据并行的并行度和 micro-batch 的累加次数相等时，数据并行和 Gradient Accumulation 在数学上完全等价。Gradient Accumulation 通过多个 micro-batch的梯度累加使得下一个 micro-batch 的前向计算不需要依赖上一个 micro-batch 的反向计算，因此可以畅通无阻的进行下去（当然在一个大 batch 的最后一次 micro-batch 还是会触发这个依赖）。

Gradient Accumulation 解决了很多问题：

在单卡下，Gradient Accumulation 可以将一个大的 batch size 拆分成等价的多个小 micro-batch ，从而达到节省显存的目的。在数据并行下，Gradient Accumulation 解决了反向梯度同步开销占比过大的问题（随着机器数和设备数的增加，梯度的 AllReduce 同步开销也加大），因为梯度同步变成了一个稀疏操作，因此可以提升数据并行的加速比。在流水并行下， Gradient Accumulation 使得不同 stage 之间可以并行执行不同的 micro-batch， 从而让各个阶段的计算不阻塞，达到流水的目的。

单纯通过 micro-batch，我们就实现了 GPipe （2018）论文中的流水并行，在 stage 数量为 4， micro-batch 数量为 8 （每个 batch 在计算 8 个 micro-batch 且 累加 8 次梯度后更新一次模型）下的时间线如下图所示：

使用梯度累加后的 Pipeline 时间线

在 GPipe 的流水并行示例中，每个“时间点” 可以在多个阶段（stage）上同时做不同的micro-batch，图中每个方块中的标号表示了第几个 micro-batch；同一个 micro-batch 还是串行的经过所有的 stage，在这种情况下，每个设备的空闲时间只有 25% 左右。

但这种流水并行存在一个问题：显存占用太大。如果每个 micro-batch 前向计算的中间结果（activation）被后向计算所消费，则需要在显存中缓存 8 份（梯度累加的次数）完整的前向 activation。这时就不得不用另一项重要的技术：重计算（Checkpointing）。

3.后向重计算

Checkpointing 是陈天奇在2016年发表的论文 Training Deep Nets with Sublinear Memory Cost 中提到的，也称之为亚线性内存优化。亚线性内存优化有两种思路，Checkpointing 和 CPU offload：

Checkpointing 的核心思想 是在前向网络中标记少量的 Tensor （被 Checkpointing 的 Tensor ），前向计算就只会保留这些被标记的 Tensor， 其余的前向的 activation，会通过在反向传播中根据 Checkpointing 的 Tensor 临时重新计算一遍前向得到。这样就使得大量的 activation 不需要一直保存到后向计算，有效减少了大量 Tensor 的生命周期，使得内存复用效率大幅提升。CPU offload 的思路类比于计算机操作系统中的“虚拟内存”技术（将不常用的内存临时换入换出到磁盘上，从而增加内存总量），在深度学习中，GPU 显存（Device Memory）的特点是昂贵、高速且容量小，而 CPU 主存（Host Memory）的特点是便宜、相对低速和大容量；那么将前向计算中的一些暂时用不到的 activation 临时换出到 CPU 主存上，等到反向计算需要时再换入到 GPU 显存里，通过这种方式也可以节省显存。

两种亚线性内存优化通过不同的方式达到了显存优化：Checkpointing 是通过额外的计算开销换显存， CPU offload 通过额外的传输开销换显存。

Checkpointing 优化

上图展示了两层 Transformer Layer 在做 Checkpointing 之前和之后的计算图对比， 其中重要的区别是前后向之间的连边从很多条变成了两条。不同框架实现Checkpointing的思路不同，Megatron 是自己重载了 torch.nn.Module ，实现了自己的 checkpointed_forward，相当于定制化了 Transformer Layer 的前后向执行逻辑；OneFlow 的 Checkpointing 就是上图中的设计， 我们在整个计算图中插入了重计算的子图，并使得后向对前向的消费转移到了对重计算子图的消费。

重计算并不是单独为流水并行设计的，并且之前大多使用在单卡或者数据并行场景下。但这个优化在流水并行下就非常关键，因为它使得前向不需要缓存所有的 activation，而只需要缓存非常少个数的（比如一层 Transformer Layer 只会缓存一个 ）、被 checkpoint 的特定 Tensor ，从而大大节省了流水并行下的显存开销。

4. 1F1B 策略

除了重计算，上述 GPipe 的流水并行策略还有另外一个内存问题，就是需要缓存几份 activation，是等于一个 batch 里有多少个 micro-batch 的（梯度累加的次数）。通常，这个累加次数都比较大（为了尽可能流水，累加次数一般大于两倍的 stage 数），那么即使缓存少数 Tensor， 这种策略仍需要较多显存。

因此，在另一篇流水并行的论文PipeDream （2018） 里就提出了改进方法，称之为 1F1B （One Forward pass followed by One Backward pass）的策略。这种改进策略可以解决缓存 activation 的份数问题，使得 activation 的缓存数量只跟 stage 数相关，从而进一步节省显存，训练更大的模型。

1F1B 策略的出发点也比较直观：由于前向计算的 activation 需要等到对应的后向计算完成后才能释放（无论有没有使用 Checkpointing 技术），因此在流水并行下，如果想尽可能节省缓存 activation 的份数，就要尽量缩短每份 activation 保存的时间，也就是让每份 activation 都尽可能早的释放，所以要让每个 micro-batch 的数据尽可能早的完成后向计算，因此需要把后向计算的优先级提高，让 micro-batch 标号小的后向比 micro-batch 标号大的前向先做。因此，如果我们让最后一个 stage 在做完一次 micro-batch 的前向后，立马就做本 micro-batch 的后向，那么我们就能让其他的 stage 尽可能早的开始后向计算，这就是 1F1B 策略。其时间线如下图所示：

NVIDIA 模型并行通信推导

流水并行的优势是带宽需求比其它并行方式低，仅需要在 stage 之间传输数据， 同时还不会阻塞整个网络的计算，因此在机器间做流水并行比较合适；但流水并行必须通过把一个 Batch 分割成若干 micro-batches 才能发挥优势， 同时它还需要额外的显存来缓存 activation，在 batch 间还会留下气泡。

NVIDIA 在论文中实验了相同的总模型并行度（ model-parallel-size = tensor-model-parallel-size * pipeline-model-parallel-size）下， 分配不同的模型并行和流水并行的 size，得出当 tensor-model-parallel-size = 8 时， 总的效率最高，这与每台机器内的卡数相同 。

一致性视角（Consistent View）抽象

理想情况下，抽象出来的超级设备（逻辑视角）的算力是所有物理视角下的设备算力之和（如果算力完全用满，就是线性加速比）；逻辑视角下的显存资源也是所有物理设备的显存资源之和。

1.流水并行 ，只需要配置 Placement 就够了

Placement 描述了逻辑上的 Op 和 物理上的 Op 的映射关系， 一个 Op 分布在哪些机器、哪些设备上，只需要在逻辑图里指定其 Placement 即可。在流水并行里，用户不需要像 PyTorch 一样关心什么时机 send/recv 、 到底是执行 forward 还是 forward backward，用户只需要配置 Placement （现阶段我们还需要用户帮忙指定 stage id）就完成了流水并行，以下就是 OneFlow 来做 GPT 的流水并行的代码（其实是包含了所有的并行方式）：

class Transformer(object): def __call__(self, hidden_states): for i in range(self.num_layers): with distribute.layer_placement_scope(i): h = self.layers[i](h)

其中 layer_placement_scope 就是在 scope 中配置了 placement 和 stage id：

def layer_placement_scope(layer_idx, device="gpu"): dist_util = get_dist_util() with flow.scope.placement( device, dist_util.get_layer_placement(layer_idx), dist_util.parallel_hierarchy, ): if dist_util.is_pipeline_parallel(): with flow.experimental.scope.config( pipeline_stage_id_hint=dist_util.get_layer_stage(layer_idx) ):

其余的所有流水并行的细节：如各个 stage 之间怎么通信， 做前向还是后向， 怎么流水起来，怎么保证正确性，这些全都不用用户操心。

在下一章节我们还会介绍 OneFlow 内部是怎么实现流水并行的，相较于 Megatron 的复杂实现， OneFlow 系统层面的工作对框架开发者而言也更加友好。

上图展示了一个可能的 Placement 例子，用于 GPU0 和 GPU1 之间的流水并行。图中负责在 CPU 和 GPU、GPU 与 GPU 之间进行数据搬运的Op（CopyH2D、CopyD2D）是 OneFlow 系统自动添加的。

2.数据 模型的混合并行，只需要配置 Variable 的 SBP 就够了

SBP 描述了 逻辑上的 Tensor 和 物理上的 Tensor 的映射关系。SBP 是三种基础映射的首字母组合：Split、Broadcast、Partial，（其中 Partial 是一个 reduce 操作，包括 PartialSum、PartialMin、PartialMax等），全称叫做 SbpParallel。其中：

Split：表示物理上的多个 Tensor 是由逻辑上的 Tensor 进行切分后得到的。Split 会包含一个参数 Axis，表示被切分的维度。如果把所有物理上的 Tensor 按照 Split 的维度进行拼接，就能还原出逻辑上的 Tensor。Broadcast：表示物理上的多个 Tensor 是逻辑上 Tensor 的复制，两者数据完全相同。PartialSum：表示物理上的多个 Tensor 跟逻辑上的 Tenso r的形状相同，但每个对应位置上元素的值 是逻辑Tensor对应位置元素的值的一部分。如果把所有物理上的 Tensor 按照对应位置相加（element-wise），即可还出原逻辑上的 Tensor。

下图展示了SBP的几种简单情形。

2D-SBP 下 逻辑 Tensor 和 物理 Tensor 的对应关系

在 2D SBP 下， 设备拓扑被称之为 ParallelConf::hierarchy ，表示如何将一维的 world_size = 8 映射成为二维的层次结构：(2, 4)。而 2-D SBP 被称之为ParallelDistribution， 是由两个 SbpParallel 组成的 list: [Broadcast, Split(0)]。只需要简单配置一下 SBP， 用户就可以轻松设定网络是数据并行、模型并行还是混合并行。

3. OneFlow：让每一位算法工程师都有能力训练 GPT

综合比较 PyTorch 和 OneFlow 在分布式并行上的设计以及开放给用户的接口，有如下总结：

PyTorch 分布式的困境：

PyTorch 只有物理视角，没有逻辑视角，所以 PyTorch 的用户想要做 分布式并行，任何时候都需要自己推导深度学习模型中哪处需要跟其他的物理设备进行通信和数据同步操作，既要推导通信所在的位置，又要推导通信的操作类型，还要推导跟其他哪些设备通信。这个在简单的数据并行下可以使用 DDP 或者 Horovod 来实现，但是在复杂的模型并行、混合并行下，做并行的门槛会非常高。PyTorch 没有将模型网络的算法逻辑和分布式并行训练的通信逻辑解耦出来，导致需要在 算子的 kernel 实现中、 搭网络的脚本里到处插入通信原语。这些手写通信原语的操作不仅繁琐、易错、而且没法复用，是根据特定模型、特定脚本位置、特定算子特判得到的。PyTorch 在非对称的并行方式里（如流水并行），各个设备的调度逻辑需要用户自己手写，用户需要自己精细的控制每个设备上的启动以及执行逻辑，且执行逻辑把前后向执行和send/recv通信操作糅合在了一起，即使在最规整的 Transformer Layer 的流水并行下也很复杂，想要扩展到其他模型上的工作量也很大。PyTorch 没有机制保证分布式并行训练中的正确性 和 数学一致性。即使用户写错了通信操作、插错了位置、 跟错误的设备进行了通信，PyTorch也检查不出来。

这些都使得用户使用 PyTorch 开发复杂分布式训练的脚本极为困难，以至于只有 NVIDIA 、 微软 等大企业的分布式训练专家才能开发出 GPT 的 PyTorch 版本。

OneFlow 分布式的易用性：

OneFlow 有一致性视角 Consistent View， 将分布式训练下的多机通信和 算法逻辑解耦，使得用户可以不用关心分布式训练的细节，降低了分布式训练的使用门槛。OneFlow 通过 Placement SBP 机制解决了分布式训练中任意并行场景的需求，用户只需要配置 op 的 Placement 就可以完成流水并行，只需要配置 Tensor 的 SBP 就可以实现数据并行、模型并行和混合并行。而且任何并行方式都是 Placement SBP 的一种特例， OneFlow 从系统层面不需要做任何的特判，SBP 才是各种分布式并行的本质。OneFlow 的通信逻辑可以复用，不需要为任何特定网络和特定算子实现相应的通信逻辑。通信逻辑由 OneFlow 的 Boxing 机制完成，与具体的算子和模型无关。OneFlow 的 SBP 还保证了数学上的一致性， 相同的逻辑上的模型脚本，使用任意的并行方式（数据并行、模型并行、流水并行）、使用任意的集群拓扑，OneFlow 都从数学上保证了模型分布式训练的正确性。

因此，我们才说 OneFlow 用一套简洁设计解决了分布式并行的各种难题。对于用户而言， OneFlow 让算法工程师不需要成为一位分布式训练的专家也有能力做复杂的分布式训练， 只要有硬件资源，任何一位算法工程师都可以训练 GPT， 任何一位算法工程师都可以开发一个新的大规模分布式训练的模型。

四、为什么分布式深度学习框架要像 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时，如果训练时间比较长（训练是整个网络的瓶颈），我们就会得到下面这种流水线的时间线：

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）模型并行的 逻辑计算图 -> 物理计算图 的映射关系如下图所示：

数据 模型 流水并行性能对比（注：第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