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 一致，且数据并行的并