Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning

文章标题：Alpa：自动化分布式深度学习中的算子间和算子内并行
作者/机构：
Lianmin Zheng, Zhuohan Li, Hao Zhang (加州大学伯克利分校); Yonghao Zhuang (上海交通大学); Zhifeng Chen, Yanping Huang, Yuanzhong Xu (谷歌); Yida Wang (亚马逊网络服务); Danyang Zhuo (杜克大学); Eric P. Xing (MBZUAI 和卡内基梅隆大学); Joseph E. Gonzalez, Ion Stoica (加州大学伯克利分校)

A1 主要贡献

核心问题: 训练如GPT-3等拥有数千亿参数的超大规模深度学习模型，需要大量的工程努力，这些努力高度依赖于特定的模型定义和集群环境。例如，为Transformer模型选择和调整数据、算子和流水线并行等多个并行维度，或为MoE模型在不同集群（TPU vs GPU）上设计并行方案，都需要深厚的机器学习和系统专业知识。手动调整并行策略虽然能带来数量级的性能提升，但其复杂性阻碍了模型开发者快速探索新的模型设计。

研究目标: 自动化大规模模型的并行化，旨在让模型开发者无需关心底层的系统挑战，从而显著加速机器学习的研究和生产。这需要在一个随着并行维度、模型大小和集群规模呈指数级增长的复杂方案空间中进行导航。现有的自动并行化工作【17, 38, 55】通常局限于单一的模型并行方法或对模型和集群做出过强假设，无法充分利用性能优化的机会。

核心思想与创新点: 本文的核心观察是，可以将不同的并行技术组织成一个层级化空间，并将这些技术映射到计算集群的层级结构上。不同的并行技术对通信带宽有不同的要求，而典型的计算集群也具有相应的结构：物理位置相近的设备通信带宽较高，而远距离设备间的通信带宽有限。

基于此观察，本文将机器学习并行方法重新划分为算子内（intra-operator）并行和算子间（inter-operator）并行两大类：
* 算子内并行：沿一个或多个张量轴（批次或非批次）对算子进行分区，并将分区后的计算分派到分布式设备上执行（图1c）。其特点是设备利用率高，但每次训练迭代中，在算子的分裂和合并处都需要通信。
* 算子间并行：将模型切分为多个不相交的阶段（stage），并在不同的设备集上进行流水线执行（图1d）。其特点是只在相邻阶段间通信，通信量可以很小，但由于调度约束会产生设备空闲时间。

Alpa利用这种层级化设计，将通信密集的算子内并行映射到高带宽连接的设备上，而将算子间并行编排在通信带宽相对较低的远距离设备之间。这种设计允许将复杂的并行优化问题分解为两个可解的子问题，从而在每个层级上近似最优地求解。

主要贡献总结:
1. 构建了一个双层并行执行计划空间（图1e）：该空间使用算子间和算子内并行来层级化地定义并行计划。
2. 设计了可行的优化算法：在每个并行层级上推导出近似最优的执行计划。
3. 实现了Alpa，一个面向GPU集群的分布式深度学习编译器系统：
* 包含一套编译遍（compilation passes），使用层级优化算法生成执行计划。
* 一个新的运行时架构，用于编排设备网格（device mesh）之间的算子间并行。
* 多项系统优化，以改进编译和解决跨网格通信问题。
4. 全面的评估：在拥有64个GPU的亚马逊EC2集群上进行的评估表明，Alpa生成的并行计划能够匹配甚至超越为特定模型（如GPT）手动调优的系统（如Megatron-LM）。对于没有手动设计策略的模型（如Wide-ResNet），Alpa也能很好地泛化，并能显著加速MoE等复杂模型的训练（相比DeepSpeed在4节点上提速9.7倍），实现了开箱即用的高效模型并行执行。

A3 背景知识与设计原则

2.1 传统的机器学习并行视角

深度学习（DL）计算通常表示为数据流图，其中边代表多维张量，节点代表计算算子。一次训练迭代包括前向传播计算损失、反向传播推导更新以及权重更新。当模型或数据过大时，需要使用分布式并行方法。现有的方法通常分为数据并行、算子并行和流水线并行。

• 数据并行：将训练数据在多个工作节点（worker）上分区，但模型是复制的。每个工作节点独立计算其数据分片上的参数更新，并在权重更新前与其他节点同步更新，以确保所有工作节点上的模型参数保持一致。

• 算子并行：当模型大到无法装入单个设备时，算子并行是一种有效的模型并行方式。它指的是沿非批次轴对特定算子（如下文中的op）的计算进行分区，例如图2b中的矩阵乘法（matmul），并在多个设备上并行计算算子的每个部分。由于输入张量是联合分区的，一个设备在计算其算子分区时可能需要从其他设备获取输入数据，这通常需要集体通信操作，如all-reduce、all-gather和all-to-all。

• 流水线并行：与对算子进行分区不同，流水线并行将模型图中的不同算子组（称为阶段, stage）放置在不同的工作节点上。同时，它将训练批次（batch）拆分为多个微批次（microbatch），并在分布式工作节点上对这些微批次的前向和后向传播进行流水线处理，如图2d所示。流水线并行使用点对点通信在不同工作节点之间传递中间激活值。

• 手动的并行组合：为了扩展当今的大型深度学习模型，需要将上述方法结合起来【40, Efficient large-scale language model training on gpu clusters using megatron-lm, 2021, Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】, 【57, Gspmd: General and scalable parallelization for ml computation graphs, 2021, arXiv】。最先进的训练系统，如Megatron-LM【40, Efficient large-scale language model training on gpu clusters using megatron-lm, 2021, Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】, 【49, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019, arXiv】，为Transformer语言模型手动设计了一种结合了这些并行方式的专门执行计划，也称为3D并行。这种方法假设模型由重复的相同Transformer层构成，为每个流水线阶段分配相同数量的层，并为所有层统一应用手工设计的算子和数据并行配置。这种手动计划不仅需要深厚的专业知识，而且无法泛化到不同的模型或集群配置。

• 自动的并行组合：各种并行方式的配置、它们之间的相互依赖关系以及它们对模型和集群设置的依赖，构成了一个难以处理的巨大空间，这使得自动组合这些并行方式变得非常困难。例如，将算子并行与数据并行结合时，每增加一个数据并行副本就需要分配一组新的设备（而不是单个设备），并需要确定这些设备内部最优的算子并行配置。当引入流水线并行时，最优的流水线方案取决于每个流水线阶段的数据和算子并行选择，以及如何为每个阶段分配设备。在这种传统视角下，先前对自动并行的探索【17, Dapple: A pipelined data parallel approach for training large models, 2021, Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming】, 【25, Beyond data and model parallelism for deep neural networks, 2018, arXiv】, 【55, Supporting very large models using automatic dataflow graph partitioning, 2019, Proceedings of the Fourteenth EuroSys Conference 2019】, 【60, Autosync: Learning to synchronize for data-parallel distributed deep learning, 2020, Advances in Neural Information Processing Systems】仅限于将数据并行与至多一种模型并行方法相结合，从而错失了大量的性能提升机会。

2.2 算子内（Intra-Operator）和算子间（Inter-Operator）并行

一种新的并行划分方法。与传统视角不同，本文将现有的并行方法重新划分为两个正交的类别：算子内并行和算子间并行。它们的区别在于是否涉及沿任何张量轴对算子进行分区。

• 算子内并行：一个算子作用于多维张量。我们可以沿某些维度对张量进行分区，将分区后的计算分配给多个设备，并让它们同时执行算子的不同部分。我们将所有使用此工作流程的并行方法定义为算子内并行。图2a-c展示了在MLP上应用的几种典型算子内并行实例。数据并行【29, One weird trick for parallelizing convolutional neural networks, 2014, arXiv】根据定义属于算子内并行——输入张量和矩阵乘法沿批次维度分区，而权重张量被复制。或者，当权重非常大时，对权重进行分区（图2b）就得到了Megatron-LM中采用的算子并行。除了前向或后向传播中的算子，还可以对权重更新阶段的算子进行分区，从而产生权重更新分片（weight update sharding）或等效的ZeRO【44, Zero: Memory optimizations toward training trillion parameter models, 2020, SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】, 【56, Automatic cross-replica sharding of weight update in data-parallel training, 2020, arXiv】技术。由于分区，算子的分裂和合并处需要集体通信。因此，算子内并行的一个关键特征是它会在分布式设备之间产生大量通信。

• 算子间并行：我们将算子间并行定义为不执行算子分区，而是将图中的不同算子分配到分布式设备上执行的正交方法类别。图2d展示了将批次拆分的流水线并行作为算子间并行的一个案例。流水线执行可以遵循不同的调度策略，如Gpipe【22, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019, Advances in neural information processing systems】, PipeDream【38, Pipedream: generalized pipeline parallelism for dnn training, 2019, Proceedings of the 27th ACM Symposium on Operating Systems Principles】和同步1F1B【17, Dapple: A pipelined data parallel approach for training large models, 2021, Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming】, 【39, Memory-efficient pipelineparallel dnn training, 2021, International Conference on Machine Learning】。本文始终采用同步1F1B调度，因为它遵循同步一致性，并且与Gpipe相比具有相同的流水线延迟但峰值内存使用更低。在算子间并行中，设备仅在流水线阶段之间通信，通常使用设备对之间的点对点通信。所需的通信量可以远小于算子内并行中的集体通信。无论使用何种调度，由于阶段之间的数据依赖性，算子间并行会导致一些设备在前向和后向计算期间处于空闲状态。

新分类的优势。通过这种分类，两种并行方式发生在深度学习计算的不同粒度上，并具有不同的通信要求，这恰好与当今计算集群的结构相匹配。我们将利用这些特性来设计层级化算法和编译遍，以自动生成执行计划。一些同期的工作【2, Varuna: scalable, low-cost training of massive deep learning models, 2022, Proceedings of the Seventeenth European Conference on Computer Systems】, 【33, Terapipe: Token-level pipeline parallelism for training large-scale language models, 2021, arXiv】, 【39, Memory-efficient pipelineparallel dnn training, 2021, International Conference on Machine Learning】, 【50, Piper: Multidimensional planner for dnn parallelization, 2021, Advances in Neural Information Processing Systems】提出了类似的分类，但Alpa是第一个利用这种分类从完整空间中自动生成并行计划的端到端系统。

3. 概览

Alpa的层级化优化。Alpa是一个编译器，它通过在两个不同层级（算子内和算子间并行）上层级化地优化计划来生成模型并行执行计划。在算子内层级，Alpa最小化在给定设备网格（device mesh，一组设备，它们之间可能有高带宽连接）上执行计算图的一个阶段（即子图）的成本，该成本是相对于其算子内并行计划而言的。不同的网格可能根据分配的工作负载拥有不同数量的计算设备。在算子间层级，Alpa最小化算子间并行延迟，这涉及到如何将模型和设备集群切分为阶段和设备网格，以及如何将它们映射为阶段-网格对。算子间优化依赖于算子内优化器报告的每个阶段-网格对的执行成本。通过这个层级优化过程，Alpa生成的执行计划由算子内和算子间计划组成，这些计划在各自的层级上是局部近似最优的。

Alpa的编译流程。为了实现这一点，Alpa实现了三个新颖的编译遍，如图3所示。给定一个模型描述（以Jax【9, JAX: composable transformations of Python+NumPy programs, 2018, http://github.com/google/jax】的中间表示（IR）形式）和一个集群配置，算子间编译遍将IR切分为多个阶段，并将设备集群切分为多个设备网格。算子间遍使用动态规划（DP）算法将阶段分配给网格，并对每个阶段-网格对调用算子内编译遍，以查询此分配的执行成本。一旦被调用，算子内遍会优化该阶段在其分配的网格上运行的算子内并行执行计划，通过使用整数线性规划（ILP）公式最小化其执行成本，并将成本报告给算子间遍。通过为每个阶段-网格对的分配重复查询算子内遍，算子间遍使用DP来最小化算子间并行执行延迟，并获得最佳的阶段和网格切分方案。

运行时编排。在得到层级化计划和指定的流水线并行调度后，每个阶段首先在其所在的网格上被编译成一个并行可执行文件。然后调用一个运行时编排遍（runtime orchestration pass），以满足两个相邻阶段之间的通信需求，这些阶段需要在它们所在的两个网格之间进行通信。运行时编排遍随后根据流水线并行调度为每个网格生成静态指令，并在所有网格上调用执行。

API设计。Alpa有一个简单的API，如图4所示。Alpa要求开发者使用Python装饰器 @parallelize 来标注需要并行的函数，例如 train_step()。在第一次调用 train_step() 时，Alpa会追踪整个函数以获取模型IR，调用编译过程，并将该函数转换为并行版本。由于算子间遍依赖于算子内遍，下文将首先描述算子内遍，然后是算子间遍，最后是运行时编排遍。

A2 方法细节

4. 算子内（Intra-Operator）并行

优化目标与方法。Alpa在一个设备网格内优化算子内并行计划。Alpa采用SPMD（单程序多数据）风格的算子内并行【31, Gshard: Scaling giant models with conditional computation and automatic sharding, 2020, arXiv】, 【57, Gspmd: General and scalable parallelization for ml computation graphs, 2021, arXiv】，该风格在设备间均匀地划分算子，并在所有设备上执行相同的指令，这是基于一个网格内的设备具有同等计算能力的事实。这种SPMD风格显著减少了算子内并行计划的空间；同时，它方便地表达和统一了许多重要的方法，如数据并行、ZeRO、Megatron-LM的算子并行及其组合，这些是现有自动算子并行系统（如Tofu【55, Supporting very large models using automatic dataflow graph partitioning, 2019, Proceedings of the Fourteenth EuroSys Conference 2019】和FlexFlow【25, Beyond data and model parallelism for deep neural networks, 2018, arXiv】）未完全覆盖的。与执行随机搜索【25, Beyond data and model parallelism for deep neural networks, 2018, arXiv】或假设线性图【55, Supporting very large models using automatic dataflow graph partitioning, 2019, Proceedings of the Fourteenth EuroSys Conference 2019】的系统不同，Alpa将问题形式化为整数线性规划（ILP），并证明对于具有数万个算子的计算图，该问题可以被高效解决。

4.1 算子内并行空间

并行算法的选择。对于计算图中的一个算子，有多种可能的并行算法可以在设备网格上运行它。例如，一个矩阵乘法 $C_{ij} = \sum_{k} A_{ik}B_{kj}$ 对应一个三层for循环。为了并行化它，我们可以在设备间并行化循环i、循环j、循环k或它们的组合，这将产生不同的计算和通信成本，需要输入张量有不同的布局（layout），并导致输出张量有不同的布局。如果一个输入张量不满足布局要求，就需要进行布局转换，这会引入额外的通信成本。算子内遍的目标是为每个算子选择一个并行算法，以最小化整个图的执行时间。

• 设备网格（Device mesh）：设备网格是一组物理设备的二维逻辑视图。网格中的每个设备具有相同的计算能力。设备可以沿第一维和第二维以不同的带宽进行通信。我们假设沿同一维度的不同设备组具有相同的通信性能。对于一组物理设备，可以有多种逻辑视图。例如，给定2个节点，每节点8个GPU（总共16个设备），我们可以将它们视为2×8、1×16、4×4、8×2或16×1的设备网格。物理设备与逻辑设备网格视图之间的映射由算子间遍优化（§5）。在本节的其余部分，我们考虑一个固定的设备网格视图。

• 分片规范（Sharding Spec）：我们使用分片规范来定义张量的布局。对于一个N维张量，其分片规范定义为 $X_0X_1 \cdots X_{n-1}$，其中 $X_i \in \{S, R\}$。如果 $X_i = S$，表示张量的第i个轴被分区。否则，该轴被复制。例如，对于一个二维张量（即矩阵），SR表示它是行分区的，RS表示它是列分区的，SS表示它是行列都分区的，RR表示它是完全复制的。定义了哪些张量轴被分区后，我们还必须将分区后的张量轴映射到网格轴上。我们只考虑二维设备网格，因此一个分区后的张量轴可以映射到设备网格的第一轴、第二轴或两者。我们用上标来表示设备分配，例如，$S^0$ 表示分区沿网格的第0轴，$S^{01}$ 表示分区同时沿两个网格轴进行。$S^0R$ 意味着张量是行分区的，第一部分在设备0和1上复制，第二部分在设备2和3上复制。表1显示了一个二维张量在2×2网格（4个设备）上的所有可能的分片规范。

• 重分片（Resharding）：当一个算子的输入张量不满足为该算子选择的并行算法的分片规范时，需要进行布局转换，即重分片，这可能需要跨设备通信。表2列出了几种重分片的情况。例如，要将一个完全复制的张量转换为任何其他分片规范（情况#1），我们可以在本地对张量进行切片而无需通信；要交换分区轴（情况#4），我们执行一次all-to-all操作。

• 算子的并行算法：有了以上定义，考虑在一个二维网格上并行化一个批处理矩阵乘法 $C_{b,i,j} = \sum_{k} A_{b,i,k}B_{b,k,j}$。表3列出了几种算子内并行算法。算法#1将循环i映射到第0个网格轴，循环j映射到第1个网格轴，导致输出张量C的分片规范为$RS^0S^1$。由于左操作数$A_{b,i,k}$和右操作数$B_{b,k,j}$都只有一个并行化索引，它们的分片规范分别为$RS^0R$和$RRS^1$。在这个算法中，每个设备本地存储了计算其输出分片所需的所有输入分片，因此没有通信成本。在表3的算法#2中，当归约循环k被并行化时，需要all-reduce通信来聚合部分和。类似地，我们可以推导出批处理矩阵乘法其他并行算法的分片规范和通信成本。对于其他基本算子，如卷积和归约，我们可以通过对其数学表达式进行类似分析，得到一系列可能的并行算法。在算子内遍中，模型图以XLA的HLO格式【51, Xla: Optimizing compiler for machine learning, 2017, https://www.tensorflow.org/xla】表示，它将常见的DL算子总结为不到80个基本算子，因此我们可以为每个基本算子手动枚举可能的并行算法。

4.2 ILP 公式化

成本最小化问题。计算图 $G = (V, E)$ 的总执行成本是所有节点 $v \in V$ 上的计算和通信成本以及所有边 $e \in E$ 上的重分片成本之和。我们将成本最小化问题形式化为一个ILP，并使用现成的求解器【18, Cbc user guide, 2005, In Emerging theory, methods, and applications】来最优地解决它。

ILP公式定义。对于节点v，可能的并行算法数量为$k_v$。它有一个长度为$k_v$的通信成本向量$c_v$，即 $c_v \in R^{k_v}$，其中$c_{vi}$是第i个算法的通信成本。类似地，节点v有一个计算成本向量 $d_v \in R^{k_v}$。对于每个节点v，我们定义一个one-hot决策向量 $s_v \in \{0, 1\}^{k_v}$ 来表示它使用的算法，$s_{vi} = 1$ 表示我们为节点v选择第i个算法。对于节点v和节点u之间的重分片成本，我们定义一个重分片成本矩阵 $R_{vu} \in R^{k_v \times k_u}$，其中$R_{vuij}$是从节点v的第i个策略的输出到节点u的第j个策略的输入的重分片成本。问题的目标是：

其中第一项是节点v的计算和通信成本，第二项是边(v, u)的重分片成本。在这个公式中，s是变量，其余是常数值。公式1中的项$s_v^T R_{vu} s_u$是二次的，不能直接输入ILP求解器。我们通过引入一个新的决策向量 $e_{vu} \in \{0,1\}^{k_v \cdot k_u}$ 来线性化【19, Computational comparison of exact solution methods for 0-1 quadratic programs: Recommendations for practitioners, 2020, Journal of Applied Mathematics】这个二次项，该向量表示节点v和u之间的重分片决策。

成本估计与图简化。虽然我们可以使用性能分析（profiling）来获得$c_v, d_v, R_{vu}$的精确成本，但为简单起见，我们使用以下方法来估计它们。对于通信成本$d_v$和$R_{vu}$，我们计算通信的字节数，并将其除以网格维度的带宽来得到成本。对于计算成本$c_v$，我们遵循【55, Supporting very large models using automatic dataflow graph partitioning, 2019, Proceedings of the Fourteenth EuroSys Conference 2019】中的相同动机，将它们全部设置为零。这是合理的，因为：（1）对于像矩阵乘法这样的重型算子，我们不允许重复计算。所有并行算法总是将工作平均分配给所有设备，因此一个算子的所有并行算法具有相同的算术复杂度；（2）对于像元素级（element-wise）算子这样的轻量级算子，我们允许它们的重复计算，但它们的计算成本可以忽略不计。

图简化与后处理优化。为了简化图，我们将计算上不重要的算子（如元素级算子、转置和归约）合并到它们的一个操作数中，并从该操作数传播分片规范。这大大减少了图中的节点数量，从而减小了ILP问题的规模。我们进行一次广度优先搜索并计算每个节点的深度，节点被合并到最深的操作数。一旦ILP决定了并行计划，我们还会应用一组后ILP通信优化，例如在适用时用reduce-scatter和all-gather替换all-reduce，因为后者减少了复制张量的数量和相应的计算，同时保持通信量不变。这实现了权重更新分片【56, Automatic cross-replica sharding of weight update in data-parallel training, 2020, arXiv】或ZeRO优化器【44, Zero: Memory optimizations toward training trillion parameter models, 2020, SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】的效果。

5. 算子间（Inter-Operator）并行

优化目标。在本节中，我们开发了将模型和设备集群切分为阶段-网格对的方法。我们的优化目标是最小化整个计算图的端到端流水线执行延迟。以前的工作【17, Dapple: A pipelined data parallel approach for training large models, 2021, Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming】, 【33, Terapipe: Token-level pipeline parallelism for training large-scale language models, 2021, arXiv】考虑了简化的问题，例如假设每个阶段的设备是预先分配的，并且所有阶段都有固定的数据或算子并行计划。Alpa通过联合考虑设备网格分配和每个阶段上存在的可变算子内并行计划来摆脱这些假设。

5.1 算子间并行空间

流水线延迟公式。假设计算图包含一系列遵循图拓扑顺序的算子，记为 $o_1, \ldots, o_K$，其中算子 $o_k$ 的输入来自算子 $o_1, \ldots, o_{k-1}$。我们将算子切分为 $S$ 个阶段 $s_1, \ldots, s_S$，其中每个阶段 $s_i$ 由算子 $(o_{l_i}, \ldots, o_{r_i})$ 组成，我们将每个阶段 $s_i$ 分配给一个大小为 $n_i \times m_i$ 的子网格，该子网格从一个包含设备、形状为 $N \times M$ 的集群网格中切分出来。令 $t_i = t_{intra}(s_i, \text{Mesh}(n_i, m_i))$ 为在 $n_i \times m_i$ 的子网格上执行阶段 $s_i$ 的延迟，该延迟由ILP最小化并由算子内遍（§4）报告。如图5所示，假设我们有 $B$ 个不同的输入微批次用于流水线，整个计算图的总最小延迟可以写为：

公式解释与约束。总延迟包含两项：第一项是所有阶段的总延迟，解释为第一个微批次通过流水线的延迟；第二项是其余 $B-1$ 个微批次的流水线执行时间，该时间受最慢阶段的限制（图5中的阶段3）。我们的目标是求解公式2，并满足两个额外约束：（1）对于图的前向传播中的一个算子，我们希望将其与其对应的后向传播算子放置在同一个子网格上。由于后向传播通常使用与前向传播相似的张量集合，这有效地减少了为后向传播获取前向传播中生成的所需张量的通信量。我们使用前向和后向延迟之和作为 $t_{intra}$，因此公式2反映了包括前向和后向传播在内的总延迟。（2）我们需要切分出的子网格 $(n_1, m_1), \ldots, (n_S, m_S)$ 完全覆盖 $N \times M$ 的集群网格——我们不浪费任何计算设备资源。

5.2 DP 公式化

子网格形状的简化。为了确保所有子网格 $(n_1, m_1), \ldots, (n_S, m_S)$ 完全覆盖 $N \times M$ 的集群网格，我们将可用的子网格形状减少为两种选择：（1）大小为 $(1, 1), (1, 2), (1, 4) \ldots (1, 2^m)$ 的一维子网格；（2）大小为 $(2, M), (3, M), \ldots, (N, M)$ 且完全使用集群网格第二维的二维子网格（即，在GPU集群上，这意味着使用每台物理机器中的所有计算设备）。附录A中的一个定理证明了这些子网格形状总能完全覆盖集群网格。为了将DP算法找到的子网格分配给集群中的物理设备，我们通过首先为较大的子网格分配设备，然后为较小的子网格分配设备来进行枚举。当有多个流水线阶段具有相同的子网格形状时，我们倾向于将相邻的流水线阶段放置在设备网格上更近的位置，以减少通信延迟。这种对子网格形状的简化适用于大多数可用的云深度学习设置，并且观察到被排除的形状（$n > 1$ 且 $m < M$）通常导致较差的结果。

DP算法设计。为了在公式2中找到 $T^*$，我们开发了一个DP算法。DP首先枚举第二项 $t_{max} = \max_{1 \le j \le S} t_j$，并为每个不同的 $t_{max}$ 最小化第一项 $t_{total}(t_{max}) = \sum_{1 \le i \le S} t_i$。具体来说，我们使用函数 $F(s, k, d; t_{max})$ 来表示将算子 $o_k$ 到 $o_K$ 切分为 $s$ 个阶段并将它们放置在 $d$ 个设备上，且每个阶段的延迟小于 $t_{max}$ 时的最小总延迟。我们从 $F(0, K + 1, 0; t_{max}) = 0$ 开始，并推导出 $F$ 的最优子结构为：

并推导出最优总延迟为：

获取阶段成本 $t_{intra}$。$t_{intra}((o_k, \ldots, o_i), \text{Mesh}(n_s, m_s), s)$ 的值由算子内遍确定。它是将在子图 $(o_k, \ldots, o_i)$ 在网格 $\text{Mesh}(n_s, m_s)$ 上执行的最低延迟，其中有s个后续阶段。我们枚举所有可能的逻辑设备网格形状 $(n_l, m_l)$，满足 $n_l \cdot m_l = n_s \cdot m_s$。对于每种选择，我们用子图、逻辑网格和其他算子内选项作为输入查询算子内遍，得到一个算子内计划。然后我们用此计划和所有其他底层编译器优化（如融合、内存规划）来编译子图，得到一个可执行文件以进行精确的性能分析。分析得到阶段延迟 $(t_l)$、运行阶段所需的设备内存 $(mem_{stage})$ 和存储中间激活值所需的内存 $(mem_{act})$。我们根据选择的流水线执行调度检查所需内存是否适合设备内存 $(mem_{device})$。例如，对于1F1B调度【17, Dapple: A pipelined data parallel approach for training large models, 2021, Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming】, 【39, Memory-efficient pipelineparallel dnn training, 2021, International Conference on Machine Learning】，我们检查：

我们选择最小化 $t_l$ 并适合设备内存的逻辑网格形状。如果都不能满足，我们设置 $t_{intra} = \infty$。

与TeraPipe的比较与复杂度。我们的算法建立在TeraPipe【33, Terapipe: Token-level pipeline parallelism for training large-scale language models, 2021, arXiv】之上。然而，TeraPipe假设所有流水线阶段都相同，而Alpa旨在将计算图的算子分组为不同的流水线阶段。此外，Alpa在DP算法中为每个流水线阶段优化了网格形状。对于固定的 $t_{max}$，我们的DP算法的计算时间复杂度为 $O(K^3NM(N + \log(M)))$。$t_{max}$ 最多有 $O(K^2(N + \log(M)))$ 种选择，因此DP算法的总复杂度为 $O(K^5NM(N + \log(M))^2)$。

性能优化 #1：提前剪枝。我们使用一种类似于TeraPipe【33, Terapipe: Token-level pipeline parallelism for training large-scale language models, 2021, arXiv】的优化。我们从小到大枚举 $t_{max}$。当 $B \cdot t_{max}$ 大于当前最优的 $T^*$ 时，我们立即停止枚举。此外，在枚举 $t_{max}$ 时，我们只评估比上一个 $t_{max}$ 大足够多（至少 $\epsilon$）的 $t_{max}$ 选择。这使得DP算法找到的解与全局最优解之间的差距最多为 $B \cdot \epsilon$。我们经验性地选择 $\epsilon = 10^{-6}$s，并发现我们的算法输出的解与真实最优解（$\epsilon=0$）在所有评估设置中都是相同的。

性能优化 #2：算子聚类。我们开发了另一个DP算法【4, Distributed balanced partitioning via linear embedding, 2016, Proceedings of the Ninth ACM International Conference on Web Search and Data Mining】来聚类相邻的算子，以减少图的总大小。我们将算子 $(o_1, \ldots, o_K)$ 聚类成一系列层 $(l_1, \ldots, l_L)$，其中 $L \ll K$。该算法的目标是合并两类算子：（1）计算量不大但拉长计算图的算子；（2）如果放在不同设备网格上可能导致大量通信的相邻算子。我们定义函数 $G(k,r)$ 为将算子 $(o_1, \ldots, o_k)$ 聚类为 $r$ 个层时，单个层接收的最大数据量的最小值。$G$ 具有以下最优子结构：

其中 $C(i, k)$ 表示 $(o_i, \ldots, o_k)$ 从 $(o_1, \ldots, o_{i-1})$ 接收的输入的总大小。我们确保每个聚类层的FLOP在每层平均FLOP的 $1 + \delta$ 倍以内，同时最小化通信。对于具有相同通信成本的解决方案，我们通过最小化每层FLOP的方差来选择结构最均匀的方案。通过我们的DP算法，我们可以在 $O(K^2L)$ 时间内计算出最佳的层聚类。在实践中，我们根据设备数量和图中重型算子的数量选择一个小的 $L$。算法1总结了算子间遍的工作流程及其与算子内遍（§4）的交互。

算法1：算子间并行遍
1: 输入：模型图 G 和形状为 (N, M) 的集群 C。
2: 输出：最小流水线执行延迟 T*。
3: // 预处理图。
4: (o1, . . . , oK) ← Flatten(G)
5: (l1, . . . , lL) ← OperatorClustering(o1, . . . , oK)
6: // 运行算子内遍以获取不同阶段-网格对的成本。
7: submesh_shapes ← {(1, 1), (1, 2), ..., (1, M)} ∪ {(2, M), ..., (N, M)}
8: for 1 ≤ i ≤ j ≤ L do
9:   stage ← (li, . . . , lj)
10:  for (n, m) ∈ submesh_shapes do
11:    for s from 1 to L do
12:      t_intra(stage, Mesh(n, m), s) ← ∞
13:    end for
14:    for (nl, ml), opt ∈ LogicalMeshShapeAndIntraOpOptions(n, m) do
15:      plan ← IntraOpPass(stage, Mesh(nl, ml), opt)
16:      tl, memstage, memact ← Profile(plan)
17:      for s 满足公式 5 do
18:        if tl < t_intra(stage, Mesh(n, m), s) then
19:          t_intra(stage, Mesh(n, m), s) ← tl
20:        end if
21:      end for
22:    end for
23:  end for
24: end for
25: // 运行算子间动态规划
26: T* ← ∞
27: for tmax ∈ SortedAndFilter(t_intra, ε) do
28:   if B · tmax ≥ T* then
29:     break
30:   end if
31:   F(0, L + 1, 0;tmax) ← 0
32:   for s from 1 to L do
33:     for l from L down to 1 do
34:       for d from 1 to N · M do
35:         根据公式 3 计算 F(s, l, d;tmax)
36:       end for
37:     end for
38:   end for
39:   T*(tmax) ← mins{F(s, 0, N · M;tmax)} + (B - 1) · tmax
40:   if T*(tmax) < T* then
41:     T* ← T*(tmax)
42:   end if
43: end for

6. 并行编排

编译与通信插入。在阶段、设备网格及其分配确定后，在算子内层级，Alpa针对每个阶段及其分配的设备网格进行编译，遵循ILP求解器输出的算子内并行计划。编译依赖于XLA【51, Xla: Optimizing compiler for machine learning, 2017, https://www.tensorflow.org/xla】和GSPMD【57, Gspmd: General and scalable parallelization for ml computation graphs, 2021, arXiv】，并为每个阶段-网格对生成并行可执行文件。需要时，编译会自动插入集体通信原语（见§4），以处理由算子内并行引起的网格内通信。

跨网格重分片。在算子间层级，Alpa实现了一个额外的并行编排遍，以处理阶段之间的跨网格通信，并为算子间并行执行生成静态指令。现有的手动系统，如Megatron-LM【45, Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020, Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining】, 【49, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019, arXiv】，限制所有流水线阶段具有相同的数据和张量模型并行度，因此流水线阶段之间的通信可以通过两个等效设备网格的对应设备之间的P2P send/recv简单实现（图6a）。在Alpa中，承载两个相邻阶段的设备网格可能具有不同的网格形状，并且在两个阶段之间通信的张量可能具有不同的分片规范（图6b和图6c）。我们将这种通信模式称为跨网格重分片，这是一个多对多的多播问题。

本地all-gather优化。给定发送方和接收方网格上张量的分片规范，Alpa通过两次迭代生成通信计划来处理跨网格分片。在第一次迭代中，Alpa计算源网格和目标网格上张量分区（即瓦片）之间的对应关系，并据此生成源设备和目标设备之间的P2P send/recv原语以完成通信。然后，在第二次迭代中，它识别出目标张量在其分片规范中存在复制的机会。在这种情况下，张量只需在两个网格之间传输一次，然后通过在目标网格上的设备之间使用其更高带宽进行all-gather交换（图6c）——它将第一次迭代中生成的send/recv重写为all-gather以避免重复通信。我们称这种方法为本地all-gather跨网格重分片。由于根据我们的设计，阶段之间的通信通常很小，我们的实验表明它的性能令人满意（§8.5）。

生成流水线执行指令。作为最后一步，Alpa生成静态执行指令以在集群上启动训练。由于每个阶段具有不同的算子集合，并且可能位于不同形状的网格上，与许多SPMD流水线并行训练系统【40, Efficient large-scale language model training on gpu clusters using megatron-lm, 2021, Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】, 【57, Gspmd: General and scalable parallelization for ml computation graphs, 2021, arXiv】相比，Alpa采用MPMD风格的运行时来编排算子间并行执行——Alpa为每个设备网格生成不同的静态执行指令。Alpa开发了一套用于算子间并行执行的指令，包括分配和释放阶段中张量内存、根据跨网格重分片计划在阶段间通信张量、同步和计算等指令。根据用户选择的流水线调度，Alpa使用一个驱动进程预先生成指令，并在执行前将整个指令列表分派给每个工作节点，从而避免了运行时的驱动-工作节点协调开销。

A7 补充细节 (局限性与讨论)

Alpa的优势。与手动组合数据、算子和流水线并行（如3D并行【45, Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020, Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining】和PTD-P【40, Efficient large-scale language model training on gpu clusters using megatron-lm, 2021, Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】）的现有工作相比，Alpa的算子间和算子内并行的层级化视图通过三个主要的灵活性显著提升了它们：（1）流水线阶段可以包含不等数量的算子或层；（2）Alpa中的流水线阶段可能被映射到不同形状的设备网格上；（3）在每个阶段内，数据和算子并行配置是根据每个算子非均匀地定制的。这些共同使得Alpa能够统一所有现有的模型并行方法，并泛化到具有更多异构性的模型架构和集群设置。

当前局限性。尽管有这些优势，Alpa的优化算法目前存在一些局限性：
* Alpa不建模不同阶段之间的通信成本，因为跨阶段通信成本本质上很小。实际上，在DP或ILP中建模该成本是可能的，但这将需要枚举指数级更多的算子内遍和DP状态。
* 算子间遍目前有一个超参数：微批次的数量B，我们当前的公式没有对其进行优化，但可以通过枚举搜索。
* 算子间遍使用静态线性调度来建模流水线并行，没有考虑更动态的调度，例如，在不同设备上并行化计算图中的不同分支。
* Alpa不优化计算和通信重叠的最佳方案；Alpa只能处理所有张量形状在编译时已知的静态计算图。

结论。尽管存在这些局限性，我们在弱扩展性（§8）上的结果表明，Alpa能够为许多著名的模型生成接近最优的执行计划。

A4 实验环境

• 模型与数据集：

  • GPT-3【10, Language models are few-shot learners, 2020, arXiv】: 齐次Transformer架构的语言模型，模型并行计划已被广泛研究。参数量从2.6B到39B不等。
  • GShard Mixture-of-Experts (MoE)【31, Gshard: Scaling giant models with conditional computation and automatic sharding, 2020, arXiv】: 混合了密集和稀疏层的语言模型，具有异构架构。专家数量从64到512不等。
  • Wide-ResNet【59, Wide residual networks, 2016, arXiv】: ResNet的变体，具有更大的通道数，架构与Transformer模型显著不同，尚无手动设计的并行策略。参数量从4.4B到60B不等。
  • 具体模型参数见附录B。

• 硬件配置：

  • 集群: 8个Amazon EC2 p3.16xlarge实例，共64个GPU。
  • 节点配置: 每个实例配备8个NVIDIA V100 16GB GPU，64个vCPU，488 GB内存。
  • 连接: 节点内的8个GPU通过NVLink连接；8个节点在同一个placement group中，节点间带宽为25Gbps。

• 软件配置：

  • 实现: 约16K行Python代码和6K行C++代码。
  • 前端/后端: Jax作为前端，XLA作为后端。
  • 编译器遍: 在Jax的Jaxpr和XLA的HLO上实现。
  • 分布式运行时: 使用Ray【37, Ray: A distributed framework for emerging ai applications, 2018, 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】的actor实现设备网格工作节点，XLA运行时执行计算。
  • 通信库: NCCL【41, The nvidia collective communication library, 2018, https://developer.nvidia.com/nccl】。

A4 实验结果

我们通过弱扩展性实验评估Alpa，即随着GPU数量的增加，模型规模也相应增大，并使用整个集群的总PFLOPS作为性能指标。

8.1 端到端性能
* GPT-3 (图7a): Alpa自动生成的计划性能与手动精调的Megatron-LM【40, Efficient large-scale language model training on gpu clusters using megatron-lm, 2021, Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】相当，甚至在某些设置下略有优势。这得益于Alpa的计划与Megatron-LM的最佳计划相似，同时还额外支持了权重更新分片。仅使用算子内并行的方案（“Intra-op only”）在跨节点时性能不佳，而仅使用算子间并行的方案（“Inter-op only”）表现出乎意料地好，可扩展至64个GPU。
* GShard MoE (图7b): 与DeepSpeed【45, Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020, Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining】相比，Alpa在2节点上实现了3.5倍的加速，在4节点上实现了9.7倍的加速。DeepSpeed由于缺乏算子间并行支持，无法有效扩展到多节点。Alpa则能自动发现结合了类似专家并行（算子内）和流水线并行（算子间）的最优执行计划，实现了线性扩展。
* Wide-ResNet (图7c): 对于这种没有现成并行策略的异构模型，Alpa依然实现了良好的可扩展性，在32个GPU上达到80%的线性扩展效率。而仅支持数据和流水线并行的基线（“PP-DP”）以及仅算子间并行的方案（“Inter-op only”）会因无法分区权重而导致内存不足。这证明了Alpa的通用性。

8.2 算子内并行消融研究 (图8)
* 实验内容: 在单节点8 GPU上，禁用了流水线并行，比较了Alpa的ILP方案与多种基线（纯数据并行、ZeRO-2、ZeRO-3、基于规则的启发式方法）的性能。
* 实验结果: Alpa的ILP方案（“Auto-sharding”）在所有情况下都表现最佳，并保持了近线性的扩展性。因为它能找到最小化通信开销的正确分区方案，而其他方法或因内存不足（Data），或因未优化通信（ZeRO），或因启发式规则不优（Heuristic）而性能较差。

8.3 算子间并行消融研究 (图9)
* 实验内容: 比较了Alpa的DP算法与两种基于规则的切分算法（“Equal operator”：每个阶段分配相同数量的算子；“Equal layer”：每个阶段分配相同数量的层）。
* 实验结果: Alpa的DP算法总是优于“Equal operator”。对于齐次模型GPT，DP算法的解与“Equal layer”相同。但对于异构模型Wide-ResNet，DP算法能找到将不同数量的层分配给不同阶段的最优解，性能分别比“Equal operator”和“Equal layer”高2.6倍和1.6倍（在32 GPU上）。

8.4 编译时间 (图10, 表5)
* 实验结果: Alpa的编译时间随着模型和集群规模线性增长，对于最大的GPT模型（39B参数，64 GPU）耗时在数小时内，这远小于实际训练时间（数周），因此是可接受的。时间主要花费在枚举阶段-网格对并进行性能分析上。通过并行编译和使用代价模型进行分析，可以加速这一过程。

8.5 跨网格重分片 (图11)
* 实验内容: 在Wide-ResNet上评估了为不同形状网格间通信设计的“本地all-gather”优化。
* 实验结果: 启用该优化后，通过将更多通信从慢速的节点间连接转移到快速的本地连接，在32个GPU上带来了2.0倍的性能提升。

8.6 案例研究：Wide-ResNet (图12)
* 内容: 可视化了Alpa为Wide-ResNet在16个GPU上找到的并行策略。
* 结论: Alpa自动发现了一个复杂的、非直观的策略。它将模型切分为3个阶段，分别分配给4、4、8个GPU。前两个阶段由于激活张量较大而偏好数据并行，而在第三个阶段，ILP求解器找到了一种非平凡的方式来分区卷积算子。这个结果表明，对于像Wide-ResNet这样的异构模型，即使是领域专家也很难手动创建如此高效的策略。

A5 结论

本文提出了Alpa，一个用于自动化模型并行分布式训练的新架构。它建立在对机器学习并行方法的新视角之上：算子内并行和算子间并行。Alpa构建了一个层级化空间，并使用一套编译遍在每个并行层级上推导出高效的并行执行计划。Alpa在两个不同的粒度上编排分布式计算设备上的并行执行。为分布式模型并行深度学习设计高效的并行计划历来是一项劳动密集型任务，我们相信Alpa将使分布式模型并行学习大众化，并加速新兴大型深度学习模型的应用。

A6 附录

A. 子网格形状覆盖证明

定理1。对于一个子网格形状列表 $(n_1, m_1), \ldots (n_S, m_S)$，如果 $\sum_i n_i \cdot m_i = N \cdot M$ 且每个 $(n_i, m_i)$ 满足以下条件之一：（1）$n_i = 1$ 且 $m_i = 2^{p_i}$ 是2的幂；或（2）$m_i = M$，那么我们总能用这些子网格形状覆盖一个完整的 $(N, M)$ 网格，其中 $M = 2^m$。

证明思路。证明过程首先放置第二类子网格（$m_i=M$），它们能完整覆盖第二维度。然后问题归约为用形状为 $(1, 2^{p_i})$ 的子网格填充剩余的网格。接着使用对 $m$ 的数学归纳法来证明。当 $m=1$ 时，所有子网格都是 $(1,1)$，显然可以覆盖。假设对 $m=k-1$ 成立，在 $m=k$ 的情况下，形状为 $(1,1)$ 的子网格数量必须是偶数，因此可以两两配对形成 $(1,2)$ 的网格。这样所有子网格的第二维都大于1，可以将所有 $p_i$ 和 $m$ 都减1，归约为 $m=k-1$ 的情况。因此，该定理通过归纳法成立。

B. 模型规格

GPT-3模型规格。所有模型使用序列长度=1024，词汇表大小=51200。其他参数见表6。

GShard MoE模型规格。所有模型使用序列长度=1024，词汇表大小=32000。其他参数见表7。

Wide-ResNet模型规格。所有模型使用输入图像大小=(224, 224, 3)，类别数=1024。其他参数见表8。

C. 额外案例研究

图13可视化了Alpa为Wide-ResNet在4个和8个GPU上找到的并行策略。
* 在4个GPU上，Alpa仅使用算子内并行。该方案在开始的几十层沿批次轴分区，然后在最后几层切换到沿通道轴分区。
* 在8个GPU上，策略与4个GPU类似，仍然是纯算子内并行，但分区方案根据设备数量和模型层特性进行了调整。