摘要

Alpa通过生成统一数据、操作符和流水并行性的执行计划，自动化大型深度学习(DL)模型的模型并行训练。现有的模型并行训练系统要么要求用户手动创建一个并行化计划，要么要求用户从有限的模型并行配置空间中自动生成一个并行化计划，这不足以在分布式计算设备上扩展复杂的DL模型。Alpa通过将并行性看作两个层次：操作符间并行性和操作符内并行性来分配大型DL模型的训练。在此基础上，Alpa构建了一个新的大规模模型层次空间并行执行计划。Alpa设计了许多编译过程，在每个独立的并行级别上自动推导出最优的并行执行计划，并实现了一个高效的运行时来协调分布式计算设备上的两级并行执行。我们的评估表明，Alpa生成的并行化计划与手工调优的模型并行训练系统相匹配，甚至优于它们设计的模型。与专门的系统不同，Alpa还可以泛化到具有异构架构的模型和不需要手工设计计划的模型。

1介绍

深度学习(DL)方面最近的一些进展都是模型尺寸显著增加的直接结果。例如，将语言模型(如GPT-3)扩展到数千亿个参数，并在更大的数据集上进行训练，从而实现了根本上的新功能。

然而，在高性能加速器集群上训练这些非常大的模型目前需要大量的工程工作，具体到模型定义和训练集群环境。例如，训练一个大型的基于transformer的语言模型需要大量调优和仔细选择多个并行度维度。在TPU集群上训练大型的混合专家(MoE)转换器模型需要手动调优每一层的分区轴，而在分布式GPU上训练相同的模型需要新的管道方案，可以依赖于分区的选择。

更一般的情况是，高效的大规模模型训练需要在网络中各个操作符的粒度上调优数据、操作符和管道并行化方法的复杂组合。正确调整并行化策略可以在训练性能上带来数量级的提升，但这取决于强大的ML和系统专业知识。

大规模模型的自动化并行化将显著加快机器学习(ML)的研究和生产，使模型开发人员能够快速探索新的模型设计，而无需考虑潜在的系统挑战。不幸的是，它需要导航一个复杂的计划空间，该空间随着并行度维度以及模型和集群的大小呈指数级增长。例如，当所有并行技术都启用时，要确定执行计划需要回答一系列相互依赖的问题，例如创建多少个数据并行副本，沿着哪个轴划分每个操作符，如何将模型划分为管道阶段，以及如何将设备映射到最终的并行可执行文件。不同并行化方法的相互作用以及它们对模型和集群设置的强烈依赖形成了一个优化计划的组合空间。最近自动并行化模型训练的努力被限制在单一模型并行化方法的空间，或依赖于对模型和集群规范的强假设。

我们的主要观察是，我们可以将不同的并行化技术组织到一个分层空间中，并将这些并行化技术映射到计算集群的分层结构中，不同的并行化技术对通信的带宽要求不同，而典型的计算集群具有相应的结构：距离近的设备通信带宽高，距离远的设备通信带宽有限。

考虑到这一点，在本文中，我们从传统的数据和模型并行的角度出发，将ML并行方法重新分类为操作符内并行和操作符间并行。操作符内并行性沿着任何张量轴(批处理或非批处理)对ML操作符进行分区，并将分区分派到分布式设备(图1c)；另一方面，操作间的并行性，将模型分割成不相交的阶段和在不同设备上执行阶段的管道(图1d)。它们发生在模型计算的两个不同粒度上，区别于是否划分运算符。

因此，可以通过在每个并行性类别中指定计划来分层表示并行执行计划，这将带来许多优势。首先，操作符内部和操作符之间的并行性具有明显的特点：操作符内部的并行性具有更好的设备利用率，但在每次分割和合并操作符、每次训练迭代时都要进行通信；而操作符间的并行性只在相邻的阶段之间通信，如果切片得当，可以很轻松，但由于调度约束，会导致设备空闲时间。我们可以利用计算集群中通信带宽的不对称特性，将操作员内部的并行性映射到具有高通信带宽的设备，同时协调具有相对较低带宽的远程设备之间的操作员之间的并行性。其次，这种分层设计允许我们将每一层作为单个可处理的子问题来最优地解决。虽然联合执行计划不能保证全局最优，但它们在训练各种大型模型时表现出了很强的经验性能。

在这个新问题公式的指导下，我们设计并实现了Alpa，Alpa是第一个自动生成并行执行计划的编译器，涵盖所有数据、运算符和管道并行性。给定模型描述和集群配置，Alpa通过将集群划分为多个设备网格来实现这一点，每个设备网格包含具有最好的高带宽连接的设备，并将模型的计算图划分为多个阶段。它为设备网格分配阶段，并自动协调设备网格上操作符内部的并行性，以及设备网格间操作符之间的并行性。

总之，我们作出了以下贡献：

• 我们构建了一个两级并行执行计划空间(图1e)，其中计划使用操作符间和操作符内的并行度分层指定。
• 我们设计了可跟踪的优化算法来导出每个级别的最优执行计划。
• 我们在GPU集群上实现了分布式DL编译系统Alpa。Alpa的特点是：(1)一组使用层次优化算法生成执行计划的编译过程；(2)一个新的运行时体系结构，它协调了各阶段和设备网格之间的互操作并行性；(3)大量的系统优化，用于改进编译和解决跨网格通信。
• 我们在训练具有数十亿个参数的大型模型时评估Alpa。我们将Alpa与最先进的分布式训练系统进行比较，该系统位于一个包含8个p3.16xlarge实例和64个gpu的AWS集群上。在GPT模型上，Alpa可以匹配专门的系统Megatron-LM。在GShard MoE模型上，与手调系统Deepspeed相比，Alpa在2个节点上实现3.5倍加速，在4个节点上实现9.7倍加速，与专门的系统不同，Alpa也推广到没有手动策略的模型，并在4节点的Wide-ResNet上实现了80%的线性缩放效率，这意味着开发人员可以使用Alpa开箱即用地对大型DL模型进行高效的模型并行执行。

图1：计算图(a)中并行化计划的生成。不同的颜色代表不同的设备，虚线框代表阶段。(b)手动创建计划。(c)和(d)自动生成计划，只使用一个操作内部和操作之间的并行。(e)展示了我们的方法，该方法创建了一个层次空间来组合操作符内部和内部的并行性。

2 背景：分布式深度学习

图2:用于训练2层多层感知器(MLP)的常见并行化技术。只显示向前传递。“x”是输入数据。“w1”和“w2”是两个权重矩阵。

DL计算通常由流行的ML框架表示为数据流图。图中的边表示多维张量；节点是计算运算符，例如矩阵乘法(matmul)，它将输入张量转换为输出张量。训练一次迭代的DL模型包括:通过图转发一批数据计算损失，通过反向反向传递获得更新，并通过权值更新操作将更新应用到参数。在实践中，模型开发人员定义数据流图，执行引擎优化并在计算设备上执行它。

当模型或数据非常大，单个设备无法在合理的时间内完成训练时，我们就求助于ML并行化方法，在分布式设备上并行化计算。

2.1 机器学习并行的传统观点

现有的ML并行方法通常分为数据并行、运算符并行和管道并行。

数据并行性。在数据并行化中，训练数据跨分布式工作者进行分区，但模型是复制的。每个worker在其独立的数据分割上计算参数更新，并在权重更新之前与其他worker同步更新，以便所有workers在整个训练过程中观察到一致的模型参数。

运算符并行性。当模型太大而无法容纳在一个设备中时，算子并行是一种有效的模型并行选择。运算符并行性是指将某个运算符(下文简称为op)的计算沿非批处理轴进行划分的方法，如图2b所示的matmul，并在多个设备上并行计算该运算符的每个部分。

因为输入张量是联合分区的，所以当设备计算它的op分区时，所需的输入张量部分可能不在本地内存中。因此，从其他设备获取输入数据需要通信。当张量被均匀划分时，即SPMD，所有设备遵循相同的集体通信模式，如all-reduce, all-gather, all-to-all。

流水并行性。流水线并行不是对操作进行分区，而是将模型图中不同的操作组(称为阶段)放在不同的worker上；同时，将训练批次分割为多个微批次，并在分布的工人上来回传递微批次，如图2d所示。与操作符并行不同的是，流水并行使用点对点通信在不同的工作者之间向前和向后传递中间激活。

并行性的手动组合。最近的发展表明，需要结合上述方法来扩展今天的大型DL模型。最先进的培训系统，如Megatron-LM，手动设计了一个专门的执行计划，将这些并行性结合到transformer语言模型中。通过假设模型中重复有相同的transformer层，它为管道的每一级分配了相同的层数，并对所有层采用相同的手工设计的操作员和数据并行配置。尽管需要很强的专业知识，但手动计划不能推广到不同的模型或不同的集群设置。

并行性的自动组合。每个单独并行的配置、它们的相互依赖以及它们对模型和集群设置的依赖形成了一个棘手的空间，这阻止了自动组合这些并行的琐碎实现。例如，当与运算符并行性结合在一起时，添加一个数据并行工作者将需要最优分配一组在其中执行运算符并行性的设备。当考虑流水并行性时，最优的流水方案取决于每个阶段的数据和操作符并行性选择以及为每个阶段分配的设备。基于这种传统的观点，之前对自动并行化的探索仅限于将数据并行化与最多一种模型并行化方法相结合，这将错过大量的性能机会。接下来，我们发展ML并行性的观点。

2.2操作符内部和操作符之间的并行度

与传统观点不同，在本文中，我们将现有的并行化方法重新分类为两个正交类别：操作符内并行和操作符间并行，它们的区别在于如果它们涉及到沿任意张量轴的分划算子，接下来我们用图2中的例子来介绍这两种并行。

Intra-operator并行。算子作用于多维张量。我们可以沿着某些维度划分张量，将划分后的计算结果分配给多个设备，并让它们同时执行运算符的不同部分。我们将使用该工作流的所有并行化方法定义为操作符内部并行化。

图2a-c说明了MLP上几个典型的操作内并行性实例的应用。数据并行性，根据定义属于操作内并行性——输入张量和matmuls沿着批处理维进行划分，权重张量进行复制。或者，当权值非常大时，对权值进行划分(图2b)，就会得到前面Megatron-LM所采用的算子并行性。除了向前或向后传递的操作符外，还可以沿着批维度从权值更新阶段对操作符进行分区，从而产生权值更新分片或等价的ZeRO技术，通常被理解为数据并行性的优化。

由于划分的原因，在经营者的拆分合并时需要进行集体通信。因此，操作内部并行的一个关键特征是它导致分布式设备之间的大量通信。

Inter-operator并行。我们将运算符间并行性定义为一组正交的方法，这些方法不执行运算符分区，而是分配图的不同运算符在分布式设备上执行。

图2d说明了批拆分流水线并行性作为操作符间并行性的一个例子，设备放置是操作符间并行的另一种情况，它对模型图进行分区，并在不同的设备上执行它们，但不会使用多个微批量使管道饱和。因此，管道并行通常被视为更好的替代方法，因为它可以减少设备空闲时间。管道执行可以遵循不同的调度，如Gpipe ， PipeDream和1F1B。我们在本文中采用了1F1B调度，因为它兼顾了同步一致性，并且与Gpipe相比具有相同的理论管道延迟，但峰值内存占用更低。

在操作符间并行性中，设备仅在流水阶段之间通信，通常使用成对设备之间的点对点通信。所需要的通信量可以比操作员内部并行的集体通信量小得多。无论使用什么调度，由于阶段之间的数据依赖，操作符间的并行性会导致一些设备在向前和向后计算期间处于空闲状态。

通过这种分类，这两种并行发生在DL计算的不同粒度上，并且具有不同的通信需求，这恰好与当今典型计算集群的结构相匹配。我们将利用这些属性来设计分层算法和编译传递来自动生成执行计划。

3 概述

图 3：编译器通道和运行时架构，分片阶段是使用内部操作传递生成的分片规范注释的阶段。

图4：演示Alpa API的示例，开发人员使用一行Python装饰器@parallelize来注释需要并行化的函数，程序的其余部分保持原样。

Alpa是一种编译器，它通过在两个不同粒度(操作内和操作间的并行性)分层优化计划来生成模型并行执行计划。在内部操作级别，Alpa将在给定的设备网格上执行计算图的一个阶段(即子图)的成本降到最低，该设备网格是一组设备之间可能具有高带宽(例如，单个服务器内的GPU)。根据分配的工作负载，不同的网格可能有不同数量的计算设备。在互操作级别，Alpa将如何将模型和设备集群分割为阶段和设备网格，以及如何将它们映射为阶段-网格对，从而最小化互操作并行化延迟。互操作优化依赖于知道由内部操作优化器报告的每个阶段-网格对的最佳执行成本。通过这个层次优化过程，Alpa生成由在各自层次上局部最优的内部操作和内部操作计划组成的执行计划。

为了实现这一点，Alpa实现了三个新的编译过程，如图3所示：给定一个模型描述，以Jax中间表示(IR)和集群配置的形式，互操作编译传递将IR切片为多个阶段，并将设备集群切片为多个设备网格。互操作传递使用动态规划(DP)算法将阶段分配给网格，并对每个阶段-网格对调用内部操作编译传递，以查询此分配的最佳执行成本。一旦被调用，intra-op pass通过使用整数线性规划(ILP)公式最小化其执行成本来优化在其指定网格上运行的阶段的intra-op并行执行计划，并将成本报告给inter-op pass。inter-op pass通过对每个阶段-网格对的分配重复查询intra-op pass，利用DP最小化各操作间并行执行延迟，得到阶段和网格的最佳切片方案。

给定输出的分层计划和指定的管道-并行调度，每个阶段首先被编译为其所在网格上的并行可执行文件。调用运行时编制传递来满足两个相邻阶段之间的通信需求，而这两个阶段需要在它们所在的两个网格之间进行通信。然后，运行时编制传递根据管道并行调度生成每个网格特定的静态指令，并在所有网格上调用执行。

API。Alpa有一个简单的API，如图4所示。Alpa要求开发人员使用Python装饰器@parallelize注释要并行化的函数，比如train_step()。在第一次调用train_step()时，Alpa跟踪整个函数以获得模型IR，调用编译，并将函数转换为并行版本。

由于互操作传递依赖于内部操作传递，因此我们首先描述内部操作传递，然后是互操作传递，最后是运行时编排传递。

4 intra-operator并行策略搜索

Alpa优化设备网格中的内部运算符并行性计划。Alpa采用spmd风格的内部操作并行，类似于许多现有的工作，通过注意到单个网格内的设备通常具有等效的计算能力，它平均地划分操作符，并在不同设备上执行相同的指令。通过这种简化，Alpa显著减少了操作内并行计划的空间，但方便地表达和统一了所有重要的并行方法，如数据并行、ZeRO、Megatron-LM的运算符并行及其组合，这是现有的一些自动运算符并行系统，如Tofu和FlexFlow所没有涵盖的。与执行随机搜索或假设线性图的系统不同，Alpa将问题形式化为整数线性规划(ILP)，并表明它可以对具有数万个节点的DL计算图进行最优求解。接下来，我们描述运算符内并行性的空间和我们的解决方案。

4.1 intra-operator并行的空间

Alpa 研究任意计算图的原始算子。对于计算图中的算子，有多种可能的并行算法可以在设备网格上运行它。例如，矩阵乘法C${}_{ij}$ =∑${}_k$ A${}_{ik}$ B${}_{bj}$ 对应到一个三级循环要并行化，我们可以并行化循环i，循环j，循环k，或者它们的组合设备，将有不同的计算和通信成本，需要不同的布局的输入张量，并产生不同布局的输出张量。如果一个输入张量不能满足布局要求，需要进行布局转换，这带D来了额外的通信成本。算子内传递的目标是为每个算子选择一种算法，最小化计算图的执行时间。接下来，我们正式定义设备网格和张量的布局，并讨论布局转换的成本。

设备网。设备网格是一组物理设备的2D逻辑视图,网格中的每个设备都具有相同的计算能力，装置可沿第一网格尺寸与第二网格尺寸进行通信不同的带宽。我们假设相同网格尺寸的不同分组设备具有相同的通信性能。对于一组物理设备，可以有多个逻辑视图。例如，给定2个节点和每个节点8个GPU(即总共16个设备)，我们可以将它们视为2×8、1×16、4×4、8×2或16×1设备网格。物理设备与逻辑设备之间的映射视图通过inter-op pass优化。在本节的其余部分，我们考虑一个固定设备网格视图。

Alpa使用ILP对单个stage-mesh pair求解。这里有两个假设：a. device mesh中所有的设备具有完全相同的计算能力；b. mesh一定是二维的。

对一个N维Tensor，定义其sharding spec（即切分方案）为X${}_0$ X${}_1$ ...X${}_{n{-}_1}$ ，此处X${}_i$ ∈{S,R}。其中，S和R分别表示sharding和replicate。解决了切分，还要解决切分后到device mesh上的映射，这里使用上脚标表示。S${}^0$ 表示该维度沿着mesh的第0维做parition。以一个示例去理解，对于一个2D Tensor，我们可以穷举其所有的sharding spec和映射关系。

Resharding。当推导切分时，遇到某个operator的某个Tensor要求的sharding spec和当前实际sharding不一致时，就要引入reshard做衔接，这时候有可能会引入集合通信。在实现层面上，这部分推导逻辑应该是复用的GSPMD组件。

不同Resharding带来的通信开销示例

Parallel Algorithms。因为切分需要推导，所以要对每种operator制定推导规则，如input的某种切分，会推导出output的何种切分，反之亦然，这种推导规则被称为parallel algorithm。每种operator可能会有多条推导规则，有些algorithm会引入一些通信补偿保持数学等价性，下表是batch matmul的多条parallel algorithm的示例

batch matmul的parallel algorithms

现在有了operator的parallel algorithm及每条的通信代价，我们就可以为每个operator选择一个parallel algorithm，力求整体stage-mesh的execution cost最小，为此我们借助ILP这个现成可用的数学工具。

4.2 ILP Formulation

• C${}_v$ ：Node v的Communication cost向量。每个元素c${}_{vi}$ 表示node v选择第i个algorithm时的communication cost；
• d${}_v$ ：Node v的Computation cost向量，定义与上面类似；
• R${}_{vu}$ ：这是一个Matrix，表示Resharding cost between v and u。每个元素R${}_{vuij}$ 表示从v输出的第i个Strategy到u的输入的第jStrategy的Resharding cost；
• s${}_v$ ：Binary decision vector，是ILP中的variable。代表对Node v使用了哪个algorithm。

ILP建模：

5 Inter-Operator并行策略搜索

给定一个阶段和一个设备网格，ILP求解器生成一个最优的阶段在网格上执行的内部操作并行性计划。在本节中，我们开发了模型图级别的方法，以最优地将模型和设备集群分割为阶段网格对。我们的优化目标是最小化整个计算图的端到端延迟。之前的作品都考虑过简化的问题，比如假设每个阶段的设备都是预先分配的，所有阶段都有固定的数据或操作员并行计划。我们开发的公式，共同考虑设备网格分配和存在不同的内部操作并行计划在每个阶段。

上面已经提到，ILP针对单个stage-mesh pair做优化，那么如何将整个computation graph和device mesh cluster切成一个个stage-mesh pair呢？这就是Inter-Operator pass的工作了。该pass使用DP解决划分问题，目标是优化整体端到端latency。

不过，这里还是要做一个假设：假设computation graph的拓扑按照user define的顺序排列，这就是要把排序限定在固定的linear顺序。否则，动态规划算法不是很好建模。

• S：stage数量，S${}_i$ 表示第i个stage
• O${}_i$ ：第i个operatorS
• Mesh（n${}_i$ ,m${}_i$ ）：Shape为(n${}_i\_i$ ,m${}_i\_i$ )的Mesh cluster
• T${}_i\_i$ ：S${}_i\_i$ 的执行耗时。这部分耗时是上一节中通过ILP计算出来的结果

一个micro-batch为B的pipeline端到端latency可以等价表示如下图。之所以说是等价表示，是因为下图中把forward stage和backward stage融合在一起算总时间了，他们被给予相同的device mesh。不过，从E2E的表示上看，当把backward stage拆出来单独排列，并观察整体latency，你会发现无论是使用1F1B还是GPipe朴素的排法，整体latency可以认为几乎相同，差的不是很重要。

图5：流水的总延迟示意图

然后，我们就可以建模了。对着上面的图，我们可以将total latency表示如下式，并以优化total latency为最终目标。这里有两个有意思的点：1. stage的数量并不是外部给定的，而是根据算法自行求解出来的；2.在实际做拆分决策的时候，只对forward部分做规划即可，backward部分可根据forward的拆分决策而自行拆分。

5.2 DP Formulation

我们认真研究上面列出的公式，列出了两项：

• 1.∑${}^s^s$ ${}_i\_i$${}_{=1}\_\{=1\}$ t${}_i\_i$ 表示所有切分stage的execution cost之和（由ILP算出)
• 2.Max1≤j≤jTj是所有execution cost中最大的开销

假设我们固定锁住第二项的值（比如就是t${}_3\_3$开销最大），那么在这种情况下，minT 等价于min∑${}^s^s$ ${}_{i=1}\_\{i=1\}$ t${}_i\_i$ 。到这里，我们就有一种很自然的求解思路：对于T${}_{\max }\_\{\backslash max\}$ 我们需要枚举每一种可能的值，在将这样的值作为t的上界情况下，优化第一项min∑${}^s^s$ ${}_{{\mathrm{i}}_{=1}}\_\{\backslash mathrm\{i\}\_\{=1\}\}$ t${}_i\_i$ ,得到当前最大值为t情况下，所有stage之和最小的方案。然后再从所有枚举的t中找到最小的即为所求。

现在定义新的函数F(s,k.d;tmax)：表示将ok到oK slice到s个stages上，并使用了d个devices，每个stage的latency不大于tmax时（即tmax当做条件给出的情况下$），所有stage的execution cost之和的最小值。那么我们能列出DP方程如下:

剪枝优化。从上面的DP方程，能够感受到这是一个不算小的搜索空间，所以必要的剪枝是不可或缺的。除了减掉无效探索空间外，在所有的规模参数中，似乎唯一可以减小的只有operator的数量，我们依然要做缩图工作。但这里的缩图和Intra-op环节中要复杂一些，因为我们要粗粒度的Operator间，通信量的传输。

• Early Pruning：减掉无用探索空间。因为t${}_{max}\_\{max\}$ 需要枚举，所以从小到大枚举可以提前剪枝。当B∗t${}_{\max }\_\{\backslash max\}$ 比当前的best T∗还要大时，停止枚举。
• Operator Clustering：为了缩小图规模，这里又用一次DP，目的是将(O${}_1\_1$ ，O${}_2\_2$ ,......,O${}_K\_K$ )缩为(l${}_1\_1$,l${}_2\_2$ ,......,l${}_L\_L$ )，其中K<<L。

动态方程如下：定义G(k,r)为：当把(O${}_1\_1$ ,O${}_2\_2$ ,......,O${}_k\_k$ )缩为r个layers时，整个系统中接受最大数据量的那个layer，把这个量压缩到最小的值，就是把系统中收到最大Input通信量的那个layer，往小压，让最大的小了，其他的也就小了，那么：

C(i,k)表示(O${}_i\_i$ ,...,O${}_k\_k$ )接收到的(O${}_1\_1$ ,...,O${}_{i-1}\_\{i-1\}$ )的总通信量，后面的约束项是为了让计算更加均匀。

5.3 讨论最优化

Alpa使用分级搜索来生成执行计划。Alpa在每个并行级别上找到最优的执行计划，然而，合并的执行计划并不是全局最优的：intra-operator pass只看到单个阶段，不考虑不同管道阶段之间的通信成本。由于各阶段之间的通信在设计上很小，我们在弱伸缩性上的结果表明，Alpa已经可以生成近乎最优的执行计划。

6 Parallelism Orchestration

在这阶段、设备网格和它们的分配确定之后，在操作内部，Alpa根据ILP求解器的最优操作内并行度计划输出，根据分配的设备网格对每个阶段进行编译。编译依赖于XLA和GSPMD，并为每个阶段-网格对生成并行可执行文件。当需要时，编译器会自动插入集合通信原语，以解决由内部操作并行性引起的内部通信问题。

在互操作级别，Alpa实现了一个额外的并行编制过程，以处理各阶段之间的跨网格通信，并生成用于互操作并行执行的静态指令。

图6:跨网格重新分片。红色箭头表示在慢速连接上发送/recv，绿色箭头表示快速连接上的所有集合，(a) Megatron-LM中等网格形状的散射-聚集优化，(b)不相等网格形状的朴素send/recv，(c)不均匀网格形状的广义局部全聚集优化。

Cross-mesh重新切分。现有的手动系统，如Megatron-LM ，约束所有管道级具有相同的数据度和张量模型并行度，因此管道级之间的通信通过两个等效设备网格对应设备之间的P2P send/recv实现(图6a)。在Alpa中，两个相邻阶段的设备网格可能具有不同的网格形状，两个阶段之间通信的张量可能具有不同的分片规格(图6b和图6c)。我们将这种通信模式称为跨网格重分片，这是一个多对多多组播问题。

根据张量在发送端和接收端网格上的分片规格，Alpa生成一个通信计划，在两次迭代中解决跨网格分片。在第一次迭代中，Alpa计算源网格和目标网格上张量分区(也就是tiles)之间的对应关系。在此基础上，在源设备和目的设备之间生成P2P发送/接收原语来实现通信。在这种情况下，张量只需要在两个网格之间传输一次，然后使用其更高的带宽在目标网格上的设备上通过all-gather交换(图6)——它将第一次迭代生成的send/recv重写为all-gather，以避免重复通信。

我们称这种方法为本地所有聚集的跨网格重分片。由于在我们的设计中，阶段之间的通信通常很小，我们的实验表明，它的性能令人满意。我们将最优的跨网格重分片计划的开发推迟到未来的工作。

生成管道执行指令。作为最后一步，Alpa生成静态执行指令，在集群上启动训练。由于每个阶段都有不同的操作集，可能位于不同形状的网格上，与许多SPMD管道并行训练系统相比，Alpa采用mpmd风格的运行时来协调操作间的并行执行——Alpa为每个设备网格生成不同的静态执行指令。

Alpa开发一组用于互操作并行执行的指令，包括在一个阶段中为张量分配和释放内存的指令，在跨网格重分片计划、同步和计算等之后在各阶段之间通信张量的指令。根据用户选择的管道调度，Alpa使用驱动进程提前生成指令，并在执行之前将整个指令列表分发给每个worker，避免了运行期间driver-worker协调的开销。

7 局限性与讨论

在本节中，我们将讨论parallelisms观点的优点以及算法的几个局限性。与现有的手工合并数据的工作相比，以及管道并行性，例如3D并行性和PTD-P， Alpa的内部和内部op并行的层次视图显著提高了他们的三个主要灵活性:

(1)管道各阶段可包含不均匀数指操作符或层;(2)管线阶段在Alpa可能映射到不同形状的设备网格;(3)在每个阶段，数据和运算符的并行度配置是根据不同的操作员进行非统一定制。它们共同允许Alpa统一所有现有的模型parallelism方法，并推广到模型架构和具有更多异构性的集群设置。尽管有这些优点，Alpa的优化算法目前有一些限制:

•Alpa不同模型之间的通信成本不同的阶段，因为跨阶段通信成本是天生小。事实上，在DP或ILP是可能的，但需要指数枚举更多的内部操作和DP状态。

•inter-op pass目前有一个超参数:微批B的number，它没有被我们当前的优化公式，但可以通过枚举进行搜索。

•inter-op pass使用静态参数来模拟管道并行性线性调度，不考虑更动态的调度，例如，并行化不同设备上的计算图中的不同分支。

•Alpa不优化重叠的最佳方案计算和通信；Alpa只能处理静态在编译时已知所有张量形状的计算图。

然而，我们对弱缩放的研究结果表明，Alpa能够为许多人生成近乎最佳的执行计划的模型。

8 评估

Alpa是用Python的16K LoC和c++的6K LoC实现的。Alpa使用Jax作为前端，XLA作为后端。编译器传递是在Jax和XLA的中间表示(即Jaxpr和HLO)上实现的。对于分布式运行时，我们使用Ray actor来实现设备网格worker, XLA运行时用于执行计算，NCCL用于通信。

我们在训练具有数十亿参数的大规模模型上评估Alpa，包括GPT-3、GShard混合专家(MoE)和Wide-ResNet。测试平台是一个典型的集群，包含8个节点和64个gpu。每个节点是一个AWS EC2 p3.16xlarge实例，包含8个NVIDIA V100 16gb GPU、64个VCPU和488GB内存。节点内的8个GPU通过NVLink连接。8个节点在一个放置组中启动，具有25Gbps跨节点带宽。

我们将Alpa与两种最先进的分布式系统进行比较，以在gpu上训练大规模模型。然后，我们分离不同的编译过程，并对我们的优化算法进行研究。我们还包括一个Alpa发现的执行计划的案例研究。

8.1端到端性能

模型和训练工作量。我们的目标是表3中列出的三种类型的模型，包括具有同构和异构架构的模型。GPT-3是一种基于同质变压器的LM，它是由多个变压器层叠加而成的。GShard MoE是一种稠密稀疏的混合LM，每两层使用混合专家层来替换变压器末端的MLP。Wide-ResNet是ResNet的一种变体，具有更大的信道大小。为了研究训练大型模型的能力，我们遵循常见的ML实践，随着GPU的数量缩放模型的大小，参数范围见表3。更准确地说，对于GPT-3，我们在之后增加了隐藏大小和层数以及GPU的数量，而对于MoE，我们主要增加了建议的专家数量。对于Wide-ResNet，我们增加了卷积层的信道大小和宽度因子。对于每个模型，我们采用建议的每ML批量大小来保持相同的统计行为。然后，我们为每个模型和系统配置调优最佳微批大小，使系统性能最大化。梯度在微批次中累积。

基线。对于每个模型，我们将Alpa与一个强大的基线系统进行比较。我们使用威震天lm v2 作为GPT-3的基线系统。Megatron-LM是一种最先进的系统，用于在GPU上训练基于同构变压器的LM。它结合了数据并行性、管道并行性和手工设计的操作符并行性(后面称为TMP)。这些技术的组合由三个整数参数控制，这些参数指定分配给每种技术的并行度。我们按照他们论文的指导对三个参数进行网格搜索，并报告了最佳配置的结果。威震天- lm专门针对类似GPT的型号，因此不支持表3中的其他型号。

我们使用DeepSpeed作为MoE的基线。DeepSpeed为在gpu上训练MoE提供了最先进的实现。它结合了面向MoE层的手工运算符并行性和基于零的数据并行性。这些技术的组合由几个整数参数控制，这些参数指定分配给每种技术的并行度。我们还对它们进行网格搜索，并报告最佳结果。DeepSpeed在GPT-3上的性能与威震天- lm相似甚至更差，所以我们在GPT-3上跳过它。

对于大型的Wide-ResNet，没有专门的系统或手工设计的方案。我们用Alpa构建了一个基线“PP-DP”，它的空间仅由数据并行性和管道并行性组成，模仿了PipeDream]和Dapple的并行性空间。

对于所有模型，我们还包括使用Alpa的结果，其中只有一个操作符内部和操作符之间的并行性。

评价指标。当随着GPU数量的增加而增加模型尺寸时，我们评估了系统的弱伸缩性。由于不同GPU数量的模型是不同的，我们不能使用吞吐量指标，如每秒令牌或每秒图像。我们遵循并使用整个集群的累计peta浮点每秒操作(PFLOPS)作为度量标准。我们所有的结果(包括后面部分的结果)都有一个标准差在0.5％，所以我们跳过了图中的误差条。

GPT-3结果。GPT-3的并行化计划已经被广泛研究。从图7a中我们可以看出，这种具有最佳网格搜索参数的人工方案能够使威震天- lm在GPT-3上实现超线性弱缩放。尽管如此，与威震天- lm相比，Alpa可以自动生成执行计划，甚至在一些设置上实现了略微更好的缩放。如果与只使用内部运算符并行的方法相比，我们的结果与最近的研究是一致的——“内部运算”在16 GPU上表现很差，因为即使是最好的计划也必须在跨节点连接上大量通信张量，使通信成为瓶颈。令人惊讶的是，“仅限互操作”在64个GPU上表现良好，并保持线性扩展。

我们研究了在威震天- lm上的人工计划的网格搜索参数，并与Alpa生成的计划进行了比较。它揭示了两个主要发现。首先，在威震天- lm中，最好的手动计划将TMP设为1，除非是在一些罕见的情况下，如在64 GPU上拟合39B模型，在这种情况下，仅靠管道并行无法在GPU内存中拟合模型(阶段) ；同时，只要内存允许，数据并行性就会最大化。在实践中，梯度积累(GA)被打开以实现所需的全局批处理大小(例如，在我们的设置中为1024)。遗传算法平摊了数据并行性的通信，减少了管道并行性的气泡，但遗传算法的通信随遗传算法步长线性增长，使遗传算法处于不利地位。其次，alpa生成的计划非常类似于Megatron-LM中性能最好的计划，具有(1)均匀大小的阶段，(2)在内存不紧张时沿着批处理维度进行分区，而在内存紧张时沿着非批处理维度进行分区。我们的计划与手动计划之间的一个关键区别是，当存在数据并行时，Alpa也会分区权值更新操作，这有助于略微提高Megatron-LM的性能。这是因为Alpa作为一个通用编译器系统，可以组成广泛的并行方法，而威震天- lm目前没有权重更新分片支持。

MoE的结果。DeepSpeed采用GShard为MoE模型开发的人工算子并行方案，称为专家并行，它使用一个简单的规则:为MoE层的算子划分专家轴，但对非专家层切换回数据并行。然后将这种专家并行性与ZeRO数据并行性和TMP结合起来。所有这些技术都属于操作符内并行。不幸的是，DeepSpeed的专用实现与任何操作符间并行方法不兼容，这是在节点间带宽较低的情况下跨多个节点伸缩所必需的。因此，Deepspeed只在一个节点内保持良好的性能(≤8gpu)。由于同样的原因，“Intra-op only”无法跨多个节点伸缩。“Inter-op only”在32和64 gpu上耗尽内存，因为当GPU的数量大于模型的层数时，不容易平均地对模型进行切片。不平衡的切片会导致一些内存密集型阶段耗尽内存。

相比之下，Alpa会自动发现结合了操作符内部和操作符之间并行性的最佳执行计划。对于操作符内部并行性，Alpa找到了一种类似于专家并行性的策略，并将其与ZeRO数据并行性结合起来，这要归功于它的基于ilp的内部操作通。然后，Alpa构建阶段，并使用操作符间的并行性，以在慢连接上实现较小的通信量。Alpa在16个GPU上保持线性扩展，在64个GPU上也能很好地扩展。与DeepSpeed相比，Alpa做到了这一点3．5×加速2个节点和a9.7×在4个节点上加速Wide-ResNet结果。与前面两种堆叠在同一层的模型不同，Wide-ResNet具有更异构的架构。当数据批处理跨层转发时，激活张量的大小会缩小，而权重张量的大小会膨胀。这将导致内存使用和计算强度在各层之间的不平衡分布。对于这种模型，如果不是不可能，手工设计方案是很困难的。然而，Alpa仍然在32个GPU上实现了80%的可扩展性能。基线“PP-DP”和“Inter-op only”在训练大型模型时耗尽内存，因为它们不能划分权重来减少内存使用，并且难以构建内存平衡阶段。“Intra-only”需要在慢速连接上进行大量通信，因此它不能跨多个节点扩展。

图7

8.2 Intra OP PARALLELISM消融研究

我们研究了算子内并行优化算法的有效性，我们将基于ilp的最佳解决方案与ZeRO优化器和基于规则的分区策略等备选方案进行比较。

实验设置。在模型大小方面，我们运行了一个类似的弱伸缩基准，但禁用了管道并行性和梯度积累来控制变量。该基准测试是在一个带有8个GPU的Amazon EC2实例上完成的。为了在一个节点中模拟大规模训练的执行环境，与模型配置相比，我们使用更大的隐藏尺寸、更小的批处理尺寸、更少的层数。

基线。我们比较了操作符内部并行性的自动解决方案。“数据”是普通的数据并行。“ZeRO-2”是数据并行性的内存高效版本，它对梯度和优化器状态进行了分区。" ZeRO-3 "额外在" ZeRO-2 "上分区参数。“Heuristic”使用了一个结合GSPMD中分片传播的规则。它将每个输入张量的最大维度标记为已分区，并运行分片传播，以获得图中所有节点的分片规格。“自动分片”是我们基于ILP求解器的解决方案。

结果。如图8所示，“Data”很快就会耗尽内存，无法训练大型模型。“0 -2”和“0 -3”解决了数据并行的内存问题。但它们不能优化通信，因为它们总是通信梯度。当梯度远大于活化时，它们的性能退化。“启发式”通过划分所有张量来解决内存问题，但可以通过更大的通信来减慢速度。“自动分片”在所有情况下都表现最佳，并保持了近似线性的伸缩，因为它找出了总是将通信开销最小化的正确分区计划。

图8

8.3 Inter-OP PARALLELISM消融研究

我们研究了算子间并行度优化算法的有效性。我们用“DP”表示我们的算法。

实验设置。我们报告了我们的DP算法在GPT和Wide-ResNet上的三种变体的性能。

基线。我们将我们的DP算法与两种基于规则的算法进行了比较。“相等运算符”禁用基于dp的运算符集群，但为每个集群分配相同数量的运算符。“等层”限制了我们的DP算法在所有阶段使用相同的层数。

结果。结果如图9所示。“DP”总是优于“相等算子”。这是因为“相等运算符”合并运算符应该放在不同的设备网格上。该算法可以根据通信开销和计算平衡对操作符进行聚类。“DP”能否胜过“Equal layer”取决于模型架构。在类似GPT的同质模型中，我们的DP算法的解决方案对所有阶段使用相同的层数，所以“Equal layer”执行与“DP”相同。在Wide-ResNet上，最优解可以将不同的层分配到不同的阶段，因此“等层”算法不如完全灵活的DP算法。对于32GPU上的Wide-ResNet，我们的算法优于“Equal operator”和“Equal layer”。

8.4 编译时间

图10显示了Alpa对所有GPT设置的编译时间。结果表明，Alpa可以很好地扩展到大型模型或大型集群，因为编译时间会随着模型的大小和集群中GPU的数量线性增长。表4报告了我们评估中最大的GPT模型(39B, 64 gpu)的编译时间分解。大部分的时间花在枚举阶段网格对和分析它们。对于编译部分，我们通过分布式工作器并行编译不同阶段来加速编译。对于分析，我们使用在XLA指令级构建的简单成本模型来加速分析，该模型用分段线性模型估计矩阵乘法和通信原语的成本。通过这些优化，编译和搜索模型最多需要几个小时，这是可以接受的，因为它比实际的训练时间要短得多，后者可能需要几周的时间。

8.5 Cross-mesh重新切分

我们在Wide-ResNet上对不同形状网格之间的跨网格重分片进行了广义局部全聚集优化，如图所示。“signal send/recv”是一个合成的情况，我们在两个阶段之间只发送1个信号字节，这可以看作是性能的上界。" w/o local all-gather "禁用我们的local all-gather优化，只使用send/recv。“w/ local all-gather”使我们的local all-gather优化将更多的通信从慢连接转移到快速本地连接，这在32GPU上带来了2.0的加速。

8.6 案例研究:WIDE-RESNET

如上图展示了Alpa在16个gpu上发现的Wide-ResNet并行化策略。我们还在附录c中包含了在4和8GPU上的结果可视化。在4 GPU上，Alpa只使用了内部运算符并行。操作符内解决方案沿批处理轴对前几十层进行分区，然后切换到对最后几层的通道轴进行分区。在16个GPU上，Alpa将模型分成3个阶段，分别分配4、4、8个GPU到阶段1、2、3。由于激活张量大于权重张量，因此在前两个阶段优先考虑数据并行性。在第三阶段，ILP求解器找到了一种非平凡的方法来划分卷积算子。这个结果表明，手动为像Wide-ResNet这样的异构模型创建这样的策略是不透明的，甚至对于领域专家也是如此。

9 结论

我们提出了一种用于自动化模型并行分布式训练的新架构Alpa，它建立在机器学习并行化方法的新观点之上：操作内部和操作间的并行化。Alpa构造一个层次空间，并使用一组编译传递来导出每个独立并行级别的最优并行执行计划。Alpa在两个不同的粒度上协调分布式计算设备上的并行执行。为分布式模型-并行深度学习提出一个高效的并行化计划历来是一项劳动密集型的任务，我们相信Alpa将使分布式模型-并行学习民主化，并加快新兴大型深度学习模型的采用，Alpa的源代码将公开。

原文链接：https://www.usenix.org/conference/osdi22/presentation/zheng-lianmin