训练 GPT-3，为什么原有的深度学习框架吃不消？

本文将重点梳理深度学习框架在支持大规模预训练模型时面临的技术挑战，以及当前各类框架的基本解决思路，帮助算法工程师对业界各类框架的分布式训练能力有更清晰的认知。

近年来，深度学习被广泛应用到各个领域，包括计算机视觉、语言理解、语音识别、广告推荐等。在这些不同领域中，一个共同的特点就是模型规模越来越大，比如 GPT-3 模型的参数量达到 1750 亿，即便拥有 1024 张 80GB A100， 那么完整训练 GPT-3 的时长都需要 1 个月。大规模预训练模型及其训练成为业界尤为关注的热点。

模型规模变大带来了哪些挑战？首先是硬件发展水平导致的内存墙问题。单一设备的算力及内存容量，受限于物理定律，持续提高芯片的集成越来越困难，难以满足大模型规模扩大的需要。为了解决算力增速不足的问题，人们考虑通过使用多节点集群进行分布式训练来提升算力，分布式训练则势在必行。

大型Transformer模型参数量和计算设备内存最近5年的增长速度

但是，简单的机器堆叠并不一定可以获取算力的增长，因为内存的带宽增长速率也大大落后于算力增长，跨计算设备之间的网络带宽更低，使得数据搬运成为整个训练的瓶颈，进而导致训练效率大大下降。

Transformer及计算机视觉、自然语言处理、语音识别模型算力增长速度

问题的核心是分布式编程门槛问题。大模型兴起时间较晚，原有的主流深度学习框架只支持数据并行，为了支持大模型的训练，人们可以对已有框架定制化地开发，针对解决某个特定领域内的分布式训练问题，如针对点击率预测的 HugeCTR，这降低了特定领域内算法工程师的门槛，但同时却削弱了通用性。主流深度学习框架也正在增加对大模型分布式训练的支持，但用户接口仍不够友好。在模型更大或者神经网络拓扑更复杂时，原有的主流深度学习通用框架的易用性和效率会大打折扣，开发者在分布式训练上的相当一部分重点需求仍未被满足。

无论是对已有框架做定制开发，还是让已有框架通用的支持大模型训练，都非易事。分布式训练涉及到性能调优、软硬件资源调度、通信带宽等工程问题，这些复杂的问题纠结在一起，造成开展超大规模预训练语言模型方面的突破式算法研究难度极高，长期以来被资金雄厚和具有扎实基础架构研发团队的科技巨头“垄断”。

在保证训练效率的前提下，降低技术门槛，让更多的算法工程师参与到大模型预训练模型中，创造更多价值，分布式训练成为包括 OneFlow 在内的业界深度学习框架共同的追逐目标。

因此，本文梳理了业界目前所采用的各类分布式并行策略，并探讨业界框架的分布式训练技术路径，帮助算法工程师对各类框架的分布式训练能力有较为清晰的认知。最后，本文还将简要介绍 OneFlow 分布式深度学习框架的发展脉络。

1、如何在各显神通的分布式并行模式中做出选择

依据网络在分布式集群中的切分方式，深度学习框架目前主要的分布式训练模式包括数据并行、模型并行和流水并行，乃至同时使用数据并行和模型并行的混合并行方法。

诚然，分布式训练通常会比单机单卡更快，且能支持更大规模的模型训练。然而，分布式训练的不同模式之间也会有优劣。在内存墙及网络墙的影响下，数据传输量成为影响分布式训练的速度以及收敛性的关键因素。

具体而言，目前不同框架对数据并行的支持也趋于成熟，性能差异不大。数据并行与模型并行相比各有优劣，简单来说，为了减少传输量，只要是训练时每个批次数据的中间结果的数量较大而参数数量少则选择数据并行，反之，若每个批次数据的计算中模型参数的数量较大而中间结果相对更少则选择模型并行。（数据并行相对于模型并行还有另外一个特别的优势，数据并行中的通信较容易和计算重叠，而模型并行中的通信不易被计算掩盖）。

一般来说，同一个神经网络的不同算子可能适合不同的并行模式，某个特定的算子只使用一种并行模式，例如在模型参数量大的地方使用模型切割，在模型参数量少的地方使用数据切割。相比于一个算子只使用单一的并行模式，一个算子也可以同时使用多样的并行模式可能进一步地减少传输量，譬如在隐藏层比较大的地方，就可能同时对数据矩阵切割以及对模型矩阵切割。

数据并行与模型并行都是让设备执行同一个层次的计算，而流水并行则是把任务划分为几个有明确先后顺序的阶段，把不同的阶段分给不同的计算设备，使得单设备只负责网络中部分层的计算。在这种多设备接力完成一个网络计算的模式下，可以支持更大的模型或者支持更大的 Batch Size。

为了应对数据并行与模型并行在内存及通信量上的局限，微软提出了 ZeRO-DP 的方式来解决，但并不能完全取代模型并行。2016 年，陈天奇团队提出亚线性内存优化、“gradient/activation checkpointing” 技术，旨在系统性减少深度学习训练中的内存占用，核心在于 “以时间换空间”，也可以大幅降低显存占用。

总的来说，流水并行在通常大模型训练情况下具有优势。流水并行的数据传输量少，仅为阶段之间需要传输的数据量之和，不像数据并行与模型并行那样大，传输量与整个计算图都有关，因此对于带宽较小的机器，会趋于使用流水并行。但某些情况下，流水并行与模型并行的结合则会优于单一的模型并行与流水并行。同时，在数据并行与模型并行中也存在计算时间掩盖传输时间的优化。

因此，对于每一个任务选择最优的并行模式是一个非常复杂的问题，需要具体情况具体分析。

2、支持大模型分布式训练的深度学习框架有何不同

既然分布式训练是大模型的核心，随着模型规模的扩大，原来不支持分布式训练的老牌深度学习框架通过增量式的逐步加入分布式训练功能，而 OneFlow 等框架则尝试通过重新设计新框架来系统解决分布式训练的难点。

考虑到老牌主流深度学习框架的影响力与市场积淀，目前一大部分方案是基于 TensorFlow 与 PyTorch 开发。

一些框架研究中，会关注如何设计并实现针对分布式的 API 接口，以求用更通用的方式满足各种复杂的分布式策略，如 Mesh-TensorFlow、 GShard。 Mesh-TensorFlow 是一种用于描述分布式任务的特定领域语言（DSL），使用 Mesh-TensorFlow 需要深度修改模型代码和训练过程，有一定上手门槛。 GShard 作为 XLA 编译器的扩展插件，借鉴了数据库技术中将大数据表切分（sharding）到不同物理设备上保存的思想。提供了 API（replicate、split、shard），用户通过它们注解张量，再由编译器完成张量到硬件设备上的映射。

还有一些工作，关注超大规模预训练模型（如 Transformer）带来的并行挑战，对模型并行、流水并行等并行策略的实现重点发力。如基于 Lingvo 开发的神经网络训练库 GPipe（Lingvo 是基于 TensorFlow 二次开发的重点针对序列模型的框架）、微软等的研究成果 PipeDream（早期版本基于 Caffe，目前基于 PyTorch 二次开发）、英伟达的 Megatron-LM（目前作为 PyTorch 的插件工作）、FaceBook 的 FairScale（PyTorch 的插件）、微软的 DeepSpeed 等。其中 Megatron-LM 和 DeepSpeed 是已开源方案中最成熟的两个，都是为大规模预训练模型高度定制的系统，不足之处在于，这些定制的功能难以应用到预训练模型之外的应用和算法上去。

此外，每个模型的并行策略候选集合是指数级的，纯手工从中挑出一种合适的并行策略，需要耗费算法工程师大量的时间以及计算资源。如何从框架层次自动解决并行策略选择问题，又称为自动并行技术，也是最近的研究热点。自动并行通过建立代价模型来预测并挑选一个较优的并行策略（暂时无法保证是最优的策略，因为挑出最优的策略是个 NP-Hard 的问题），有希望将算法工程师从并行策略的选择和配置中解放出来。

某些框架，如 FlexFlow 就进行了自动并行的尝试。FlexFlow 提供了一套简单可行的代码替换方案，使得 PyTorch 和 TensorFlow Keras 的用户只需要少量需改代码，就可以运行 FlexFlow。

OneFlow 则通过重新设计和开发，试图系统地解决分布式训练的效率、易用性等问题。OneFlow 在系统设计之初就通过 Actor 模型将数据搬运作为整体计算图的一部分，方便全局优化；OneFlow 还发明了 SBP 概念（与 Google GShard 有一些重叠，但是 GShard 的超集），让算子的数学逻辑与分布式训练时物理设备的分配解耦，将分布式集群抽象成逻辑上的“超级计算机”，使算法科学家可以像操作单机单卡设备那样完成分布式训练。

同时，OneFlow 的自动并行完全不需要用户去做复杂的配置，也允许用户对某几个计算节点定制并行模式。 （详细内容，请查看《GPT-3模型为何难以复现？这也许是分布式AI框架的最优设计》）

3、OneFlow 致力于解决分布式训练大模型的难题

不同于原有的深度学习框架以一种增量式、渐进的方式迭代自己的设计，逐步从支持单设备，到支持数据并行，再试图支持模型并行、流水并行等更挑战的问题，作为新一代深度学习框架，OneFlow 自成立的第一天就致力于以一种通用的方式解决大数据、大计算和大模型的难题，自设计之初就纳入了支持数据并行、模型并行和流水并行的想法。

我们相信，OneFlow 的技术路线是解决深度学习横向扩展难题的必由之路。实际上，在 OneFlow 走通这条路径之后，其他一些技术社区的团队已经开始沿着这个方向进发。我们获得的经验是，创新和创造是我们最可依赖的法宝，仅仅跟随已有框架走过的路是不可能实现超越的，唯有创新才有机会。

无疑，OneFlow 团队会持续做出积极的创新和探索。