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.小结

感兴趣的可以继续戳原文。

来源：公众号「OneFlow」

内容中包含的图片若涉及版权问题，请及时与我们联系删除