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）。

*博客内容为网友个人发布，仅代表博主个人观点，如有侵权请联系工作人员删除。