为了高效地搜索如此庞大的搜索空间，该研究首先提出了以下观察作为指导：
- 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 种。
下图给出了 8 卡 GPU 场景不同 PP 并行度（8/4/2/1）下的决策树示意图。
8 卡 GPU 场景不同 PP 并行度（8/4/2/1）下的决策树示意图
2. 基于动态规划的并行优化算法
现有系统如 Megatron 或者 DeepSpeed 通常由用户指定全局并行方案及其对应的并行度，严重限制了对分布式执行计划的表达能力。Galvatron 的优化目标是在用户给定模型定义和分布式环境的情况下，无需用户指定任何并行配置，自动生成最优的分布式执行计划。具体来说，给定一个 L 层的模型 M 和内存容量为 E 的 N 个 GPU 设备，Galvatron 的优化目标是搜索得到最高的系统吞吐 T_pt 并返回对应的并行方案，这里的并行方案指以层（或算子）为基础单位的细粒度混合并行策略。
算法 1: Galvatron 优化过程
优化过程：Galvatron 的优化过程如算法 1 所示。Galvatron 最外层循环逐步增加搜索的 batch size 直到超出设备内存；给定每个候选 batch size B，Galvatron 首先根据 Takeaway#1 对模型进行 PP 切分并搜索不同的并行度 P（第 4 行），选定 P- way PP 后模型被切分为 P 个阶段（第 6 行），与之对应的所有设备被切分为 P 组，每组含有 N/P 个设备；接着 Galvatron 构建对应的决策树，其能不重不漏地表示出 DP、SDP、TP 的任意组合，从而获得策略集合 S；然后对于每个模型阶段 M_i，在设备内存限制 E 下，Galvatron 使用动态规划搜索得到其中每一层的最优混合并行策略并返回最小时间开销（第 9 行）；最后，Galvatron 在所有可能的 PP 并行度和 batch size 中选择吞吐率最高的策略并返回（第 15 行）。
动态规划搜索：下面介绍 Galvatron 并行优化工作流中的动态规划搜索算法。对于一个给定的包含 L 层的模型阶段，用开销函数 C(L,E) 表示 L 层模型在设备内存限制 E 下的总执行时间，
表示第 L 层使用策略 S_j 的执行时间，其中策略 S_j 是并行策略候选集合 S 中的策略。设定初始值
，Galvatron 的动态规划搜索遵循以下状态转移方程（公式 1）：
其中，
是第 L 层使用策略 S_j 的内存开销，
是第 L 层使用策略 S_j、其前一层使用策略 S_i 带来的转换开销。在状态转移过程中，当内存开销超过设备内存限制设备内存限制 E，开销函数 C 返回无穷大。
复杂度分析：Galvatron 使用的动态规划搜索（公式 1）计算复杂度为 O(LE|S|)。可见，每一层的搜索空间 S 的大小对于整体搜索复杂度非常重要，Galvatron 提出的基于决策树的搜索空间分解能够显著减小搜索空间，控制搜索开销在合理范围内。
3. 基于混合建模的执行代价估计方法
Galvatron 使用一个策略开销估计模块来对混合并行策略的计算、通信、内存开销进行估计。现有的开销估计方法主要包括测量（profiling）和模拟（simulating）两种，Galvatron 吸取了两者的长处并设计了一种代价低廉、高效且准确的开销估计方法。具体来说，对于内存开销，Galvatron 使用张量的形状和数据类型来完成估计；对于计算时间，Galvatron 通过在单设备上通过 profiling 测量逐样本计算时间，结合 batch size 和拟合函数从而估计总体计算时间；对于通信时间，Galvatron 通过通信量除以设备通信带宽得到预估的通信时间，其中通信量由理论计算得到，通信带宽由 profiling 测量得到。
基于以上估计结果，Galvatron 通过 simulating 执行过程来计算给定层使用给定策略的开销 c(l,s)。和现有分布式训练系统的代价模型不同，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 作为实验对象。
Galvatron 和基线系统在 8 GPUs 20G 显存下的吞吐对比情况
实验对比效果：该研究首先在八卡 Nvidia RTX TITAN 24GB 环境下进行了实验。实验表明，在不同模型规模和不同内存限制下，Galvatron 相比于总是取得最优的吞吐率，与现有的最先进的单一并行方法和混合并行方法相比，训练吞吐量均有显著提升。具体而言，在 ViT 模型上，Galvatron 相比于单一策略的吞吐率加速比最高可达 338%，相比于其他混合并行策略的吞吐量从加速比最高可达 55%；在其他三个模型上，Galvatron 相比于单一策略和现存的混合策略加速比最高可达 200%-334% 和 28%-52%。
Galvatron 搜索得到的部分最优并行策略示意图
可解释性实验：该研究选取了一些由 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 具有良好的可扩展性，详细结果可以参见论文原文。
北大河图团队简介
河图（Hetu）开发团队来自于北京大学数据与智能实验室（Data and Intelligence Research Lab at Peking Univeristy，以下简称：实验室），实验室由北京大学计算机学院崔斌教授带领，多年来主要在人工智能、大数据等领域进行前沿研究，在理论和技术创新以及系统研发上取得多项成果，已在国际顶级学术会议和期刊发表学术论文 100 余篇。
河图系统是一个面向超大模型的分布式深度学习系统，相比现有的老牌分布式深度学习框架，在系统功能性、系统复杂性和系统易用性上有诸多创新贡献，如自动分布式并行策略、一致性协议和通信架构、GPU 算子优化等方面。Hetu 团队目前已在多种分布式机器学习或深度学习场景下进行了学术创新，相关成果被 SIGMOD、VLDB、ICML、KDD 等国际顶级会议收录，其中稀疏大模型分布式训练系统 HET 荣获 VLDB 2022 最佳可扩展数据科学论文奖。本次被 VLDB 2023 录用的论文 Galvatron 则是河图团队在稠密大模型分布式训练场景上的又一突破成果，已经集成到 Hetu 系统并对外开源。目前，Hetu 团队已经与腾讯、阿里巴巴、快手、字节跳动等多家知名企业展开科研合作和应用落地。
