撰文 | 成诚

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 次梯度后更新一次模型）下的时间线如下图所示：

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

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 在机器内做模型并行。

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 中的隐藏层神经元数量 以及 输出神经元数量。

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

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

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

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

模型并行数据对比

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