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

A1 主要贡献

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

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

研究目标：本文将跨设备网格重分片问题形式化为一个多对多多播（many-to-many multicast）通信问题，并揭示了在模型并行深度学习背景下的几个独特挑战：
- 复杂的通信需求：跨设备网格重分片中传输的消息是一个分布式张量，它在两个独立的网格上分片。这些张量在源和目标网格上可能表现出不同的布局，这要求在最小化通信成本的同时，仔细选择通信原语、发送和接收设备以及通信路径。
- 异构网络：深度学习集群包含异构网络硬件（如NVLink, PICe, NIC, InfiniBand），导致设备在网格内和网格间通信时带宽不均。因此，需要仔细分配通信任务以充分利用不均衡的带宽并平衡网络流量。
- 与计算调度的依赖关系：一个完整的模型并行作业包含许多跨设备网格重分片任务，这些任务的触发依赖于计算依赖。为了提高模型的整体性能，需要对这些计算和通信任务进行最优的调度和重叠。

创新点与主要贡献：
为了应对这些挑战，本文提出了新的解决方案，以优化单个任意的跨设备网格重分片任务，并集体优化模型并行作业中的所有跨设备网格重分片任务。主要贡献总结如下：
* 形式化问题：形式化了模型并行深度学习作业中的跨设备网格重分片问题，并揭示了其特性。
* 优化单个重分片任务：提出了一种通用且有原则的解决方案来优化单个跨设备网格重分片任务。具体而言，首先证明了现有解决方案中的通信原语是次优的，然后提出使用广播（broadcast）作为替代方案，并证明其在理论上是最优的。基于此，开发了新算法来调度和负载均衡单个跨设备网格重分片任务中的多个广播原语。
* 新的流水线调度：提出了一种新的流水线调度方案，该方案可以将模型并行作业中多个跨设备网格重分片通信任务与流水线并行计算进行重叠。
* 实现与评估：将所提出的技术实现为一个名为 AlpaComm 的通信库，并在各种微基准测试中展示了其比现有解决方案快高达10倍的性能。
* 端到端性能提升：将该库与先进的模型并行系统 Alpa 集成。本文的技术将两个模型 GPT 和 U-Transformer 的端到端多节点训练吞吐量分别提高了10%和50%。


图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集群设定：在深入细节之前，我们首先考虑具有以下特性的典型机器学习集群设置：
* 快速的节点内通信和慢速的节点间通信：例如，GPU集群中，节点内的GPU通过高速NVLink互连，而跨节点的GPU需要通过较慢的以太网或InfiniBand连接通信。
* 节点间全连接拓扑：我们假设一对节点之间的网络带宽不会受到其他节点间通信的影响，并且任何节点对都具有相同的带宽。
* 主机通信瓶颈：设备（如GPU）的带宽通常远高于其主机的带宽。当单个主机中的多个设备向另一主机发送数据时，它们会竞争主机的网络接口带宽。我们假设每个主机只有一个NIC。
* 独立的发送/接收带宽（全双工）：我们假设每个设备具有独立的发送和接收带宽，可以同时以全带宽发送和接收数据。

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 实验

实验环境

  • 实现：本文的优化在一个名为 AlpaComm 的通信库中实现，包含约900行C++和2.5k行Python代码。该库基于NCCL【24, The nvidia collective communication library, 2018, URL: https://developer.nvidia.com/nccl】为GPU实现，并已集成到Alpa【45, Alpa: Automating inter-and intra-operator parallelism for distributed deep learning, 2022, arXiv preprint】中。eager-1F1B调度和重叠优化在Alpa之上实现。
  • 硬件配置：实验在AWS集群上运行，每个节点是p3.8xlarge实例，配备4个NVIDIA V100（16GB）GPU和32个vCPU。节点内的GPU通过NVLink连接，节点间带宽为10Gbps。
  • 模型：端到端评估使用了两个模型：
    • GPT：一个类似GPT-3的语言模型，由结构相同的Transformer层堆叠而成。
    • U-Transformer：一个U形的卷积神经网络，带有长跳跃连接，并在卷积层后插入注意力层。
    • 具体模型配置和并行策略见下表。

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

实验结果

5.1 微基准测试

5.1.1 单设备到多设备
* 实验内容：发送方网格只有一个GPU，接收方网格的GPU数量变化（同一节点内1-4个GPU，或固定每节点2个GPU，节点数1-4个）。消息大小为1GB。
* 基线：
1. "Send/Recv"：仅使用点对点send/recv原语。
2. "Alpa"：基于all-gather的方法，用于Alpa和Megatron-LM等系统。
* 实验结果与分析：如图5所示，当接收方在单个节点时，"Send/Recv"的延迟随GPU数量线性增长，而我们的AlpaComm和Alpa的延迟增长不到1%。当接收方跨越多个节点时，Alpa中的all-gather必须使用节点间连接，开销显著增加。而AlpaComm几乎没有性能下降，因为它能有效处理分块、填充和流水线化，并且能处理Alpa无法处理的非均匀分区情况。


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

5.1.2 多设备到多设备
* 实验内容：发送方和接收方网格都包含多个节点，测试了多种来自常见深度学习工作负载的分片规范（见表2）。张量形状为(1024, 1024, 512)。
* 基线：与5.1.1相同，但基线也使用贪心方法进行负载均衡。
* 实验结果与分析：如图6所示，在多种情况下，AlpaComm表现与Alpa相似或更好。在案例7和8中，Alpa的跨节点all-gather很慢，而AlpaComm通过流水线化，比Alpa快2.5倍。在案例3、4和9中，AlpaComm比Alpa快3到10倍，因为它通过重新排序通信，充分利用了发送方两个节点的带宽，而Alpa的调度会导致一个发送节点空闲。

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


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

5.2 端到端性能

  • 实验内容：在GPT和U-Transformer模型上进行端到端训练性能评估。
  • 基线：
    1. "Send/Recv"
    2. "Alpa"
    3. "Broadcast"：使用基于广播的重分片，但没有第4节介绍的重叠优化。
    4. "Signal Send/Recv"：作为性能上限。
  • 实验结果与分析：如图7所示，在GPT模型上，AlpaComm和Alpa都接近上限，因为它们能将通信卸载到高带宽的NVLink。AlpaComm由于重叠优化，实现了1.1倍的加速。在U-Transformer模型上，U形跳跃连接使跨网格通信成为瓶颈，我们的eager 1F1B调度能有效重叠通信，与Alpa相比实现了1.5倍的加速。在所有情况下，AlpaComm都达到了上限性能的75%以上。


图7. 端到端评估结果。

5.3 消融研究

5.3.1 负载均衡
* 实验内容：在多对多微基准测试上评估负载均衡算法的效果。
* 基线：
1. "Naive algorithm"：选择索引最低的发送方。
2. "Load balance only"：仅平衡每个节点发送的数据量。
3. AlpaComm：结合了DFS剪枝和随机化贪心搜索。
* 实验结果与分析：如图8所示，在大多数情况下，"Naive"和"Load balance only"都会遇到拥塞，而AlpaComm能找到最优的顺序，避免任何发送或接收节点空闲。


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

5.3.2 重叠
* 实验内容：在U-Transformer模型上评估重叠友好调度的效果。
* 基线：
1. "Broadcast"：无重叠。
2. "Overlap"：在Broadcast基础上增加了通信和计算的重叠，但未使用eager-1F1B。
3. "Eager-1F1B"：在"Overlap"基础上使用eager-1F1B调度。
* 实验结果与分析：如图9所示，对于典型情况（微批次数量较大），"Overlap"比"Broadcast"快30%，而"Eager-1F1B"在此基础上又增加了15%的性能提升。


图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}$。