论文链接：https://arxiv.org/abs/2211.13878
代码链接：https://github.com/PKU-DAIR/Hetu/tree/main/tools/Galvatron
一、研究背景
最近一段时间，「大模型」在 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。
二、大模型自动并行到底难在哪里
研究者认为，大模型的自动并行之难主要体现在以下三个方面：
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 用户带来额外的使用负担；同时也不需要用户付出额外的系统安装或者调试代价，使用时只需要添加几行代码，就可以轻松完成自动并行的整个流程。
四、关键技术
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 卡的设备。在这种情况下，一个设备组内的最优混合并行策略与其他组内的最优策略保持一致。
基于以上重要观察，该研究提出了一种基于决策树的搜索空间构建方法：
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 种。
下图给出了 8 卡 GPU 场景不同 PP 并行度（8/4/2/1）下的决策树示意图。
2. 基于动态规划的并行优化算法
现有系统如Megatron或者DeepSpeed通常由用户指定全局并行方案及其对应的并行度，严重限制了对分布式执行计划的表达能力。Galvatron的优化目标是在用户给定模型定义和分布式环境的情况下，无需用户指定任何并行配置，自动生成最优的分布式执行计划。具体来说，给定一个L层的模型M和内存容量为E的N个GPU设备，Galvatron的优化目标是搜索得到最高的系统吞吐
并返回对应的并行方案，这里的并行方案指以层（或算子）为基础单位的细粒度混合并行策略。
复杂度分析：Galvatron使用的动态规划搜索（公式1）计算复杂度为
。可见，每一层的搜索空间S的大小对于整体搜索复杂度非常重要，Galvatron提出的基于决策树的搜索空间分解能够显著减小搜索空间，控制搜索开销在合理范围内。
3. 基于混合建模的执行代价估计方法
Galvatron使用一个策略开销估计模块来对混合并行策略的计算、通信、内存开销进行估计。现有的开销估计方法主要包括测量（profiling）和模拟（simulating） 两种，Galvatron吸取了两者的长处并设计了一种代价低廉、高效且准确的开销估计方法。具体来说，
- 对于内存开销，我们使用张量的形状和数据类型来完成估计；
- 对于计算时间，我们通过在单设备上通过profiling测量逐样本计算时间，结合batch size和拟合函数从而估计总体计算时间；
- 对于通信时间，我们通过通信量除以设备通信带宽得到预估的通信时间，其中通信量由理论计算得到，通信带宽由profiling测量得到。
基于以上估计结果，Galvatron通过simulating执行过程来计算给定层使用给定策略的开销
和现有分布式训练系统的代价模型不同，Galvatron在建模中首次考虑了计算和通信的重叠对GPU性能下降带来的影响。我们通过实验发现，由于重叠导致的GPU性能下降会显著影响系统的执行效率，而这在以往的工作中都是被忽视的。因此，Galvatron的开销估计更加准确，并行优化效果更佳。
五、实验结果
实验设置： 在实验中，我们将Galvatron和四种使用单一并行策略的基线系统（DP、SDP、TP、PP）以及专家设置的DeepSpeed 3D Parallelism进行对比，同时额外设置了两个弱化版本的Galvatron作为辅助基线，在有限的并行策略组合空间（即TP+DP，PP+DP），展开自动并行搜索。我们选用了NLP Transformer模型BERT和T5，CV Transformer模型ViT和Swin Transformer作为实验对象。
实验对比效果： 我们首先在八卡Nvidia RTX TITAN 24GB环境下进行了实验。实验表明，在不同模型规模和不同内存限制下，Galvatron相比于总是取得最优的吞吐率，与现有的最先进的单一并行方法和混合并行方法相比，训练吞吐量均有显著提升。具体而言，在ViT模型上，Galvatron相比于单一策略的吞吐率加速比最高可达338%，相比于其他混合并行策略的吞吐量从加速比最高可达55%；在其他三个模型上，Galvatron相比于单一策略和现存的混合策略加速比最高可达200%-334%和28%-52%。
可解释性实验： 我们选取了一些由Galvatron搜索得到的最优并行策略进行展示。对于BERT模型在8GB情况下(Case A)，Galvatron选择了两种混合并行策略PP-TP-DP以及PP-TP-SDP，而当可用显存增大到12GB时，Galvatron放弃了PP，选择使用更多的DP，同时引入SDP来节省显存空间。在Swin Transformer上情况则稍有不同，该模型的不同层呈现出明显的异质性，当显存比较紧缺时（Case C），浅层SDP的并行度较高，随着层数加深，每层的激活变小，参数变多，所以TP逐渐取代SDP。当显存增多时（Case D）,不仅重新启用了PP替代了一部分低效的SDP，而且浅层倾向于使用DP的趋势更为明显。
可扩展性实验： 我们进一步在更大的集群上测试了Galvatron，包括16卡Nvidia RTX TITAN GPUs的环境和64卡Nvidia A100 GPUs的环境。在16卡环境下，Galvatron相比于其他策略依然拥有最高吞吐率，与相同内存限制的8卡实验结果相比，由于更加多样化的混合并行策略，Galvatron在16卡上能获得超过2倍的加速比。在64卡实验中，Galvatron的吞吐率同样是其他策略中最高的。这说明了Galvatron具有良好的可扩展性，详细结果可以参见论文原文。
六、总结
以Transformer代表的「大模型」为分布式深度学习系统带来了巨大的机遇和挑战，我们提出的Galvatron系统是目前业界首个集成众多并行维度（包括数据并行、分片数据并行、张量并行、流水并行）的自动并行分布式训练系统，相比于传统系统DeepSpeed、Megatron，以及搜索空间受限的其他基线系统，在CV和NLP领域的主流Transformer模型上，Galvatron都展现出了显著的性能优势。目前该系统已经开源（https://github.com/PKU-DAIR/Hetu/tree/main/tools/Galvatron），用户无需了解分布式训练的专业知识，就可以一键完成大模型的自动并行高效训练。未来我们将在Galvatron基础上持续集成更多深度学习系统优化技术，在最大化硬件利用率的同时保证用户便捷的使用体验。
北大河图团队简介
河图（Hetu）开发团队来自于北京大学数据与智能实验室（Data and Intelligence Research Lab at Peking Univeristy，以下简称：实验室），实验室由北京大学计算机学院崔斌教授带领，多年来主要在人工智能、大数据等领域进行前沿研究，在理论和技术创新以及系统研发上取得多项成果，已在国际顶级学术会议和期刊发表学术论文 100 余篇。
河图系统是一个面向超大模型的分布式深度学习系统，相比现有的老牌分布式深度学习框架，在系统功能性、系统复杂性和系统易用性上有诸多创新贡献，如自动分布式并行策略、一致性协议和通信架构、GPU 算子优化等方面。Hetu 团队目前已在多种分布式机器学习或深度学习场景下进行了学术创新，相关成果被 SIGMOD、VLDB、ICML、KDD 等国际顶级会议收录，其中稀疏大模型分布式训练系统 HET 荣获 VLDB 2022 最佳可扩展数据科学论文奖。本次被 VLDB 2023 录用的论文 Galvatron 则是河图团队在稠密大模型分布式训练场景上的又一突破成果，已经集成到 Hetu 系统并对外开源。目前，Hetu 团队已经与腾讯、阿里巴巴、快手、字节跳动等多家知名企业展开科研合作和应用落地。
作者：苗旭鹏Illustration by Pixel True from IconScout
快抢沙发！
