Alpa：用于分布式深度学习的算子间与算子内并行自动化

介绍

本文是Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning的译文与笔记。

摘要

Alpa 通过生成统一 数据并行、算子并行 与 流水线并行 的执行计划，实现了对大规模深度学习（DL）模型的模型并行训练自动化。现有的模型并行训练系统要么要求用户手动创建并行化方案，要么只能从受限的模型并行配置空间中自动生成方案，难以在分布式计算设备上扩展复杂的 DL 模型。

Alpa 将大规模 DL 模型的训练分发视为两个层次的并行性：算子间并行 与 算子内并行。基于这一视角，Alpa 构建了一个新的分层模型并行执行计划空间，并设计了一系列编译阶段，在每个并行层次上自动推导高效的并行执行计划。Alpa 还实现了一个高效的运行时系统，用于在分布式计算设备上协同这两个层次的并行执行。

我们的评估结果表明，即使在其目标模型上，Alpa 生成的并行化方案也能够达到或超过人工调优的模型并行训练系统。与专用系统不同，Alpa 还能泛化到具有异构架构的模型以及没有人工设计并行方案的模型。Alpa 的源代码已公开发布，地址为 https://github.com/alpa-projects/alpa。

引言

近年来深度学习（DL）的多项进展，直接得益于模型规模的显著增长。例如，将语言模型（如 GPT-3）扩展到数千亿参数，并在更大规模的数据集上进行训练，使模型具备了全新的能力。然而，在分布式集群上训练这些超大规模模型，仍然需要大量针对具体模型定义和集群环境的工程工作。

例如，训练大规模基于 Transformer 的语言模型需要对多种并行维度进行大量调优，并谨慎选择其组合方式。训练大规模 Mixture-of-Experts（MoE） Transformer 模型时，在 TPU 集群上需要为每一层手动调整分区轴；而在 AWS GPU 集群上训练同一模型，则需要新的流水线方案，并且这些方案可能依赖于分区选择。

更普遍地说，高效的大规模模型训练需要在数据并行、算子并行与流水线并行之间进行复杂组合，并且粒度细化到单个张量算子层面。已有研究表明，正确调优并行化策略可以带来数量级的训练性能提升，但这通常依赖于深厚的机器学习与系统领域专业知识。

如果能够实现大规模模型并行化的自动化，将极大加速机器学习研究与生产实践，使模型开发者能够在不受底层系统挑战影响的情况下，快速探索新的模型设计。然而，这一目标需要在一个极其复杂的方案空间中进行搜索，该空间会随着并行维度、模型规模和集群规模的增加而呈指数级增长。

例如，在启用所有并行技术的情况下，确定执行计划需要回答一系列相互依赖的问题，包括：创建多少数据并行副本、每个算子沿哪个轴进行分区、如何将模型划分为流水线阶段，以及如何将设备映射到最终的并行可执行单元。不同并行化方法之间的相互作用，以及它们对模型与集群配置的高度依赖，共同形成了一个组合爆炸的优化空间。近期的一些自动模型并行化工作，要么仅局限于单一模型并行方式的空间，要么对模型和集群规格做出了较强的假设。

我们的关键观察是：可以将不同的并行化技术组织到一个分层空间中，并将这些并行化技术映射到计算集群的分层结构上。不同的并行化技术对通信带宽的需求不同，而典型的计算集群也具有相应的结构特征：物理位置接近的设备可以进行高带宽通信，而距离较远的设备则受限于较低的通信带宽。

基于上述观察，本文采用一种不同于传统数据并行与模型并行的视角，将机器学习并行化方法重新划分为算子内并行与算子间并行。算子内并行沿着一个或多个张量轴（批次轴或非批次轴）对机器学习算子进行划分，并将这些分区分发到分布式设备上（见 图 c）；而算子间并行则将模型切分为若干互不重叠的阶段，并在不同设备集合上以流水线方式执行这些阶段（见 图 d）。这两种并行方式发生在模型计算的不同粒度上，其区别在于是否对算子本身进行划分。

在此基础上，一个并行执行计划可以通过在每一类并行方式中分别指定方案，以分层结构的形式进行表达，从而带来多方面的优势。首先，算子内并行与算子间并行具有显著不同的特性：算子内并行能够实现更高的设备利用率，但在每次训练迭代中，每当分区算子发生拆分与合并时都需要进行通信；而算子间并行仅在相邻阶段之间进行通信，如果切分得当，这种通信开销可以较小，但会由于调度约束而引入设备空闲时间。我们可以利用计算集群中通信带宽的非对称特性，将算子内并行映射到高带宽互连的设备上，而在带宽相对较低、距离较远的设备之间协调算子间并行。

其次，这种分层设计使我们能够将每一层视为一个独立且可处理的子问题，从而在每个层次上近似地获得最优解。尽管整体执行计划并不保证全局最优，但在多种大规模模型训练中，该方法在经验上展现了强劲的性能。

在这一新的问题表述指导下，我们设计并实现了 Alpa，这是首个能够自动生成同时覆盖数据并行、算子并行与流水线并行的并行执行计划的编译器。给定模型描述和集群配置，Alpa 通过将集群划分为多个设备网格来实现这一目标，每个设备网格优先包含具有高带宽连接的设备；同时，它将模型的计算图划分为多个阶段，并将这些阶段分配到设备网格上，从而在设备网格内部自动协调算子内并行，在设备网格之间自动协调算子间并行。

我们的主要贡献包括：

• 构建了一个两层并行执行计划空间（见 图 e），通过算子间并行与算子内并行的分层方式来描述执行计划。
• 设计了可处理的优化算法，在每一层次上推导近似最优的执行计划。
• 实现了 Alpa，一个面向 GPU 集群的分布式深度学习编译系统。Alpa 的特性包括： （1） 一组编译阶段，利用分层优化算法生成执行计划； （2） 一种新的运行时架构，用于在设备网格之间协调算子间并行； （3） 多项系统级优化，用于提升编译效率并处理跨设备网格通信。

图 1： 针对 (a) 所示计算图生成并行化方案。不同颜色表示不同设备，虚线框表示流水线阶段。(b) 为手动创建的方案。(c) 与 (d) 分别仅使用算子内并行或算子间并行自动生成方案。(e) 展示了我们的方法，通过构建分层空间来结合算子内并行与算子间并行。

• 在包含数十亿参数的大模型训练中对 Alpa 进行了评估。我们在由 8 个 p3.16xlarge 实例（共 64 块 GPU）组成的 Amazon EC2 集群上，将 Alpa 与最先进的分布式训练系统进行了比较。在 GPT 模型上，Alpa 的性能可以与专用系统 Megatron-LM 相当；在 GShard MoE 模型上，相比于人工调优的 Deepspeed 系统，Alpa 在 2 个节点上实现了 3.5× 的加速，在 4 个节点上实现了 9.7× 的加速。不同于专用系统，Alpa 还能泛化到没有人工并行策略的模型，在 4 个节点上训练 Wide-ResNet 时实现了 80% 的线性扩展效率。这意味着开发者可以使用 Alpa 开箱即用 地获得高效的大规模深度学习模型并行执行。

背景：分布式深度学习

在常见的机器学习框架中，深度学习计算通常被表示为一个数据流图。图中的边表示多维张量，节点则是计算算子，例如将输入张量转换为输出张量的矩阵乘法（matmul）。一次深度学习模型的训练迭代包括：将一个批次的数据前向传播通过计算图以计算损失，通过反向传播计算梯度更新，并通过权重更新操作将这些更新应用到模型参数上。在实践中，模型开发者负责定义数据流图，而执行引擎则对其进行优化并在计算设备上执行。

当模型规模或数据规模过大，以至于单个设备无法在合理时间内完成训练时，我们就需要采用机器学习并行化方法，将计算并行地分布到多个设备上。

图 2： 用于训练一个 两层多层感知机（MLP） 的常见并行化技术示例。图中仅展示了前向传播过程。“x” 表示输入数据，“w1” 和 “w2” 表示两个权重矩阵。

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

现有的机器学习并行化方法通常被划分为数据并行、算子并行和流水线并行。

数据并行。 在数据并行中，训练数据被划分并分配到分布式工作节点上，而模型本身会被复制到每个节点。每个工作节点在其独立的数据分片上计算参数更新，并在权重更新之前与其他节点同步这些更新，从而确保在整个训练过程中，所有节点观察到一致的模型参数。

算子并行。 当模型过大而无法放入单个设备时，算子并行是一种有效的模型并行方案。算子并行指的是将某个特定算子（下文简称 op）的计算沿着非批次轴进行划分，例如 图 b 所示的矩阵乘法（matmul），并在多个设备上并行计算该算子的不同部分。

由于输入张量被联合划分，当某个设备计算其算子分区时，所需的输入张量片段可能并不驻留在本地内存中，因此需要通过通信从其他设备获取输入数据。当张量被均匀划分时，即 SPMD 模式，所有设备将遵循相同的集合通信模式，例如 all-reduce、all-gather 和 all-to-all。

流水线并行。 与算子划分不同，流水线并行将模型计算图中的不同算子组（称为 阶段）放置在不同的工作节点上；同时，将训练批次拆分为多个 微批次，并在分布式工作节点上对这些微批次的前向和反向传播进行流水线执行，如 图 d 所示。与算子并行不同，流水线并行在前向和反向传播过程中，通过点对点通信在不同工作节点之间传递中间激活值。

并行化方式的手动组合。 近期的发展表明，上述并行化方法需要进行组合，才能扩展当今的大规模深度学习模型。最先进的训练系统（如 Megatron-LM）为 Transformer 语言模型手动设计了结合多种并行方式的专用执行计划，这通常被称为 三维并行。该方法假设模型由重复的 Transformer 层构成，将相同数量的层分配给每个流水线阶段，并为所有层统一应用人工设计的算子并行和数据并行配置。尽管这种方法需要高度专业的经验，但其手动设计的方案难以泛化到不同模型或不同的集群配置中。

并行化方式的自动组合。 各类并行方式的配置、它们之间的相互依赖关系，以及它们对模型与集群配置的依赖，共同构成了一个难以处理的复杂空间，这使得自动组合这些并行方式并非易事。例如，当与算子并行结合时，每增加一个数据并行副本，就需要为该工作节点分配一整组设备（而非单个设备），并在这些设备内部重新确定最优的算子并行配置；当引入流水线并行时，最优的流水线方案又依赖于每个流水线阶段的数据并行和算子并行选择，以及设备在阶段之间的分配方式。在这种传统视角下，已有的自动并行化研究通常仅限于将数据并行与至多一种模型并行方式相结合，从而错失了大量潜在的性能提升空间。接下来，我们将介绍我们对机器学习并行化的视角。

算子内并行与算子间并行

不同于传统视角，本文将现有的并行化方法重新划分为两个正交的类别：算子内并行与算子间并行。它们的区别在于是否沿任意张量轴对算子进行划分。下面我们结合 图 2 中的示例来介绍这两类并行方式。

算子内并行。 算子作用于多维张量。我们可以沿着某些维度对张量进行划分，将得到的分区计算分配给多个设备，并让这些设备同时执行算子的不同部分。我们将所有遵循这一工作流程的并行化方法定义为算子内并行。

图 a–c 展示了在一个 MLP 上应用多种典型算子内并行实例的情况。按照定义，数据并行属于算子内并行：输入张量和矩阵乘法沿着批次维度进行划分，而权重张量被复制。另一方面，当权重非常大时，对权重进行划分（见 图 b）就形成了前文提到的算子并行方式，这也是 Megatron-LM 所采用的方法。除了前向或反向传播阶段的算子之外，还可以对权重更新阶段的算子进行划分，从而得到权重更新分片，或等价的 ZeRO 技术，这通常被视为数据并行的一种优化。

由于存在算子划分，在算子的拆分与合并处需要进行集合通信。因此，算子内并行的一个关键特征是：它会在分布式设备之间引入大量通信。

算子间并行。 我们将算子间并行定义为另一类正交的并行化方法，它不对算子本身进行划分，而是将计算图中的不同算子分配到分布式设备上执行。

图 d 展示了批次拆分的流水线并行，这是算子间并行的一种典型形式。流水线执行可以采用不同的调度方式，例如 Gpipe、PipeDream 以及 同步 1F1B。本文全程采用同步 1F1B 调度，因为它保持同步一致性，并且在具有相同流水线延迟的情况下，相比 Gpipe 具有更低的峰值内存占用。

在算子间并行中，设备仅在流水线阶段之间进行通信，通常使用设备对之间的点对点通信。其所需的通信量通常显著小于算子内并行中的集合通信。无论采用何种调度方式，由于阶段之间存在数据依赖关系，算子间并行都会在前向与反向计算过程中导致部分设备处于空闲状态。

通过这种分类，这两类并行方式发生在深度学习计算的不同粒度上，并且具有不同的通信需求，而这恰好与当今计算集群的结构相匹配。我们将利用这些特性来设计分层算法和编译阶段，以自动生成执行计划。已有若干并行工作也提出了类似的分类方法，但 Alpa 是首个端到端系统，能够基于这一分类从完整的并行空间中自动生成并行执行计划。

注：设备放置 也是算子间并行的一种形式，它对模型计算图进行划分并在不同设备上执行，但不会通过多个微批次来填满流水线。因此，相比之下，流水线并行通常被认为是更优的替代方案，因为它能减少设备空闲时间。

总览

Alpa 是一个通过在两个不同层次上进行分层优化来生成模型并行执行计划的编译器：算子内并行层与算子间并行层。

在算子内层面，Alpa 在给定的设备网格上，针对某一阶段（即计算图的子图），最小化其在对应算子内并行方案下的执行代价。设备网格由一组彼此之间具有较高带宽连接的设备组成（例如单台服务器内的 GPU）。不同网格可以根据所分配的工作负载，包含不同数量的计算设备。

在算子间层面，Alpa 通过决定如何切分模型与设备集群为多个阶段与设备网格，并确定阶段与网格之间的映射关系，从而最小化算子间并行的执行延迟。算子间优化依赖于算子内优化器所报告的每个阶段–网格对的执行代价。通过这一分层优化过程，Alpa 生成由算子内与算子间并行方案共同组成的执行计划，这些方案在各自层次上都是局部近似最优的。

图 3： 编译阶段与运行时架构示意图。分片阶段 是指经过算子内编译阶段生成分片规格并完成标注的阶段。

为实现上述目标，Alpa 实现了三个新的编译阶段，如 图 3 所示。给定以 Jax 中间表示（IR）形式描述的模型以及集群配置，算子间编译阶段首先将 IR 切分为多个阶段，并将设备集群切分为多个设备网格。该阶段使用动态规划（DP）算法为阶段分配设备网格，并针对每个阶段–网格对调用算子内编译阶段，以查询该分配的执行代价。

当被调用时，算子内编译阶段通过一个 整数线性规划（ILP） 模型，对运行在指定设备网格上的阶段进行算子内并行优化，以最小化其执行代价，并将结果反馈给算子间编译阶段。算子间编译阶段通过反复查询不同阶段–网格分配下的算子内执行代价，利用动态规划最小化整体的算子间并行执行延迟，并最终获得最优的阶段与网格切分方案。

在生成最终的分层执行计划并选定流水线并行调度策略后，每个阶段会首先在其对应的设备网格上被编译为一个并行可执行程序。随后，运行时编排阶段会被调用，用于满足相邻阶段之间、跨设备网格的通信需求。该运行时阶段根据流水线并行调度策略，为每个设备网格生成特定的静态指令，并在所有网格上触发执行。

Alpa 提供了一个简洁的 API，如 图 4 所示。开发者只需使用 Python 装饰器 @parallelize 标注需要并行化的函数（例如 train_step()），其余程序保持不变。

图 4： 用于展示 Alpa 在 Jax 中 API 的示例。开发者通过 Python 装饰器 @parallelize 标注需要并行化的函数，其余代码无需修改。

# 在 Jax 函数上添加 @parallelize 装饰器
@parallelize
def train_step(state, batch):
    def loss_func(params):
        out = state.forward(params, batch["x"])
        return jax.numpy.mean((out - batch["y"]) ** 2)

    grads = grad(loss_func)(state.params)
    new_state = state.apply_gradient(grads)
    return new_state

# 一个典型的训练循环
state = create_train_state()
for batch in data_loader:
    state = train_step(state, batch)

在首次调用 train_step() 时，Alpa 会对整个函数进行追踪以获取模型 IR，随后触发编译流程，并将该函数转换为并行版本。

由于算子间编译阶段依赖于算子内编译阶段，接下来的内容将首先介绍算子内编译阶段，然后是算子间编译阶段，最后介绍运行时编排阶段。

算子内并行

Alpa 在设备网格内优化算子内并行方案。Alpa 采用 SPMD 风格的算子内并行方式，即在设备之间对算子进行均匀划分，并在所有设备上执行相同的指令，这与单个设备网格内各设备计算能力等价的事实相一致。这种 SPMD 风格显著缩小了算子内并行方案的搜索空间；同时，它能够方便地表达并统一多种重要方法，例如数据并行、ZeRO、Megatron-LM 的算子并行，以及它们的组合——而这些方法并未被现有的自动算子并行系统（如 Tofu 和 FlexFlow）完全覆盖。

不同于采用随机搜索或假设计算图为线性结构的系统，Alpa 将该问题形式化为一个整数线性规划（ILP），并证明即使对于包含数万个算子的计算图，也可以高效求解。下面我们将描述算子内并行的空间以及我们的解决方案。

算子内并行空间

对于计算图中的一个算子，在设备网格上运行它时，可能存在多种并行算法选择。以矩阵乘法 $C_{ij}=\sum_k A_{ik}B_{kj}$ 为例，它对应一个三层嵌套循环。为了实现并行化，可以将循环 $i$、$j$、$k$ 中的任意一个，或它们的组合，映射到不同设备上执行。不同的并行方式会带来不同的计算与通信开销，并且对输入张量的布局有不同要求，同时生成具有不同布局的输出张量。

如果某个输入张量不满足算子所需的布局要求，就需要进行布局转换，从而引入额外的通信开销。算子内编译阶段的目标，是为每一个算子选择一种并行算法，使整个计算图的执行时间最小化。下面我们将形式化地定义设备网格与张量布局，并讨论布局转换的代价。

设备网格

设备网格是对一组物理设备的二维逻辑视图。网格中的每个设备具有相同的计算能力。设备可以沿着第一维或第二维进行通信，并且这两种方向上的通信带宽可能不同。我们假设沿着同一网格维度的不同设备组具有相同的通信性能。

对于一组物理设备，可能存在多种逻辑视图。例如，给定 2 个节点、每个节点 8 块 GPU（共 16 个设备），可以将其视为 2×8、1×16、4×4、8×2 或 16×1 的设备网格。物理设备到逻辑设备网格视图的映射由算子间编译阶段进行优化。在本节的其余部分中，我们假设设备网格视图是固定的。

分片规格

我们使用分片规格（sharding spec）来定义张量的布局。对于一个 $N$ 维张量，其分片规格定义为 $X_0X_1\cdots X_{n-1}$，其中 $X_i\in{S,R}$。若 $X_i=S$，表示该张量在第 $i$ 个轴上被划分；否则（$X_i=R$）表示该轴被复制。

例如，对于一个二维张量（即矩阵）：

• SR：按行划分
• RS：按列划分
• SS：同时按行和列划分
• RR：完全复制，不进行划分

在确定哪些张量轴被划分之后，还需要将这些被划分的张量轴映射到设备网格的轴上。由于只考虑二维设备网格，一个被划分的张量轴可以映射到设备网格的第一轴、第二轴，或者同时映射到两者。我们使用上标来标注设备分配方式。

例如：

• $S^0$ 表示分区沿着设备网格的第 0 轴
• $S^{01}$ 表示分区同时沿着两个设备网格轴
• $S^0R$ 表示张量按行划分为两部分：第一部分在设备 0 和设备 1 上复制，第二部分在设备 2 和设备 3 上复制

表 1： 在 2×2 设备网格上，一个二维张量的分片规格示例。张量 $A$ 的形状为 $(N,M)$，设备网格为 [[Device 0, Device 1], [Device 2, Device 3]]。每个设备存储 $A$ 的一个分区。第一列为分片规格名称，其余列使用 Numpy 语法描述各设备上存储的分区。

表 2： 若干重新分片（resharding）的示例。all-gather(x,i) 表示沿设备网格第 $i$ 个轴对 $x$ 字节的数据执行一次 all-gather 操作。$M$ 为张量大小，$(n_0,n_1)$ 为设备网格形状。

表 3： 批量矩阵乘法 $C_{b,i,j}=\sum_k A_{b,i,k}B_{b,k,j}$ 的若干并行算法示例。记号 all-reduce(x,i) 表示沿设备网格第 $i$ 个轴对 $x$ 字节的数据执行一次 all-reduce 操作。$M$ 为输出张量大小，$(n_0,n_1)$ 为设备网格形状。

重新分片

当某个算子的输入张量不满足其所选并行算法的分片规格时，就需要进行一种布局转换，即重新分片（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 格式表示，该格式将常见的深度学习算子归纳为不足 80 种基础算子，因此我们可以为每一种基础算子手动枚举其可能的并行算法。

ILP 形式化

计算图 $G=(V,E)$ 的总执行代价，等于所有节点 $v\in V$ 上的计算与通信代价之和，以及所有边 $e\in E$ 上的重新分片代价之和。我们将这一代价最小化问题形式化为一个整数线性规划（ILP），并使用现成的求解器进行最优求解。

对于节点 $v$，其可选的并行算法数量为 $k_v$。对应地，它具有一个长度为 $k_v$ 的通信代价向量 $c_v\in\mathbb{R}^{k_v}$，其中 $c_{v,i}$ 表示第 $i$ 个算法的通信代价；以及一个计算代价向量 $d_v\in\mathbb{R}^{k_v}$。我们为每个节点 $v$ 定义一个 one-hot 决策向量 $$ s_v\in{0,1}^{k_v} $$ 其中 $s_{v,i}=1$ 表示选择节点 $v$ 的第 $i$ 个并行算法。

对于节点 $v$ 与节点 $u$ 之间的重新分片代价，我们定义一个重新分片代价矩阵 $$ R_{vu}\in\mathbb{R}^{k_v\times k_u}, $$ 其中 $R_{vu,ij}$ 表示从节点 $v$ 的第 $i$ 种策略输出，转换为节点 $u$ 的第 $j$ 种策略输入所需的重新分片代价。

优化目标为： $$ \min_s \sum_{v\in V} s_v^{\top}(c_v+d_v) +\sum_{(v,u)\in E} s_v^{\top}R_{vu}s_u, $$ 其中第一项为节点 $v$ 的计算与通信代价，第二项为边 $(v,u)$ 的重新分片代价。在该公式中，$s$ 是优化变量，其余均为常量。

由于式中项 $s_v^{\top}R_{vu}s_u$ 是二次项，无法直接输入 ILP 求解器，因此我们通过引入新的决策向量 $$ e_{vu}\in{0,1}^{k_v\cdot k_u} $$ 来对二次项进行线性化，该向量表示节点 $v$ 与 $u$ 之间的重新分片决策。

尽管可以通过性能剖析获得 $c_v$、$d_v$ 与 $R_{vu}$ 的精确代价，本文为了简化采用如下估计方式：对于通信代价 $d_v$ 与 $R_{vu}$，通过计算通信字节数并除以设备网格对应维度的带宽来估算；对于计算代价 $c_v$，按照已有工作的方法，将其统一设为 0。这是合理的，因为：（1）对于矩阵乘法等重计算算子，不允许复制计算，所有并行算法都会将工作量均匀分配到设备上，因此算术复杂度相同；（2）对于逐元素等轻量算子，尽管允许复制计算，但其计算开销可以忽略不计。

为简化计算图，我们将计算代价可忽略的算子（如逐元素算子、转置和归约）合并到其某个操作数中，并从该操作数传播分片规格。这显著减少了图中的节点数量，从而降低了 ILP 问题规模。具体做法是对计算图执行一次广度优先搜索，计算每个节点的深度，并将节点合并到其最深的操作数中。

在通过 ILP 决定并行方案之后，我们还会应用一组 ILP 后通信优化，例如在适用的情况下，将 all-reduce 替换为 reduce-scatter 与 all-gather 的组合。这样可以在保持通信总量不变的同时，减少复制张量的数量及其对应的计算量，从而实现权重更新分片或 ZeRO 优化的效果。

算子间并行

在本节中，我们提出将模型与设备集群切分为阶段–网格对的方法。我们的优化目标是最小化整个计算图的端到端流水线执行延迟。以往工作仅考虑了简化问题，例如假设每个阶段的设备已预先分配，且所有阶段都具有固定的数据或算子并行方案。Alpa 通过联合考虑设备网格分配以及各阶段可能存在的不同算子内并行方案，消除了这些假设。

算子间并行的空间

假设计算图中的算子按拓扑顺序排列，记为 $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_{\text{intra}}(s_i,\text{Mesh}(n_i,m_i)) $$ 表示在 $n_i\times m_i$ 子网格上执行阶段 $s_i$ 的延迟，该延迟由算子内编译阶段（§4）通过 ILP 最小化并返回。如下 图 5 所示，假设流水线包含 $B$ 个不同的输入微批次，则整个计算图的最小总延迟写为： $$ T^*=\min_{s_1,\ldots,s_S;,(n_1,m_1),\ldots,(n_S,m_S)} \left( \sum_{i=1}^{S} t_i + (B-1)\cdot \max_{1\le j\le S}{t_j} \right). $$

整体延迟由两部分构成：第一项是所有阶段的延迟之和，可解释为第一个微批次通过流水线的延迟；第二项是其余 $B-1$ 个微批次的流水化执行时间，其上界由最慢阶段决定（图 5 中为阶段 3）。

我们在求解上述目标时加入两个额外约束：（1）对于前向传播中的某个算子，我们希望将其与对应的反向算子共置在同一子网格上。由于反向传播通常使用前向传播中生成的相似张量，这样可以有效减少为反向传播获取所需张量的通信量。我们将前向与反向的延迟求和作为 $t_{\text{intra}}$，因此上式反映了包含前向与反向传播在内的总延迟。（2）切分得到的子网格 $(n_1,m_1),\ldots,(n_S,m_S)$ 必须完全覆盖 $N\times M$ 的集群网格，确保不浪费任何计算设备资源。下面我们将详细介绍动态规划（DP）形式化。

图 5： 流水线总延迟示意图，由两部分决定：所有阶段的总延迟 $(t_1+t_2+t_3+t_4)$，以及最慢阶段的延迟 $((B-1)\cdot t_3)$。

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 算法得到的子网格，我们采用一种枚举策略：优先为较大的子网格分配设备，再为较小的子网格分配设备。当存在多个具有相同子网格形状的流水线阶段时，我们倾向于将相邻的流水线阶段在设备网格上放置得更近，以降低通信延迟。

这种对子网格形状的简化在大多数云端深度学习环境中都表现良好：在 AWS 上，GPU 实例通常包含 1、2、4 或 8 块 GPU；在 GCP 上，TPU 实例通常包含 8、32、128、256 或 512 个 TPU。被该假设排除的子网格形状为 $n>1$ 且 $m<M$ 的情况；我们观察到这些形状通常效果较差，因为可以用形状为 $(n',M)$ 的替代子网格（其中 $n'\cdot M=n\cdot m$），其拥有更多通过高带宽互连的设备。通过这一缩减，我们只需确保 $$ \sum_{i=1}^{S} n_i\cdot m_i = N\cdot M. $$

为求解上式中的 $T^*$，我们设计了一个 DP 算法。DP 首先枚举第二项 $$ t_{\max}=\max_{1\le j\le S} t_j $$ 并在给定 $t_{\max}$ 的情况下最小化第一项 $$ t_{\text{total}}(t_{\max})=\sum_{i=1}^{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(s,k,d;t_{\max})= \min_{\substack{ k\le i\le K\ n_s\cdot m_s\le d }} \left{ \begin{aligned} & t_{\text{intra}}\big((o_k,\ldots,o_i),\text{Mesh}(n_s,m_s),s\big) \ &\quad + F\big(s-1,i+1,d-n_s\cdot m_s;t_{\max}\big) \end{aligned} ;\middle|; t_{\text{intra}}\big((o_k,\ldots,o_i),\text{Mesh}(n_s,m_s),s\big)\le t_{\max} \right}. $$

最终，总体最优延迟为： $$ T^*(t_{\max})=\min_s {F(s,0,N\cdot M;t_{\max})} + (B-1)\cdot t_{\max}. $$

$t_{\text{intra}}((o_k,\ldots,o_i),\text{Mesh}(n_s,m_s),s)$ 的取值由算子内编译阶段决定。它表示在具有 $s$ 个后续阶段的情况下，将子图 $(o_k,\ldots,o_i)$ 映射到子网格 $\text{Mesh}(n_s,m_s)$ 上执行所能达到的最小延迟。

需要注意的是，$\text{Mesh}(n_s,m_s)$ 表示的是一组物理设备。因此，我们会枚举所有满足 $n_l\cdot m_l=n_s\cdot m_s$ 的逻辑设备网格形状 $(n_l,m_l)$。对于每一种选择，我们以子图 $(o_k,\ldots,o_i)$、逻辑网格 $(n_l,m_l)$ 以及其他算子内并行选项作为输入，调用算子内编译阶段以获得一个算子内并行方案。随后，我们使用该方案并结合所有其他底层编译优化（例如算子融合、内存规划）对该子图进行编译，生成一个可执行程序，以便进行精确性能剖析。

通过对该可执行程序进行剖析，我们可以得到阶段的执行延迟 $t_l$，以及在每个设备上运行该阶段所需的内存 $mem_{\text{stage}}$，和用于存储中间激活值的内存 $mem_{\text{act}}$。随后，根据所选的流水线执行调度策略，检查所需内存是否能够满足设备内存 $mem_{\text{device}}$ 的限制。例如，对于 1F1B 调度，需要满足： $$ mem_{\text{stage}} + s\cdot mem_{\text{act}} \le mem_{\text{device}}. $$

我们选择在满足内存约束的前提下，使 $t_l$ 最小的逻辑设备网格形状；如果不存在任何满足条件的选择，则设定 $t_{\text{intra}}=\infty$。

我们的算法在设计上借鉴了 TeraPipe 的思想。然而，TeraPipe 假设所有流水线阶段是相同的，其目标是寻找将输入 token 划分为不同大小微批次的最优方式；而 Alpa 的目标是将计算图中的算子分组为不同的流水线阶段，并假设输入微批次大小相同。此外，Alpa 还在算子间并行的 DP 算法中，为每个流水线阶段优化其设备网格形状。

算法复杂度

对于一个固定的 $t_{\max}$，DP 算法在 $$ O(K^3NM(N+\log M)) $$ 时间内完成切分计算。$t_{\max}$ 的候选值数量至多为 $$ O(K^2(N+\log M)), $$ 对应于所有 $i,j=1,\ldots,K$ 以及所有子网格选择下的 $t_{\text{intra}}((o_i,\ldots,o_j),\text{Mesh}(n_s,m_s))$。

因此，该 DP 算法的总体复杂度为： $$ O\big(K^5NM(N+\log M)^2\big). $$ 对于包含上万个算子的计算图而言，这一复杂度是不可接受的。为此，我们引入了若干实用的性能优化策略。

性能优化一：提前剪枝

我们采用了一种与 TeraPipe 类似的优化策略，从小到大枚举 $t_{\max}$。当 $B\cdot t_{\max}$ 已经大于当前最优的 $T^*$ 时，立即停止枚举，因为更大的 $t_{\max}$ 不可能再产生更优解。

此外，在枚举 $t_{\max}$ 时，我们仅当当前 $t_{\max}$ 相比上一次至少大于 $\varepsilon$ 时才进行评估。这样可以保证 DP 算法得到的解与全局最优解之间的差距最多为 $B\cdot\varepsilon$。在实践中，我们取 $\varepsilon=10^{-6},\text{s}$，并发现对于所有评估设置，算法输出的解与真实最优解（即 $\varepsilon=0$）完全一致。

算法： 算子间编译阶段总结。

性能优化二：算子聚类

在计算图中，许多算子并不消耗大量计算资源（例如 ReLU），它们的具体放置位置对总体执行时间影响较小。为此，我们设计了另一种 DP 算法，对相邻算子进行聚类，以减小公式所作用的计算图规模。

我们将算子 $(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$ 具有如下最优子结构：

$$ G(k,r)= \min_{1\le i\le k} \left{ \max\bigl(G(i-1,r-1),C(i,k)\bigr) ;\middle|; \text{FLOP}(o_i,\ldots,o_k)\le(1+\delta)\frac{\text{FLOP}_{\text{total}}}{L} \right}. $$

其中，$C(i,k)$ 表示算子 $(o_i,\ldots,o_k)$ 从 $(o_1,\ldots,o_{i-1})$ 接收到的输入总大小，而 $$ \text{FLOP}_{\text{total}}=\text{FLOP}(o_1,\ldots,o_K) $$ 表示整个计算图的总 FLOP。我们确保每个聚类后的层，其 FLOP 不超过平均每层 FLOP 的 $(1+\delta)$ 倍，同时最小化通信开销。对于通信代价相同的解，我们进一步选择结构最均匀的方案，即最小化各层 FLOP 的方差。

通过该 DP 算法，我们可以在 $$ O(K^2L) $$ 时间内计算出最优的层聚类结果。需要注意的是，这里的 $L$ 是算法的一个超参数。在实践中，我们会根据设备数量以及计算图中重计算算子的数量来选择一个较小的 $L$。实验表明，不同 $L$ 的取值对最终性能影响不大。

算法 1 总结了算子间编译阶段的整体流程，并展示了其与 §4 中算子内编译阶段的交互关系。

并行化编排

在确定了阶段、设备网格及其映射关系之后，在算子内层面，Alpa 会针对每个阶段及其分配的设备网格进行编译，并遵循由 ILP 求解器输出的算子内并行方案。该编译过程依赖 XLA 与 GSPMD，为每个阶段–网格对生成并行可执行程序。当需要时，编译器会自动插入集合通信原语，以处理算子内并行引起的网格内通信。

在算子间层面，Alpa 还实现了一个额外的并行化编排阶段，用于处理阶段之间的跨网格通信，并生成用于算子间并行执行的静态指令。

跨网格重新分片

现有的手动系统（如 Megatron-LM）通常限制所有流水线阶段具有相同的数据并行与张量模型并行度，因此阶段之间的通信可以简单地通过等价设备网格上对应设备之间的点对点 send/recv 实现（见 图 a）。而在 Alpa 中，相邻两个阶段所在的设备网格可能具有不同的网格形状，并且阶段之间传输的张量也可能具有不同的分片规格（见 图 b 与 图 c）。我们将这种通信模式称为跨网格重新分片，它本质上是一个多对多的多播通信问题。

图 6： 跨网格重新分片。红色箭头表示在低速连接上的 send/recv，绿色箭头表示在高速连接上的 all-gather。(a) Megatron-LM 中用于相同网格形状的 scatter–gather 优化。(b) 针对不同网格形状的朴素 send/recv。(c) 针对不同网格形状的广义局部 all-gather 优化。

给定发送端与接收端设备网格上张量的分片规格，Alpa 通过两轮处理生成跨网格通信方案。第一轮中，Alpa 计算源网格与目标网格之间张量分区（亦称 tile）的对应关系，并据此生成源设备与目标设备之间的点对点 send/recv 原语，以完成所需通信。第二轮中，Alpa 识别目标张量分片规格中存在复制维度的情况。在这种情况下，张量只需在两个网格之间传输一次，随后再利用目标网格内部的高带宽连接，通过 all-gather 在目标网格内交换数据（见 图）。在此过程中，Alpa 会将第一轮生成的部分 send/recv 重写为 all-gather，以避免重复通信。

我们将该方法称为局部 all-gather 的跨网格重新分片。由于按设计阶段之间的通信量通常较小，实验结果表明该方法具有令人满意的性能。关于最优跨网格重新分片方案的研究将留作未来工作。

生成流水线执行指令

作为最后一步，Alpa 会生成用于在集群上启动训练的静态执行指令。由于不同阶段包含的算子集合不同，且可能位于形状不同的设备网格上，与许多 SPMD 风格的流水线并行训练系统不同，Alpa 采用一种 MPMD 风格的运行时 来编排算子间并行执行，即为每个设备网格生成不同的静态执行指令。

Alpa 设计了一组用于算子间并行执行的指令，包括：阶段内张量的内存分配与释放、按照跨网格重新分片方案在阶段之间进行张量通信、同步以及计算等。根据用户选择的流水线调度策略，Alpa 使用一个驱动进程提前生成这些指令，并在执行前将完整的指令列表分发到各个工作节点，从而避免运行时的驱动–工作节点协调开销。

局限性与讨论

在本节中，我们讨论这种并行化视角的优势以及算法中的若干局限性。

与现有手动组合数据并行、算子并行和流水线并行的方法（如 3D 并行 和 PTD-P）相比，Alpa 提出的算子间并行与算子内并行的分层视角在以下三个方面显著提升了灵活性：（1）流水线阶段可以包含不均匀数量的算子或层；（2）Alpa 中的流水线阶段可以映射到形状不同的设备网格；（3）在每个阶段内部，数据并行与算子并行配置可以在逐算子粒度上进行非均匀定制。综合来看，这些特性使 Alpa 能够统一现有的模型并行方法，并泛化到具有更高异构性的模型架构和集群配置。

尽管具备上述优势，Alpa 当前的优化算法仍存在一些局限性：

• 未建模跨阶段通信代价。 由于跨阶段通信在设计上通常较小，Alpa 并未显式建模不同阶段之间的通信代价。事实上，无论是在 DP 还是 ILP 中建模该代价都是可行的，但这会要求枚举指数级数量的算子内并行方案与 DP 状态。
• 微批次数量为超参数。 当前算子间编译阶段包含一个超参数——微批次数 $B$，该参数未被纳入现有形式化中进行优化，而是通过枚举搜索获得。
• 流水线调度模型受限。 算子间编译阶段采用的是静态线性流水线调度，尚未考虑更动态的调度方式，例如在计算图中将不同分支并行地映射到不同设备上。
• 未优化计算与通信重叠。 Alpa 尚未针对计算与通信的重叠进行优化，并且只能处理在编译期已知所有张量形状的静态计算图。

尽管如此，我们在弱扩展实验中的结果表明，Alpa 仍能够为多种重要模型生成近似最优的执行计划。

实验评估

Alpa 的实现规模约为 1.6 万行 Python 与 6000 行 C++。Alpa 以 Jax 作为前端、XLA 作为后端，编译阶段基于 Jaxpr 与 HLO 等中间表示实现。在分布式运行时方面，我们使用 Ray 的 actor 来实现设备网格工作进程，使用 XLA runtime 执行计算，并使用 NCCL 进行通信。

我们在多个参数规模达数十亿的大模型上对 Alpa 进行了评估，包括 GPT-3、GShard Mixture-of-Experts（MoE） 和 Wide-ResNet。测试平台为一个典型的集群环境，由 8 个节点、64 块 GPU 组成。每个节点为一台 Amazon EC2 p3.16xlarge 实例，包含 8 块 NVIDIA V100 16GB GPU、64 个 vCPU 和 488GB 内存。节点内的 8 块 GPU 通过 NVLink 互连，8 个节点位于同一个放置组中，节点间带宽为 25Gbps。

表 4： 端到端评估中使用的模型。LM 表示语言模型，IC 表示图像分类。

我们将 Alpa 与两个用于 GPU 大规模模型训练的最先进分布式系统进行比较。此外，我们还分别隔离不同的编译阶段，对优化算法进行了消融实验，并给出了 Alpa 所生成执行计划的案例分析。

端到端性能

模型与训练负载

我们选取了 表 4 中列出的三类模型，涵盖了同构架构与异构架构的模型。GPT-3 是一种通过堆叠大量 Transformer 层构成的同构 Transformer 语言模型，其模型并行方案已被广泛研究。GShard MoE 是一种结合稠密层与稀疏层的语言模型，其中在每两层 Transformer 的末尾使用 Mixture-of-Experts 层替换 MLP。Wide-ResNet 是一种具有更大通道数的 ResNet 变体，其结构与 Transformer 模型差异显著，且此前并不存在人工设计的并行化方案。

为研究大模型训练能力，我们遵循常见的机器学习实践，随着 GPU 数量的增加同步扩展模型规模，其参数范围如 表 4 所示。具体而言，对于 GPT-3，我们按照已有工作的方法同时增加隐藏层维度与层数；对于 MoE，主要增加专家数量；对于 Wide-ResNet，则增加卷积层的通道数与宽度因子。对每种模型，我们采用文献中推荐的全局批大小以保持一致的统计行为，并针对每个模型与系统配置调优最佳微批大小以最大化系统性能。梯度在微批之间进行累积。更详细的模型规格见附录。

对比基线

针对每种模型，我们选择了一个强有力的对比基线。对于 GPT-3，我们使用 Megatron-LM v2 作为基线系统。Megatron-LM 是目前在 GPU 上训练同构 Transformer 语言模型的最先进系统，它结合了数据并行、流水线并行以及人工设计的算子并行。这些并行方式的组合由三个整数参数控制，分别指定各类并行的并行度。我们按照其论文中的建议，对这三个参数进行网格搜索，并报告最佳配置的结果。由于 Megatron-LM 专门针对 GPT 类模型设计，因此不支持 表 4 中的其他模型。

图 7： 端到端评估结果。“×” 表示内存不足（out-of-memory）。黑色方框表示线性扩展。

我们使用 DeepSpeed 作为 MoE 的对比基线。DeepSpeed 提供了当前最先进的 GPU 上 MoE 训练实现，结合了为 MoE 层手工设计的算子并行方案以及基于 ZeRO 的数据并行。这些技术的组合由若干整数参数控制，用于指定各类并行方式的并行度。我们同样对这些参数进行网格搜索并报告最佳结果。由于 DeepSpeed 在 GPT-3 上的性能与 Megatron-LM 相当或更差，因此我们在 GPT-3 上不再报告 DeepSpeed 的结果。需要注意的是，原始的 GShard-MoE 实现仅在 TPU 上可用，因此我们未纳入其结果，尽管其并行策略已被 Alpa 的策略空间所覆盖。

对于大规模 Wide-ResNet，目前不存在专用系统或人工设计的并行方案。我们使用 Alpa 构建了一个基线方案 PP-DP，其并行空间仅包含数据并行与流水线并行，以模拟 PipeDream 与 Dapple 的并行空间。

对于所有模型，我们还给出了仅使用算子内并行或仅使用算子间并行的 Alpa 结果，以模拟部分自动并行系统的性能。开源系统 FlexFlow 不支持我们评估的模型，因为其缺乏多项必要算子支持（例如 Layer Normalization、混合精度算子）。Tofu 仅支持单节点执行且未开源。由于理论与实践上的限制，我们未纳入这些系统的结果，并且也不期望 FlexFlow 或 Tofu 在本评估中超越最先进的人工基线。

评估指标

Alpa 不会改变同步梯度下降算法的语义，因此不会影响模型收敛性。我们在评估中以训练吞吐量作为指标，研究在增加 GPU 数量的同时扩展模型规模时的弱扩展性。按照已有工作的方法，我们使用整个集群的聚合 PFLOPS 作为性能指标。该指标通过在充分预热后、使用虚拟数据运行若干批次来测量。所有实验结果的标准差均在 0.5% 以内，因此图中未绘制误差条。

GPT-3 结果

GPT-3 的并行化方案已被广泛研究。我们在 图 7a 中观察到，通过网格搜索得到的最佳人工方案，使 Megatron-LM 在 GPT-3 上实现了超线性弱扩展。尽管如此，Alpa 能够自动生成执行计划，并在多个设置下取得了略优于 Megatron-LM 的扩展性能。

与仅使用算子内并行的方法相比，结果与近期研究一致：当 GPU 数量超过 16 时，“仅算子内并行”性能较差，因为即使是最优方案，也需要在跨节点连接上进行大量张量通信，使通信成为瓶颈。令人惊讶的是，“仅算子间并行”在最多 64 GPU 的情况下仍能保持线性扩展。

我们进一步分析了 Megatron-LM 中人工方案的最优参数配置，并将其与 Alpa 生成的方案进行比较，得到两个主要结论。首先，在 Megatron-LM 中，除少数特殊情况（例如在 64 GPU 上拟合 39B 模型）外，最优方案通常将 TMP 设为 1；当流水线并行无法单独满足显存约束时，才会引入算子并行。同时，只要内存允许，数据并行度会被最大化。在实践中，通过启用 梯度累积（GA） 来满足期望的全局批大小。GA 能摊薄数据并行的通信开销，并减少流水线并行中的空泡，但算子并行的通信量会随 GA 步数线性增长，从而使 TMP 处于不利地位。

其次，Alpa 生成的方案与 Megatron-LM 中表现最好的人工方案高度相似，具有以下特征：（1）阶段大小均衡；（2）在内存压力较小时沿批次维度进行划分，而在内存受限时沿非批次维度进行划分。一个关键差异在于：当存在数据并行时，Alpa 还会对权重更新算子进行划分，这也是其性能略优于 Megatron-LM 的原因之一。这得益于 Alpa 作为通用编译系统，能够组合更丰富的并行化方法，而 Megatron-LM 目前尚未支持权重更新分片。

MoE 结果

DeepSpeed 采用了一种由 GShard 提出的、针对 MoE 模型的人工算子并行方案，称为专家并行。该方法遵循一条简单规则：在 MoE 层中沿专家维度对算子进行划分，而在非专家层中切换回数据并行。随后，该专家并行与 ZeRO 数据并行及 TMP 相结合。这些技术全部属于算子内并行范畴。

不幸的是，DeepSpeed 的专用实现并未包含任何算子间并行方法，而算子间并行对于在跨节点、低带宽环境下扩展至多节点是必需的。因此，在本集群上，DeepSpeed 只能在单节点（≤ 8 GPU）内保持良好性能。“仅算子内并行”同样由于相同原因，无法在多节点上扩展。

另一方面，“仅算子间并行”在 32 GPU 与 64 GPU 情况下会出现内存不足，这是因为当 GPU 数量大于模型层数时，难以对模型进行均匀切分。不均衡的切分会使部分内存密集型阶段发生内存溢出。

图 8： 算子内并行消融实验结果。“×” 表示内存不足，黑色方框表示线性扩展。

相比之下，Alpa 能够自动发现同时结合算子内并行与算子间并行的最优执行计划。在算子内并行方面，得益于基于 ILP 的算子内编译阶段，Alpa 找到了一种类似专家并行的策略，并将其与 ZeRO 数据并行相结合。随后，Alpa 构建流水线阶段，并利用算子间并行，在低速连接上尽量减少通信量。Alpa 在 16 GPU 上保持线性扩展，并能够良好地扩展到 64 GPU。与 DeepSpeed 相比，Alpa 在 2 个节点上实现了 3.5× 加速，在 4 个节点上实现了 9.7× 加速。

Wide-ResNet 结果

不同于前两类通过堆叠相同层构成的模型，Wide-ResNet 具有更加异构的架构特征。随着数据批次在网络层中前向传播，激活张量的大小逐渐减小，而权重张量的大小不断增大，这导致各层之间的内存使用和计算强度分布极不均衡。对于这种模型，人工设计并行方案非常困难，甚至几乎不可能。然而，Alpa 仍然在 32 GPU 上实现了 80% 的扩展效率。

基线方法 “PP-DP” 与 “Inter-op only” 在训练大模型时会出现内存不足，因为它们无法对权重进行划分以降低内存占用，并且难以构建内存均衡的流水线阶段。“Intra-only” 则需要在低速连接上进行大量通信，因此无法扩展到多节点。关于 Wide-ResNet 的执行计划案例分析见 §8.6。

算子内并行消融实验

我们研究了算子内并行优化算法的有效性，将基于 ILP 的方案与 ZeRO 优化器及基于规则的划分策略进行比较。

实验设置

我们在模型规模的弱扩展场景下进行基准测试，与 §8.1 类似，但禁用流水线并行和梯度累积以控制变量。实验在一台 AWS p3.16xlarge 实例（8 GPU）上完成。为了模拟单节点大规模训练环境，我们采用了更大的隐藏维度、更小的批大小和更少的层数，相较于 §8.1 中的模型配置。

基线方法

我们比较了多种自动化算子内并行方案：

• Data：标准数据并行
• ZeRO-2：对梯度和优化器状态进行分片的数据并行
• ZeRO-3：在 ZeRO-2 基础上进一步对参数进行分片
• Heuristic：结合规则与 GSPMD 中分片传播的方法，对每个输入张量的最大维度进行划分
• ILP：本文提出的基于 ILP 的方案

实验结果

如 图 8 所示，“Data” 很快出现内存不足，无法训练大模型。“ZeRO-2” 与 “ZeRO-3” 虽然解决了内存问题，但由于始终需要通信梯度，未对通信进行优化，当梯度远大于激活值时性能显著下降。“Heuristic” 通过对所有张量进行划分解决了内存问题，但可能因通信量过大而变慢。ILP 自动分片在所有情况下表现最佳，并保持近线性扩展，因为它能够找到始终最小化通信开销的正确划分方案。

算子间并行消融实验

我们进一步研究了算子间并行优化算法的有效性，将该算法记为 DP。

实验设置

我们在 GPT 和 Wide-ResNet 上报告三种 DP 变体的性能，实验设置与 §8.1 相同。

基线方法

我们将 DP 算法与两种基于规则的方法进行比较：

• Equal operator：禁用 DP 算子聚类，为每个阶段分配相同数量的算子
• Equal layer：限制 DP，使所有流水线阶段包含相同数量的层

实验结果

如 图 9 所示，“DP” 始终优于 “Equal operator”，因为后者会将本应映射到不同设备网格的算子合并在一起，而 Alpa 能够基于通信代价与计算均衡性对算子进行聚类。“DP” 是否优于 “Equal layer” 取决于模型结构：在 GPT 这类同构模型上，最优解为各阶段使用相同层数，因此 “Equal layer” 与 “DP” 表现相同；而在 Wide-ResNet 上，最优方案需要为不同阶段分配不同层数，因此 “Equal layer” 明显劣于完整的 DP 算法。在 Wide-ResNet 的 32 GPU 设置下，DP 分别比 “Equal operator” 和 “Equal layer” 快 2.6× 和 1.6×。

图 9： 算子间并行消融实验结果。

图 10： Alpa 在所有 GPT 模型上的编译时间。模型规模与 GPU 数量同时扩展。

编译时间

图 10 展示了 §8.1 中所有 GPT 设置下 Alpa 的编译时间。编译时间对应于在给定微批次数 $B$ 的情况下，完整运行一次 算法 1。从结果可以看出，Alpa 能够良好地扩展到大模型或大规模集群，因为编译时间随着模型规模和 GPU 数量线性增长。

表 5 给出了我们评估中最大 GPT 模型（39B，64 GPU）的编译时间分解。大部分时间花在枚举阶段–网格对并对其进行性能剖析上。在编译阶段，我们通过使用分布式工作进程并行编译不同阶段来加速；在剖析阶段，我们采用了一个构建在 XLA 指令级别 的简单代价模型来加速，该模型使用分段线性模型估计矩阵乘法与通信原语的代价。借助这些优化，模型的编译与搜索最多只需数小时，这是可以接受的，因为相比之下实际训练时间往往需要数周。

表 5： GPT-39B 的编译时间分解。

跨网格重新分片

我们在 Wide-ResNet 上评估了针对不同形状设备网格的跨网格重新分片中的广义局部 all-gather 优化，结果如 图 11 所示。“signal send/recv” 是一个合成场景，仅在阶段之间发送 1 字节信号，可视为性能上界。“w/o local all-gather” 禁用了局部 all-gather 优化，仅使用 send/recv；“w/ local all-gather” 启用了局部 all-gather 优化，将更多通信从慢速连接迁移到高速本地连接上，在 32 GPU 上带来了 2.0× 的加速。

图 11： Wide-ResNet 上的跨网格重新分片。

案例研究：Wide-ResNet

我们在 图 12 中可视化了 Alpa 在 16 GPU 上为 Wide-ResNet 找到的并行化策略；4 GPU 与 8 GPU 的结果见附录 C。在 4 GPU 上，Alpa 仅使用算子内并行：前几十层沿批次轴进行划分，而在最后几层切换为沿通道轴划分。在 16 GPU 上，Alpa 将模型切分为 3 个阶段，分别为阶段 1、2、3 分配 4、4、8 GPU。由于激活张量大于权重张量，前两个阶段优先采用数据并行；在第三阶段，ILP 求解器为卷积算子找到了非平凡的划分方式。结果表明，对于像 Wide-ResNet 这样具有异构结构的模型，即便是领域专家，手动设计此类策略也十分困难。

图 12： Wide-ResNet 在 16 GPU 上的并行策略可视化。不同颜色表示张量分布到的设备；灰色块表示张量在设备间复制。每个卷积层与全连接层的输入数据及其激活可以沿批次轴和隐藏轴划分；权重可以沿输入通道轴与输出通道轴划分。

相关工作

数据并行训练系统

Horovod 与 PyTorch DDP 是两种常用的数据并行训练系统，通过 all-reduce 同步梯度。BytePS 统一了 all-reduce 与参数服务器，并利用数据中心集群中的异构资源。AutoDist 使用学习方法来组合数据并行训练策略。ZeRO 通过减少复制张量来改进数据并行的内存使用；MiCS 在 ZeRO 之上进一步最小化通信规模，以提升公有云上的可扩展性。在 Alpa 中，数据并行可视为算子内并行的一种特例——沿批次轴进行划分。

模型并行训练系统

模型并行的两大类别已在 §2 中讨论。TensorFlow、GSPMD 与 OneFlow 提供了注解式 API，允许用户手动指定算子内并行方案。ColocRL 将互不重叠的模型分区放置在不同设备上而不进行流水线化，因此只有在模型存在并行分支时才会产生并发；相比之下，GPipe 通过将输入数据拆分为微批次来形成流水线并行。PipeDream 通过使用异步训练、降低内存占用并与数据并行结合来改进 GPipe，但 PipeDream 是异步系统，而 Alpa 是同步训练系统。TeraPipe 为基于 Transformer 的语言模型发现了一种新的流水线并行维度。Google 的 Pathway 系统是与 Alpa 并行开展的工作，其倡导一种结合 SPMD 与 MPMD 的单控制器运行时架构；这与 Alpa 的运行时设计相似——Alpa 在算子内并行中使用 SPMD，在算子间并行中使用 MPMD。

模型并行方案的自动搜索

另一类相关工作聚焦于自动发现模型并行训练方案。Tofu 提出了一种动态规划算法，用于在单节点上线性计算图中生成最优的算子内并行策略。FlexFlow 提出了 SOAP 形式化，并开发了基于 MCMC 的随机搜索算法，但其仅支持设备放置而不支持流水线并行；其搜索算法难以扩展到大规模计算图或集群，并且缺乏最优性保证。TensorOpt 提出了一种动态规划算法，用于自动搜索同时考虑内存与计算代价的算子内并行策略。Varuna 面向低带宽集群，重点在于自动化流水线并行与数据并行。Piper 也能够同时寻找算子间与算子内并行策略，但其依赖人工设计的算子内并行方案，并且假设均匀网络拓扑以及异步流水线调度。

大规模模型训练技术

除并行化之外，还有许多互补技术可用于训练大规模模型，例如内存优化、通信压缩以及低精度训练。Alpa 能够整合其中的多种技术。例如，Alpa 使用重计算（rematerialization）来降低内存占用，并采用混合精度训练来加速计算。

深度学习编译器

编译器技术已被广泛用于优化深度学习模型的执行。大多数相关工作主要关注单设备上的计算优化。相比之下，Alpa 是一个面向分布式训练的编译器，支持一个全面的执行计划空间。

其他领域的分布式张量计算

除深度学习之外，针对线性代数和模板（stencil）计算等领域，也已开发了分布式张量计算的库与编译器。但与 Alpa 不同的是，这些系统并未考虑深度学习所需的关键并行化技术。

结论

我们提出了 Alpa，一种用于自动化模型并行分布式训练的新体系结构，建立在对机器学习并行化方法的新视角之上：算子内并行与算子间并行。Alpa 构建了一个分层并行空间，并通过一组编译阶段，在每个并行层次上推导高效的并行执行计划。Alpa 在两个不同的计算粒度上，对分布式计算设备上的并行执行进行协调。

为分布式模型并行深度学习设计高效的并行化方案，历来是一项劳动密集型的任务。我们相信，Alpa 将降低分布式模型并行训练的门槛，并加速新一代大规模深度学习模型的应用与普及。