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

撰文 | 成诚

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] 的矩阵，同时使用数据并行 & 模型并行 & 流水并行 进行训练。

3072 张 A100 集群拓扑

1.流水并行

从上述的机器拓扑中可以发现，流水并行是 3072 块 A100 能训练 GPT 的关键。因为无论是数据并行还是模型并行，都会在相应的机器之间进行全连接的通信，当机器数量增大时，通信开销和时延会大到难以忍受。而流水并行既解决了超大模型无法在单设备上装下的难题，又很好解决了机器之间的通信开销的问题，每个阶段（stage） 和下一个阶段之间仅有相邻的某一个 Tensor 数据需要传输，每台机器的数据传输量跟总的网络大小、机器总数、并行规模无关。但流水并行为了多个阶段之间可以流水起来，还依赖两个重要的特性： 梯度累加（Gradient Accumulation） 和 亚线性内存优化（ Sublinear Memory Cost 2016, 陈天奇）。

近期，百度和华为相继发了自己的千亿级中文预训练模型的宣传文。其中，百度提出了 "4D混合并行"，本质上是 Megatron 里的数据并行 + 模型并行 + 流水并行 + DeepSpeed 里的 ZeRO 优化 ；华为文章中的 “5D混合并行”，是将重计算（Checkpointing, 亚线性内存优化的一种）作为了第5维 （其实百度也做了重计算，只是没有将其列为多维并行中的一维）。

在介绍这两个特性之前，我们先简单解释一下深度学习训练和模型更新的两种约束：BSP (Bulk Synchronous Parallel) 和 SSP (Stale Synchronous Parallel ) ，其中 BSP 是最常见的模型更新规则：每个 batch 的前向计算都需要使用最新的模型，那么就要求上一个 batch 的后向计算结束且在模型更新后，下一个 batch 的前向才能开始。如果使用 BSP 去做流水并行，我们就会发现每个阶段的前向和后向是完全串行的，其中一个设备在工作时，其他所有设备都在等待，那么分布式的优势就完全没有被发挥出来：

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 策略。其时间线如下图所示：

1F1B 策略下的 Pipeline 时间线

从上图 1F1B 和之前 GPipe 的流水线对比可知， GPipe 需要缓存 8 份的 activation 供后向使用，而 1F1B 策略只需要缓存 4 份。二者虽然空闲时间的占比是一样的，但节省显存就可以跑更多的 Layer 层数 和 更大的 micro-batch size，从而提升性能。

以上几个关键技术（GPipe、梯度累加、重计算和 1F1B）的介绍就是分布式训练 GPT 的流水并行的核心技术（数据&模型并行我们放在下一章节详细介绍）。无论是 NVIDIA 的Megatron（PyTorch），还是 OneFlow、PaddlePaddle、MindSpore ，都是通过不同的设计实现了上述相同的功能，而且 Megatron 在 NVIDIA 的深度优化下， 在 GPU 上的性能表现已经非常优异了。那么 OneFlow 再搞一套 GPT 的意义何在？别急，看了下一章节，你就知道 PyTorch 做到上述这些技术的痛点在哪儿了。

二、Megatron ：PyTorch 分布式训练的极限、痛点在哪儿？

NVIDIA 基于 PyTorch 开发了 Megatron，本质上是一个专用于 GPT 的模型库，所有的代码都是 Python 脚本，NVIDIA 为 GPT 专门定制了分布式训练所需的算子、 流水并行调度器、模型并行所需的通信原语等功能。可以说，NVIDIA 在使用 PyTorch 做分布式训练上已经做到极致了。

在本章节，我们会简单介绍一下 Megatron 是如何使用 PyTorch 的，当你也了解 Megatron 的设计以后，你就可以回答这个问题： PyTorch 做分布式训练，真的好用吗？

1.流水并行，PyTorch 需要人工排线和精细控制流水

PyTorch 是单卡视角，一个设备上的 Tensor、模型脚本跟另一个设备上的 Tensor、模型脚本并无直接关系，对于每个设备上的模型脚本都完全对称的（Mirror）最简单的数据并行来说，PyTorch 这样的设计没有什么明显的缺陷。每个设备上的脚本运行到相同 batch 的模型更新部分（Optimizer），统一做一次模型同步（AllReduce 操作）就完成了数据并行，这就是 PyTorch 的 DDP（DistributedDataParallel）模块。

而流水并行，模型网络分布在各个设备上是非对称的，各个设备“接力”执行网络的一部分，这种并行方式用 PyTorch 要如何实现呢？

流水并行 2 卡接力执行网络

上图展示了流水并行下，前两个 stage 分布在 GPU 0 和 GPU 1 上时，网络的拓扑关系。GPU 0 和 GPU 1 是接力执行的， GPU 0 上的 T2 Layer 的输出 Tensor 需要发给 GPU 1 上的 T3 Layer 作为输入。

首先，你需要根据 stage 阶段的不同，分别在各个设备上定义只属于自己那部分的模型网络，而由于第一个 stage 和最后一个 stage 在执行时序上的特殊性，这里 Megatron 还需要进行特判 megatron/training.py 。

def train_step(...):
if mpu.is_pipeline_first_stage():
unwrapped_model = model[0]
elif mpu.is_pipeline_last_stage():
unwrapped_model = model[-1]

在每个设备根据自己的那部分网络启动以后， Megatron 需要给每个设备上的每一次执行前后都调用 NCCL 的通信操作，前一个 stage 的输出需要通过 NCCL p2p的 ncclSend 操作发给 下一个 stage， 下一个 stage 必须同时调用 ncclRecv 进行接收。当这两个操作成对出现时，这次传输才会成功。（megatron/schedules.py）

def forward_backward_pipelining_without_interleaving(...):
for i in range(num_microbatches_remaining):
output_tensor = forward_step(...)
if forward_only:
p2p_communication.send_forward(output_tensor, timers)
else:
output_tensor_grad = p2p_communication.send_forward_recv_backward(output_tensor, timers)

# Add input_tensor and output_tensor to end of list, then pop from the
# start of the list for backward pass.
input_tensors.append(input_tensor)
output_tensors.append(output_tensor)

if forward_only:
if not last_iteration:
input_tensor = p2p_communication.recv_forward(timers)
else:
input_tensor, output_tensor = input_tensors.pop(0), output_tensors.pop(0)
input_tensor_grad = backward_step(...)
if last_iteration:
input_tensor = None
p2p_communication.send_backward(input_tensor_grad, timers)
else:
input_tensor = p2p_communication.send_backward_recv_forward(input_tensor_grad, timers)

因此对于 PyTorch 用户而言，用户自己需要关心每个 stage 在什么时机需要 recv，什么时机要 send， 发给谁；同时根据 Pipeline 的执行时序，需要特判在前多少个 step，都是需要只做前向（因为后向还没来）， 但又有一些 step，我需要既做前向又做后向，因此你可以看到在 megatron/p2p_communication.py 里，你会发现 Megatron 向用户提供了这些操作：

def recv_forward(...):
"""Receive tensor from previous rank in pipeline (forward receive)."""

def recv_backward(...):
"""Receive tensor from next rank in pipeline (backward receive)."""

def send_forward(...):
"""Send tensor to next rank in pipeline (forward send)."""

def send_backward(...):
"""Send tensor to previous rank in pipeline (backward send)."""

def send_forward_recv_backward(...):
"""Batched send and recv with next rank in pipeline."""

def send_backward_recv_forward(...):
"""Batched send and recv with previous rank in pipeline."""

def send_forward_recv_forward(...):
"""Batched recv from previous rank and send to next rank in pipeline."""

def send_backward_recv_backward(...):
"""Batched recv from next rank and send to previous rank in pipeline."""

def send_forward_backward_recv_forward_backward(...):
"""Batched send and recv with previous and next ranks in pipeline."""

通过这些接口，你就会发现，算法工程师如果想用 PyTorch 做流水并行，他需要精细的控制所有的流水细节， 包括每个 stage 的每个时刻是只做前向，还是前向后向一起做， 同时还需要管理不同 stage 之间收/发数据的节奏，这个要求对于用户而言就太高了。

更让人头痛的是，PyTorch 并没有机制保证这些流水并行中的各个设备之间数据交互的正确性 ，所以用户不仅可能写的不高效， 还可能写错，即使写错了，PyTorch 也无从检查。 这些都给用户带来了极大的使用门槛。因此，也只有 NVIDIA 、 微软等大企业的分布式训练专家可以搞得定 PyTorch 做流水并行。

2.模型并行，PyTorch 需要用户在 kernel 中手写通信原语操作，需要用户推导所有的通信位置

GPT 的大规模训练需要同时用到数据并行、模型并行和流水并行， 对于一个逻辑上的 Transformer Layer，需要同时对一个层做数据并行和模型并行，这个在 Megatron 和 DeepSpeed 的语义里称之为 data-parallel-size 和 tensor-model-parallel-size 。

为什么要既做数据并行，又做模型并行？其实是为了节省显存，并充分利用 GPU 之间的高速互联（NVLink 和 NVSwitch）带宽与机器之间的 IB 网络带宽的差别，NVIDIA 设计了一种在机器间做数据并行， 在机器内做模型并行的混合并行。

在什么样的网络结构、参数规模、网络拓扑下该用数据并行、模型并行还是流水并行，是一个非常复杂的问题。不同的并行方式导致的设备之间、机器之间的通信量是不同的；同时又需要考虑设备显存的约束、 GPU 通信带宽和网络通信带宽的占比、 总的 Batch Size 大小对模型收敛速度的影响等等。目前还没有一个严格的理论来指导具体模型在具体网络拓扑下究竟该用哪种并行配置最优。对于并行策略的研究，我们会在未来专门出一篇文章来探讨这个话题。

对于大部分情况而言，数据并行的效率一般是最高的，但在 GPT-3 这样的网络参数规模下，单个 GPU 根本装不下这么大的模型，所以必须要用到模型并行和 流水并行来降低每个 GPU 上的显存需求。又基于 NVLink 和 IB 网络通信带宽的差别，NVIDIA 设计了一种折中的的方案，对整个集群拓扑做分组，分为机器间和机器内，机器间的网络传输速度较慢，往往是分布式并行的瓶颈，所以适合做流水并行和数据并行；机器内的 NVLink 延迟低、带宽高，正好符合模型并行的要求，由于 GPT-3 必须使用模型并行，因此被放在了机器内做。

于 GPT-3 必须使用模型并行，因此被放在了机器内做。

数据并行是在反向的梯度更新时需要插入 AllReduce 操作，而模型更新在 Gradient Accumulation 里是一个低频操作，多个 micro-batch 只会做一次，所以数据并行在机器间做是比较合适的。
模型并行（Tensor Model Parallelism）， NVIDIA 推导了 Transformer Layer 里的 MLP 和 Self-Attention 操作，模型并行下需要在特定位置插入 AllReduce 来实现前向、后向的数据同步工作。由于模型并行需要在每个 micro-batch 的前向、后向都需要做数据同步， 属于高频操作， 所以 模型并行 适合在机器内做。

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 时， 总的效率最高，这与每台机器内的卡数相同 。

模型并行度和流水并行度对性能的影响

用 PyTorch 做模型并行的痛点是什么？如果你去了解一下 Megatron 搭 GPT 的模型脚本就立马清楚了，我们知道模型并行需要在前后向插入一些数据同步的操作，但是在哪里插入？NVIDIA 给出了最主要的同步操作推导结果：在 RowParallelLinear 里需要将这个同步写在 kernel 的 forward 函数里：

class RowParallelLinear(torch.nn.Module):
def forward(self, input_):
output_parallel = F.linear(input_parallel, self.weight)
# All-reduce across all the partitions.
output_ = reduce_from_tensor_model_parallel_region(output_parallel)

这是最关键的一处数据同步操作。但即使在 GPT 这样全部由 Transformer Layer 组成的非常规整的网络里，模型并行需要插入的同步操作就包括但不限于：

• AllReduce ：train_step 、 calc_params_l2_norm、CrossEntropy

• Scatter ：RowParallelLinear

• AllGather ：ColumnParallelLinear

• 等等...

那么问题来了，算法工程师怎么知道这么长的模型脚本里，到底：

• 哪处需要插入通信操作？（现在 GPT 的脚本里 NVIDIA 给推导了需要插入通信的位置，如果用户想改网络结构，想加/换一个Op，推导是不是都得重来？）

• 该插入什么通信操作？（除了 AllReduce，集合通信还有 ReduceScatter、AllGather、Reduce、Broadcast、All2All 等操作，除了集合通信，还有 Scatter 、Gather 等非对称的切分、拼接操作，切分/拼接还要考虑对 Tensor 的哪个维度操作...）

• 通信操作要跟谁通信？（数据并行和模型并行同时做时， 整个 GPU 集群会被分组，每一组组内做 AllReduce 同步数据， 组间在模型更新时 才同步模型梯度，这意味着每个 rank 的 GPU 想要通信时，是需要跟其他特定对应的 rank 做通信的，这更加增加了实现难度）

更要命的是，如果插入了通信操作，怎么保证正确性？PyTorch 没法保证。PyTorch 将所有的操作都交给了用户， 即使用户插入了一个错误的通信原语（比如将该插入 AllGather 操作的位置插入了 AllReduce），PyTorch 也没法检查出来。

所以这就是为什么只有 NVIDIA 可以用得了 PyTorch 做 Megatron，普通用户只能直接用 megatron/pretrain_gpt.py，想基于 Megatron 做其他模型/网络的迁移、二次开发和研究，是非常困难的。

其实，NVIDIA、 微软、 PyTorch 都被绕进一个大坑里去了：在没有一致性视角（ Consistent View ）的情况下做复杂的分布式并行是非常困难的，往往只能做一些具体网络具体场景具体算子的特判和分析，通过简单的通信原语来实现分布式。而 OneFlow 通过一致性视角下的 Placement + SBP 就非常简单的实现了通用的复杂并行支持。

三、OneFlow 用一致性视角轻松填平分布式训练难的鸿沟

对于上面 PyTorch 遇到的分布式难题， OneFlow 如何轻松解决？

OneFlow 的两大独特设计是：
1. 运行时 Actor 机制
2. 编译期一致性视角 Consistent View，通过 Placement + SBP + Boxing 解决了分布式易用性的问题。

对于分布式集群环境（多机多卡训练场景），OneFlow 会把整个分布式集群抽象成为一个超级设备，用户只需要在这个超级设备上搭建深度学习模型即可。这个虚拟出来的超级设备我们就称之为逻辑视角（Logical View），而实际上的分布式集群的多机多卡就是物理视角（Physical View），OneFlow维护逻辑视角和物理视角之间的数学上的正确性就称之为一致性视角（Consistent View）。

一致性视角（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的几种简单情形。

SBP 逻辑与物理 Tensor 的对应关系

需要注意的是，对于同一个逻辑上的 Tensor，其物理上的 Tensor 的映射关系可能会有多种，这取决于生产这个 Tensor 的 Op 和消费这个 Tensor 的 Ops 是如何看待这个逻辑上的 Tensor 的。

那么用 OneFlow 做数据并行、模型并行，需要用户做什么呢？其实只需要配置 Variable 的 SBP 即可。我们简单介绍一下数据并行和模型并行在 OneFlow 里的配置方式：

• 数据并行下，每个设备上都有整份的模型，所以 Variable 的 SbpParallel 是 Broadcast， 表示物理上的每个设备上的模型都是逻辑上的完整模型的复制。其余的用户就不用再做任何操作了（其实数据并行下，反向梯度更新的同步操作 AllReduce 是 OneFlow 系统内部根据 SBP 的推导规则自动插入的。）

• 模型并行下，每个设备都把模型切分并只保留一部分， 所以 Variable 的 SbpParallel 是 Split(0)， 表示物理上的每个设备上的模型都是逻辑上的完整模型经过第0维切分后的。其余的用户就不用再做任何操作了。前后向的数据同步操作也是 OneFlow 系统内部根据 SBP 推导规则自动插入的。

其实对于 Linear Layer 里的 Variable （假设是 row major 存储），Split(0) 和 Split(1) 分别表示了两种不同的模型并行方式。如果是 Split(0) 前后向就需要插入 AllReduce， 如果是 Split(1) ，前后向就需要插入 AllGather 操作了。至于为什么要插入 AllReduce 或者 AllGather，我会在下一章节介绍 SBP 推导的时候详细解释。另外，其实数据并行梯度更新要插入 AllReduce 做梯度同步，在 OneFlow 里也是自动推导出来的，并不是一个像 PyTorch DDP 一样的模块做特判。

而且 OneFlow 的 Consistent View 还保证了：任何配置 SBP 得到的并行结果， OneFlow 都保证了 其计算在数学上是完全一致的， 我们从机制上保证了分布式训练的正确性难题， 这一点是现在的 PyTorch 无法做到的。

3. 2D SBP

那么如何同时让一个 Op 既做数据并行，又做模型并行（分组）？这里就用到了 2-D SBP。

在 2-D SBP 下，（其实 OneFlow 还支持扩展到任意维度 N-D）集群和设备呈现为一个 2-D 的拓扑结构。比如我们一共有 2 机 8 卡（每台机器 4 张 GPU），我们可以将 8 个设备表示成一个 (2 x 4) 的矩阵， 那么 如何在机器间数据并行、机器内模型并行呢？用户只需要将 Variable 的 2-D SBP 配置成： [ Broadcast, Split(0) ] 即可，那么实际上 各个设备上的 物理 Tensor 跟 逻辑 Tensor 的映射关系如下图所示。在第一维上的 Broadcast， 表示 GPU0 和 GPU4、GPU1 和 GPu 5、 GPU2 和 GPU6、 GPU3 和 GPU7 在机器间做数据并行，在第二维上的 Split(0)， 表示 GPU0，1，2，3 、 GPU4，5，6，7 在机器内做模型并行。

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时，如果训练时间比较长（训练是整个网络的瓶颈），我们就会得到下面这种流水线的时间线：

数据预处理 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 表示参数规模。

模型并行数据对比

3.混合并行（数据&模型）性能对比

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

4.流水并行 + 混合并行（数据&模型） 性能对比

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