VLDB 2023 | 北大河图发布分布式训练神器Galvatron, 一键实现大模型高效自动并行

机器之心 2022-11-30 20:53 发表于北京

机器之心专栏

机器之心编辑部

北大河图团队提出了一套面向大模型的自动并行分布式训练系统Galvatron，相比于现有工作在多样性、复杂性、实用性方面均具有显著优势，论文成果已经被 VLDB 2023 接收。

最近一段时间，「大模型」在 AI 领域的各种应用场景都大放异彩，其中基于 Transformer 的大规模预训练模型是最典型的大模型之一，已经成为了当前基础模型（Foundation Model）的核心架构。例如 NLP 领域的 BERT、GPT 系列，CV 领域的 ViT、Swin Transformer 系列，以及最近非常火热的多专家混合模型 MoE 以及多模态模型 CLIP 等等，都是以 Transformer 作为核心的基础架构。与之相应的，这类稠密大模型拥有着动辄数十亿、百亿甚至万亿规模的参数量，面临高昂的计算、存储、以及通信开销，也为 AI 基础设施带来了巨大的挑战。

为了支持大模型的训练，人们研发了很多工具（如 NVIDIA 提出的「威震天」Megatron，Microsoft 提出的 DeepSpeed，Meta 提出的 FairSeq 等等）来实现各种各样的并行方式，数据并行、张量模型并行、流水并行、分片数据并行等等。这些系统对上述并行方式提供了良好的封装，对外屏蔽相应的实现细节，使得用户能够通过添加配置的方式来实现混合并行策略。

基于上述思想，已经有很多工作围绕如何在张量或者算子层面表达各种并行方式进行探究，这类工作的 “自动化” 主要体现在并行 API 到执行层的转化过程。但是如果只是局限在设计并行 API 或者中间表达上，这种工程上的封装并没有从根本上解决分布式训练的问题。最直观的结果就是，用户仍然无法从分布式部署的难题中解放出来。实际上，大模型的分布式部署是一个非常复杂的问题，目前的绝大多数的分布式训练系统，都依赖用户人工反复尝试以及系统专家经验来进行部署，造成严重的资源利用效率低下的问题，与真正的 “自动并行” 存在相当大的差距。

基于此，北大河图团队提出了一款分布式训练神器 Galvatron，以实现大模型的高效自动并行，研究论文入选国际顶会 VLDB 2023。

• 论文地址：https://arxiv.org/abs/2211.13878
  
• 项目代码链接：https://github.com/PKU-DAIR/Hetu/tree/main/tools/Galvatron
  

大模型自动并行到底难在哪里

研究者认为，大模型的自动并行之难主要体现在以下三个方面：

（1）多样性：首先，在并行方式方面，目前大模型的并行方式呈现出百花齐放的态势，即使是对于同一个算子，不考虑混合并行方式，不同的基础并行方式也会存在显著的差异，从而导致不同的内存开销、通信代价以及计算效率。下图展示了四种最主要的基础并行方式，即数据并行（Data Parallelism）、张量模型并行（Tensor Parallelism）、流水模型并行（Pipeline Parallelism）、分片数据并行 (Sharded Data Parallelism)，在两张 GPU 上对于简单的矩阵乘法算子分布式执行的过程。

并行方法对比示意图

其次，在模型方面，各种各样的模型架构最近也是层出不穷，这往往也伴随着不同的模型配置（例如不同输入序列长度，模型层数，模型隐层宽度等），从而造成计算负载上的差异。另外，在硬件方面，用户往往面临着非常差异化的集群环境，可能会面临不同的内存容量、通信带宽、计算能力等等。总体上来看，由于上述多样性的存在，没有哪种并行技术总是能够获得最佳训练效率，“自动并行” 也就成为了分布式训练的核心挑战。

（2）复杂性：上述分析还相对比较单一，实际上哪怕是对于同一个算子也可以同时应用多种不同的基础并行方式，如果考虑到由这些基础并行方式复合所构成的混合并行方式，则会导致问题变得非常复杂。更重要的是，大模型的计算图往往结构非常庞大，对应的也需要更大规模的集群，如果对每个算子都进行探索（包括选取集群中合适的计算资源以及设计相应的混合并行方式），会带来组合空间爆炸的问题，寻找整个模型的最优分布式执行方案变得难以求解。

（3）实用性：除此之外，实用性也是非常重要的问题。一方面，在进行自动并行搜索的过程中，对于各种分布式执行方案，必须提供比较精确的内存、通信、计算开销，否则会导致结果与实际执行偏差过大，产生次优解或者根本无法使用。为此，就需要非常精准的代价模型，对不同的模型结构和硬件条件进行建模。另一方面，系统提供自动并行能力所带来的额外时间开销必须在一个可以接受的范围内，过于高昂的搜索代价同样也无法接受。

分布式训练神器 Galvatron，一键实现大模型高效自动并行

系统特性：

为了解决上述问题，研究者们提出了一些系列工作来探索混合并行的自动搜索：一类工作主要讨论了同时考虑数据并行和模型并行的搜索空间，代表性工作包括 FlexFlow，Tofu，另一类工作则产生于流水并行场景，将其与数据并行相结合，代表性工作包括 PipeDream，DAPPLE。在此基础上还有一些衍生工作，如 Unity、Alpa，进一步扩展了自动并行的探索范围。北大河图团队提出的系统「惊破天」Galvatron 同样属于自动并行搜索的研究领域，但相比于现有工作，该系统主要拥有以下三方面优势：

（1）在多样性方面，现有工作能够支持的并行维度仍然比较有限，而 Galvatron 不仅可以支持更多的并行维度，并且面对更加差异化的 Transformer 模型结构也能够精准建模，以及在不同的集群硬件条件下验证了其自适应调优的能力。

大模型分布式训练系统对比示意图

（2）在复杂性方面，由于具备在多样性上的优势，Galvatron 面临了前所未有的庞大搜索空间，为此，研究者探究了几条目前大规模分布式训练过程中的一些经过实验性或理论性验证的重要观察，作为搜索空间的剪枝准则，从而实现高效的分布式执行计划优化。

（3）在实用性方面，该研究结合了理论建模和实验测量两者的优势，实现对内存、通信、计算开销的精确估计，甚至考虑到了计算与通信重叠所导致的 GPU 执行效率下降问题，从而保证能够得到足够准确的自动并行优化结果。

另外，Galvatron 底层选择 PyTorch 作为执行引擎，兼容 Huggingface 等常见的主流 Transformer 模型实现，所以完全不会对 PyTorch 用户带来额外的使用负担；同时也不需要用户付出额外的系统安装或者调试代价，使用时只需要添加几行代码，就可以轻松完成自动并行的整个流程。

Galvatron 工作流程及用户使用接口展示

关键技术

1. 基于决策树的搜索空间分解

Galvatron 的设计目标是高效地在复杂而庞大的并行策略空间内自动搜索，并为给定的 Transformer 模型和分布式环境生成最佳的并行执行计划。在搜索空间上，Galvatron 是业界首个考虑 4 种主流并行方法的自动并行训练系统，包括包括数据并行（data parallelism, DP）、分片数据并行（sharded data parallelism, SDP）、张量并行（tensor parallelism, TP）和流水并行（pipeline parallelism, PP）。由于混合并行策略会包含以上四种并行算法的任意组合，在多 GPU 的场景下这种组合带来的搜索空间十分庞大。例如，对于双机四卡场景，一种可行的策略是在机内使用 2-way TP，机间使用 2-way PP，另一种可行的策略是在机内使用 2-way PP，机间使用 2-way DP。当节点内 GPU 数量扩展到 8 卡时，模型每一层的候选策略都多达上百种。随着模型层数的增加，其搜索空间规模指数增长，难以有效探索。

为了高效地搜索如此庞大的搜索空间，该研究首先提出了以下观察作为指导：

• Takeway#1：PP 倾向于被跨设备岛放置。此处 “设备岛” 指具有高内部带宽的一组设备，在绝大多数 Transformer 模型中，PP 的通信量相比于其它并行方式，显著更少。因此，人们通常优先对模型进行 PP 切分并放置于设备岛之间。
  
• Takeway#2：在同构设备的前提下，并行策略倾向于将设备均匀切分。例如，对于 4 卡 GPU 的 2-way DP 倾向于将设备切分两组 2 卡的设备，而不是一组 1 卡和一组 3 卡的设备。在这种情况下，一个设备组内的最优混合并行策略与其他组内的最优策略保持一致。
  
• Takeway#3：一般来说，在能够混合使用 DP 和 SDP 的情况下，只使用 SDP 在理论上性能更优。根据分析结果，N-way SDP 的通信开销和内存开销均优于 和 的组合，其中 。
  

基于以上重要观察，该研究提出了一种基于决策树的搜索空间构建方法：

（1）给定一个 Transformer 模型，基于 Takeway#1 和 Takeway#2，Galvatron 首先用 PP 将模型切分成多个阶段，同时将设备均匀且连续地切分为多个设备组。例如 8 卡场景下，模型被切分为 1/2/4/8-way PP，分别对应设备组大小为 8/4/2/1。

（2）每种 PP 切分对应着一棵决策树及一个子搜索空间，决策树叶结点总数为设备组大小，决策树高度为可用的并行方法数量，即决策树每一层可应用一种并行策略。

（3）并行策略在决策树不同层间不可重复使用。

（4）非叶结点的度数默认在 2 的指数次幂 {2,4,8,…} 中选择。

基于以上决策树构建规则，Galvatron 构建的决策树能表示以上并行的任意组合。Takeway#1 和 Takeway#2 帮助 Galvatron 规避了低效的并行组合，缩小了搜索空间。对于 8 卡 GPU 训练一层模型的场景，以上规则将产出 34 种候选的混合并行策略。进一步，使用 Takeway#3 将 DP 和 SDP 同时出现在一棵决策树的情形剪枝后，8 卡候选策略数降至 22 种。