ON OPTIMIZING THE COMMUNICATION OF MODEL PARALLELISM

发表时间: 2022-11 · arXiv:2211.05322

原文: https://arxiv.org/abs/2211.05322

文章标题：论模型并行通信的优化
作者/机构：Yonghao Zhuang, Hexu Zhao, Lianmin Zheng, Zhuohan Li, Eric P. Xing, Qirong Ho, Joseph E. Gonzalez, Ion Stoica, Hao Zhang

速读

一句话结论 本文针对大模型训练中算子内与算子间并行结合时产生的跨设备网格重分片通信瓶颈，提出了一套从底层通信原语到顶层流水线调度的全栈优化方案 AlpaComm，在端到端多节点训练中实现了最高 1.5 倍的吞吐量提升。

要解决什么问题 在训练超大模型时，通常需要将算子内并行（张量切片）和算子间并行（流水线）结合使用。这就导致在相邻流水线阶段的边界处，张量需要从一个设备网格发送到另一个设备网格。由于源网格和目标网格对张量的切片布局要求可能不同，这种数据传输不仅是简单的点对点发送，还伴随着复杂的张量布局转换，即跨设备网格重分片（cross-mesh resharding）。原有的深度学习系统在处理这种多对多多播通信时存在严重卡点：首先，它们往往使用次优的通信原语，导致跨节点慢速网络（如以太网）的带宽被无效占用，且多个 GPU 竞争同一个主机的网卡带宽造成拥塞；其次，现有的同步流水线调度（如 1F1B）具有严格的数据依赖，使得这种昂贵的跨网格通信时间无法与计算重叠，最终在网络拓扑复杂或模型存在大量跨网格跳跃连接时，引发严重的通信延迟和流水线气泡。

怎么做的 核心思路是将复杂的跨网格重分片任务自底向上分为三个层级进行优化：单元通信任务的底层原语替换、多任务的主机级负载均衡、以及全局的流水线调度重叠。首先，方法将整个张量的重分片分解为多个独立的单元通信任务，每个任务负责将一个唯一的数据切片从源设备子集发送到目标设备子集。针对单个单元任务，作者发现传统的点对点发送或依赖全局 all-gather 的策略会成倍增加跨主机延迟，因此提出使用广播（broadcast）作为替代方案。具体做法是将数据切片分成 $K$ 个分区，接收设备在收到一个分区后立刻作为发送方将其传递给下一个设备，形成流水线。假设发送固定大小对象耗时为 $t$，跨 $A$ 个节点的总延迟为： $$T_{bc}(A, B) = t + \frac{(A-1)t}{K}$$ 当 $K$ 足够大时，总延迟逼近物理极限 $t$。其次，针对多个单元任务并发时的网卡竞争问题，在主机级别建立负载均衡与调度模型。目标是最小化最后一个任务的完成时间： $$\min_{n_i^*, S_i} \max_i (S_i + T_i)$$ 为此，设计了结合带剪枝的深度优先搜索（处理小规模）和带随机化的贪心搜索（处理大规模）算法，为每个任务分配最优的发送主机并决定执行顺序，避免任何节点空闲或拥塞。最后，为了打破计算与通信的严格依赖，提出了一种对重叠友好的 eager-1F1B 流水线调度方案。该方案在预热阶段让每个阶段提前执行更多微批次的前向计算，从而在稳定的交替执行阶段中，将当前微批次的通信与后续微批次的计算重叠。对于后向传播，引入后向权重延迟技术，优先计算不依赖权重的激活梯度以触发跨网格通信，将权重梯度的计算延后以掩盖通信时间。

效果如何 实验在 AWS 集群上搭建，节点内配备 4 张 16GB 显存的 V100 GPU 并通过 NVLink 互连，节点间网络带宽为 10Gbps。评估选用了两种模型：堆叠标准层的 GPT 模型，以及带有长跳跃连接的 U-Transformer 模型。对比基线包括代表纯点对点通信路线的 Send/Recv、代表基于 all-gather 路线的先进系统 Alpa、仅使用广播但不做流水线重叠的 Broadcast，以及仅发送一字节以探测性能上限的 Signal Send/Recv。在多对多设备的微基准测试中，由于有效规避了跨节点 all-gather 的开销并优化了发送节点带宽，该方法比 Alpa 快 2.5 到 10 倍。在端到端多节点训练中，该方法在 GPT 模型上实现了 10% 的吞吐量提升；而在跨网格通信成为严重瓶颈的 U-Transformer 模型上，eager-1F1B 调度成功重叠了大量通信，相比 Alpa 实现了 50% 的加速，整体性能达到了理论上限的 75% 以上。该方法的主要代价是 eager-1F1B 调度和后向权重延迟会增加每个 GPU 缓存临时激活张量的数量，导致显存占用略有增加，但单卡的增加量仅为流水线阶段数与单层激活大小的乘积，远小于模型参数本身的体积。

A1 主要贡献

本文研究了一种在规模化模型并行深度学习（DL）中新颖且重要的通信模式，我们称之为跨设备网格重分片（cross-mesh resharding）。当模型并行的两种范式——算子内并行（intra-operator parallelism）和算子间并行（inter-operator parallelism）——结合起来以支持大型集群上的大模型时，这种模式就会出现。在跨设备网格重分片中，一个分片的张量需要从一个源设备网格（source device mesh）发送到一个目标设备网格（destination device mesh），而在目标设备网格上，该张量可能以相同或不同的布局进行分布。

核心问题：现有的解决方案要么是次优的，要么无法泛化到由不同模型架构和并行策略产生的不同网络拓扑或张量布局。例如，将它们应用于新兴模型如 U-Transformer 时，会导致通信速度大幅下降。因此，本文的核心问题是：给定任意模型及其组合的算子间和算子内并行计划，执行跨设备网-格重分片的最佳方式是什么？

研究目标：本文将跨设备网格重分片问题形式化为一个多对多多播（many-to-many multicast）通信问题，并揭示了在模型并行深度学习背景下的几个独特挑战：

创新点与主要贡献：
为了应对这些挑战，本文提出了新的解决方案，以优化单个任意的跨设备网格重分片任务，并集体优化模型并行作业中的所有跨设备网格重分片任务。主要贡献总结如下：


图1. 模型并行应用于一个多层感知机（MLP），其中每个节点代表一个算子及其输出张量，实线箭头表示数据流向。(b)-(d)中带虚线边界的节点显示了matmul2所需的输入布局。在(d)中结合两种并行方式时，为了在两个设备网格之间交换张量和转换其布局，出现了跨设备网格重分片。

A3 背景知识

本节介绍模型并行深度学习中的基本通信特性，并描述一种在大型模型并行工作负载中出现的新通信模式，即跨设备网格重分片。

2.1 模型并行

算子内并行（Intra-op parallelism）：算子内并行是指沿着某些张量轴对算子及其输入输出张量进行分片或复制，并将算子计算的不同区域分配给并行设备的方法，如图1(b)所示。在实践中，算子内并行以SPMD（单程序多数据）方式实现【18, Gshard: Scaling giant models with conditional computation and automatic sharding, 2020, arXiv preprint】【43, Gspmd: general and scalable parallelization for ml computation graphs, 2021, arXiv preprint】。当张量沿着不同的轴进行分片时，它们可能以不同的分布式布局存储在多个设备上。如果输入张量的布局与算子所需的输入布局不一致，就需要进行布局转换（也称为重分片 resharding）。在算子内并行中，这种布局转换很容易通过所有参与设备之间的集体通信操作（如all-reduce, all-gather, all-to-all）来完成。

算子间并行（Inter-op parallelism）：算子间并行是对计算图（而不是单个算子或张量）进行分区，并将每个生成子图（称为一个阶段 stage）放置在一个设备上。如果一个算子的输入张量由另一台设备上的另一算子生成，则需要进行点对点（P2P）通信，将张量从源设备发送到目标设备，如图1(c)所示。流水线并行【10, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019, Advances in neural information processing systems】【22, Memory-efficient pipeline-parallel dnn training, 2021, International Conference on Machine Learning】是算子间并行的一个特例，其中所有设备被组织成一个流水线，通过让不同流水线阶段的设备以“流水线”方式同时计算来饱和系统；P2P通信在阶段之间使用。由于分区之间的数据依赖性，一些设备在等待其输入或最后一个分区完成时会处于空闲状态，从而产生“气泡”（bubbles）。

结合算子内和算子间并行：单独使用算子内或算子间并行都不足以训练像GPT-3这样的大型模型【4, Language models are few-shot learners, 2020, Advances in neural information processing systems】。在实践中，必须将它们结合起来。这种组合策略在许多模型并行系统【29, 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】【45, Alpa: Automating inter-and intra-operator parallelism for distributed deep learning, 2022, arXiv preprint】【43, Gspmd: general and scalable parallelization for ml computation graphs, 2021, arXiv preprint】中实现，首先使用算子间并行对计算图进行分区，然后使用算子内并行对每个阶段进行分片，如图1(d)所示。具体来说，图首先被划分为多个阶段，每个阶段分配给一组从集群中划分出来的设备，称为设备网格（device mesh）。一个阶段的算子和张量根据选定的算子内并行计划在该阶段分配的网格上进行并行化；集体通信仅在每个网格内的设备之间发生。在任何两个相邻阶段的边界处，需要在它们的网格之间交换张量。与算子间并行不同，张量可能在源和目标网格上以不同的布局分片——在这种情况下，通信不仅涉及传输张量，还涉及在源和目标设备组之间执行张量布局转换。我们将这种通信模式称为跨设备网格重分片（cross-mesh resharding），这是本文的重点。实践中，模型并行系统可能会次优地将跨设备网格重分片通信映射到深度学习集群中较慢的互连上，例如不同节点/主机之间带宽有限的以太网连接，这会严重影响整体训练或推理性能。

2.2 跨设备网格重分片

形式化定义：一个模型并行作业可能涉及许多跨设备网格重分片通信。我们将一个跨设备网格重分片任务定义为将单个张量从源网格发送到目标网格，同时满足布局要求所需的通信。为此，我们首先形式化设备网格和张量布局的定义。

设备网格（Device mesh）：我们遵循GSPMD【43, Gspmd: general and scalable parallelization for ml computation graphs, 2021, arXiv preprint】中对网格和张量布局的定义：网格是一个由相同处理器组成的n维数组。我们将网格MeshA的2D视图表示为(m1, m2)，其中 m1 × m2 = |MeshA|。

分片规范（Sharding spec）：一个分片的N维张量D在MeshA上的布局可以用一个分片规范来描述，记作一个N元素字符串：$X = X_{d_0}^0 X_1^{d_1} \dots X_{N-1}^{d_{N-1}}$，其中 $X_i \in \{S, R\}, d_i \in \{0, 1, 01\}, 0 \le i \le N-1$。S代表分片（sharded），R代表复制（replicated）。当张量维度被分片时，上标$d_i$表示分片映射到的网格维度。例如，$X^{d_i}=S^0$表示D的第i维沿着MeshA的第一维分片。图2说明了三种分片规范及其产生的张量布局。

数据切片（Data slice）：由于分片或复制，MeshA中的每个设备可能持有一个D的数据切片或副本。例如，在图2的第一个分片规范中，MeshA中的每个设备持有一个唯一的4×1数据切片。

跨设备网格重分片（Cross-mesh resharding）：一个跨设备网格重分片是一个在两个不重叠的设备网格MeshA和MeshB之间（$MeshA \cap MeshB = \emptyset$）的多对多多播通信任务。该任务旨在将分片在MeshA上（遵循分片规范$X_A$）的张量D发送到MeshB，使其在MeshB上的分片规范变为$X_B$。图2展示了两个跨设备网格重分片的例子。


图2. 跨设备网格重分片的两个例子。给定一个4×4矩阵和两个2×2的网格MeshA和MeshB，每个虚线框显示一个分片规范及其在网格每个设备上的张量布局。当在规范1和2之间以及规范2和3之间通信时，形成了两个跨设备网格重分片任务。

任务分解：无论$X_A$和$X_B$如何，两个网格之间传输的消息大小的下限是D的大小。因此，我们可以将一个跨设备网格重分片分解为一组单元通信任务，每个任务对应源MeshA上D的一个唯一数据切片$DS_i$，并负责将$DS_i$从MeshA发送到MeshB中的一个设备子集。由于源网格中的多个设备可能持有$DS_i$的副本，而目标网格中的多个设备可能需要$DS_i$，因此存在多种可能的通信路径。

全局视角：一个模型并行作业在两个流水线阶段的边界处交换的每个张量都有许多跨设备网格重分片任务。每个任务在每个前向和后向传播中都会根据阶段的计算依赖性重复执行。这为通过调度多个跨设备网格重分片任务和计算任务来隐藏通信时间提供了重叠的机会。

A2 方法细节

3 优化单个跨设备网格重分片任务

问题框架：一个跨设备网格重分片问题可以分解为许多单元通信任务，每个任务负责一个数据切片。我们的目标是尽快完成所有任务。为此，我们将原问题构建为一个两级优化：

  1. 优化单个单元通信任务：我们分析不同方法，并提出使用广播（broadcast）来实现这部分，并证明基于广播的策略可以获得最优性能。
  2. 负载均衡和调度：我们为跨设备网格重分片中的所有单元任务制定负载均衡和调度问题，并设计了两种算法来解决该问题。

ML集群设定：在深入细节之前，我们首先考虑具有以下特性的典型机器学习集群设置：

3.1 优化单个单元通信任务

单元任务定义：我们用$DS_i$表示第i个单元通信任务对应的数据切片。$DS_i$在设备集$N_i \subseteq MeshA$中有副本，我们希望将$DS_i$发送到设备集$M_i \subseteq MeshB$。一个单元任务不能直接由集体通信原语处理，因为其发送方和接收方不重叠。一个好的通信策略应该（1）最小化跨主机通信量，（2）最大化所有参与设备的总发送和接收带宽利用率。在以下分析中，我们假设发送一个固定大小的对象需要时间t，并忽略节点内通信时间。我们用$T(A, B)$表示将对象从一个设备发送到A个节点、每个节点B个设备的A×B个设备所需的时间。

Send/recv策略：最简单的策略是让唯一的发送方逐一将数据切片$DS_i$发送给$M_i$中的每个接收设备。总通信延迟为$T_{sr}(A, B) = ABt$，与设备总数成正比，如图3a所示。当有多个潜在发送方时，可以将工作负载分配给不同发送方来加速。

带本地allgather的Send/recv策略：一种直接的改进是让每个接收节点只接收一份$DS_i$的副本，然后使用快速的节点内连接在节点内的所有设备间传输切片。为此，我们将$DS_i$分成B个部分，将每个部分发送到特定接收节点上的B个设备之一。然后同一节点上的接收设备通过all-gather集体通信原语组装$DS_i$的完整副本。因此，一个发送方的总延迟为$T_{srla}(A, B) = At$，如图3b所示。

带全局allgather的Send/recv策略：我们可以将$DS_i$切成更细粒度的块，具体来说，可以把$DS_i$分成A×B个切片，并将一个切片发送给A×B个设备中的每一个。这一步的延迟是t。然后这些设备执行一个全局all-gather来为每个设备收集所有切片，使用典型的环形all-gather算法【24, The nvidia collective communication library, 2018, URL: https://developer.nvidia.com/nccl】，其延迟也是t。因此，理想的总延迟为$T_{srga}(x, y) = 2t$，不随节点或设备数量增长，如图3c所示。

基于广播的重分片策略：我们观察到可以通过重叠数据发送和接收设备间的数据交换来进一步优化。具体来说，当一个接收方开始从发送方接收数据时，它可以充当另一个发送方，将数据发送给其余的接收设备，如图3d所示。最终，这个策略可以看作是从发送设备到所有接收设备的广播。假设我们将数据切片$DS_i$分成K个分区，每个设备在接收到一个完整分区后就将其发送给下一个设备。发送每个分区的时间是$t/K$，总延迟是$T_{bc}(A, B) = \frac{t}{K} \cdot (K+A-1) = t + \frac{(A-1)t}{K}$，如图3d所示。当我们选择K相对于A足够大时（例如，实验中K≈100），总延迟约等于t，这是所有策略中最优的，并已达到通信上限。因为每个接收节点只有一个NIC，接收一个副本至少需要时间t。


图3. 单个通信任务的通信策略。橙色箭头代表跨主机通信，其上的数字代表通信延迟。蓝色箭头代表节点内通信。

3.2 多个单元通信任务的负载均衡和调度

问题形式化：由于主机是主要的通信瓶颈，我们在主机级别而不是单个设备级别进行负载均衡和调度。假设一个单元通信任务在主机集$n_i \subseteq MeshA$中有副本$DS_i$，我们想将其发送到主机集$m_i \subseteq MeshB$。我们的目标是为每个任务选择一个发送主机$n_i^* \in n_i$，并安排不同任务的执行顺序。假设第i个任务的持续时间为$T_i$，执行开始时间为$S_i$，我们可以将负载均衡和调度问题形式化为以下优化问题：

优化目标（公式1）是最小化最后一个单元通信任务的完成时间。约束条件（公式2）确保每个任务的发送方拥有正确的数据切片，约束条件（公式3）确保共享相同发送方或接收方设备的任务不会相互重叠。

朴素算法（Naive algorithm）：最直接的算法是为每个任务分配索引最低的设备作为发送方，并按照任意的全局顺序调度所有任务。我们用此算法作为基线。

仅负载均衡（Load balance only）：我们首先只考虑跨多个发送主机的负载均衡。在这种情况下，原问题简化为：


这个极小化极大问题可以通过一个经典的贪心算法来最优解决：我们将所有任务按持续时间$T_i$降序排序，然后遍历排序后的任务，将每个任务分配给当前工作负载最轻的发送主机。

带剪枝的深度优先搜索（DFS with Pruning）：为了同时考虑调度，我们首先注意到，调度问题（即决定$S_i$）可以简化为：为每个主机上的所有发送/接收任务分配一个执行顺序。然后，每个任务的开始时间$S_i$可以设置为在其发送主机$n_i^*$和接收主机$m_i$上所有前序任务都完成的最早时间。我们提出一个基于搜索的算法：构建一个关于发送方分配和任务顺序两种决策的深度优先搜索树。当算法到达一个叶节点时，我们计算其完成时间。在搜索过程中，我们剪掉那些执行时间下限（由单个设备上的任务持续时间总和衡量，如公式4）超过当前最优完成时间的分支。该算法计算复杂度为指数级，在实践中，我们设置时间预算，当单元通信任务超过20个时，该算法无法在预算内产生有效调度。

带随机化的贪心搜索（Greedy Search with Randomization）：我们提出一种迭代贪心算法，在每次迭代中，选择一个包含最大设备数量的非重叠任务集。为了找到这样的集合，我们使用随机化算法：首先，生成一个待调度任务的随机排序。然后，按此顺序遍历每个任务，并选择与之前已选任务不重叠的任务。我们重复此过程多次，并选择迭代中最大的候选集。由于跨设备网格重分片任务中许多单元通信任务是相同的，并且任务均匀分布在不同设备上，因此该贪心随机算法可以在短时间内找到良好甚至最优的解决方案。

4 对重叠友好的流水线调度

动机：通过实现一个假设性的上限“Signal Send/Recv”（只通信一个字节以保持数据依赖性），我们发现训练吞吐量可以提高1.1倍到1.5倍。这表明仍然有机会通过重叠来减少通信成本。我们的关键观察是，现有的同步流水线调度，如1F1B【23, Efficient large-scale language model training on gpu clusters using megatronlm, 2021, Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】和GPipe【10, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019, Advances in neural information processing systems】，对重叠并不友好。如图4(a)所示，在1F1B中，同一微批次的计算和通信任务是连续调度的，由于严格的数据依赖性，通信无法被重叠。

Eager-1F1B调度：为解决此问题，我们开发了一种新的、更对重叠友好的调度，称为eager-1F1B。其主要思想是积极地运行计算，为重叠创造机会。如图4(b)所示，eager-1F1B调度基于1F1B，但将前向计算任务提前了几个时间步。在eager-1F1B的预热阶段，阶段i会运行更多微批次的前向计算（具体为 (2×(#stages−i)+1) 个），然后再进入稳定的“一前一后”阶段。当通信成本不可忽略时，1F1B无法隐藏该成本，而eager-1F1B调度在两个依赖任务之间插入了其他微批次的计算任务，从而可以隐藏通信成本。


图4. 1F1B和eager-1F1B调度的时间线。每一行是一个流水线阶段。数字表示微批次索引。微批次4从阶段2到阶段3的通信被放大显示。

内存开销：eager-1F1B调度的缺点是内存使用略有增加，因为每个GPU为更多的微批次存储临时激活张量。然而，在通常情况下，这种增加非常小。如表1所示，新的调度使每个GPU的内存使用最多增加#stages × size_activation，这远小于参数大小。

Table 1. 混合精度训练中一个GPT-3层的参数、优化器状态和每个GPU的激活大小。序列长度S=1024，隐藏层大小H=12288，每个GPU的微批次大小B=2，张量模型并行度TMP=8。

后向传播优化：eager-1F1B调度可以重叠前向部分的通信，但对后向部分没有帮助。为了重叠后向通信，我们将后向计算分为两部分：计算激活的梯度和计算权重的梯度。第一部分不依赖于第二部分，而跨设备网格通信仅使用第一部分的输出。因此，我们延迟第二部分的计算，使其与通信重叠。我们称之为后向权重延迟（backward weight delaying）。类似地，这项技术会略微增加峰值内存使用，因此我们使用一个简单的成本模型来估计计算和通信时间，并进行最小程度的延迟以覆盖所有通信。

A4 实验

实验环境

Table 3. 端到端评估中的模型。并行配置是一个元组（数据并行度，算子并行度，流水线并行度）。

实验结果

5.1 微基准测试

5.1.1 单设备到多设备


图5. 单设备到多设备微基准测试结果。

5.1.2 多设备到多设备

Table 2. 多对多设备微基准测试中发送方和接收方的分片规范。


图6. 多对多设备微基准测试结果。

5.2 端到端性能


图7. 端到端评估结果。

5.3 消融研究

5.3.1 负载均衡


图8. 负载均衡算法的消融研究。

5.3.2 重叠


图9. 重叠友好调度的消融研究。

A5 结论

本文介绍了用于加速跨设备网格重分片的通信优化，这是一种在大型神经网络分布式训练中新颖且重要的通信模式。我们在三个粒度上优化了该问题，包括单个单元通信任务、一个跨设备网格重分片任务中所有单元任务的负载均衡，以及将所有跨设备网格重分片任务与计算进行重叠。我们将一个先进的分布式训练框架Alpa移植到了我们的工作之上。我们的工作将大型模型的训练速度提高了高达1.5倍，并取得了非常接近假设性上限的性能。

A6 附录

B 分解为单元任务

B.1 一个跨设备网格重分片的所有单元任务示例

图10和图11分别列出了从图2中的跨设备网格重分片任务1和任务2分解出的所有单元任务的示例。


任务1: $S^{01}R$, MeshA → $S^0R$, MeshB
图10. 跨设备网格重分片案例1的单元任务。


图11. 跨设备网格重分片案例2的单元任务。

B.2 分解跨设备网格重分片的算法

分解步骤：我们将一个跨设备网格重分片任务通过两步分解为单元任务：

  1. 创建切片：首先，我们根据发送方和接收方网格的分片规范将张量分割成切片。
  2. 创建单元任务：然后，对于每个切片，我们创建一个单元任务，其发送方是所有拥有该切片的发送设备，接收方是所有需要该切片的接收设备。

切片创建细节：为了创建切片，我们遍历张量的所有维度。每个维度d有一组切点，记为$Cut(d)$，初始值为$\{0, l_d\}$，其中$l_d$是维度d的长度。如果维度d在发送方网格上以k度分片，该维度被分成k个等大的副本，因此我们将$l_d/k, 2l_d/k, \dots$加入切点集$Cut_d$。同样的操作也适用于接收方维度的分片规范。最终，对于一个n维张量，我们得到n个对应的切点集。对于每个切点集$Cut_d = \{p_0, p_1, \dots, p_{c-1}\}$，我们创建对应的区间集$I_d = \{(p_0, p_1), (p_1, p_2), \dots, (p_{c-2}, p_{c-1})\}$。最终的切片集由所有张量维度的区间集的笛卡尔积表示，即$I_0 \times I_1 \times \dots \times I_{d-1}$。