Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models

发表时间: 2023-12 · ASPLOS'23 (Google)

作者/机构: Shibo Wang, Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, Sameer Kumar, Tongfei Guo, Yuanzhong Xu, and Zongwei Zhou

A1 主要贡献

核心问题与研究目标

大型深度学习模型因其庞大的体积，远超单个加速器（如GPU或TPU）的内存容量，因此必须采用模型并行技术才能运行。其中，层内（张量）模型并行是一种常用技术，它将模型的单个层或操作符划分到分布式加速器集群的多个设备上。然而，这种并行方式会产生大量跨设备的数据通信，这些通信开销在总执行时间中占比巨大，严重损害了计算效率。图1展示了多个大型模型在TPU v4 Pods上训练时，通信时间占据了相当大的比例。本文的目标是提出一种新技术，通过将通信与计算重叠执行，来有效隐藏层内模型并行带来的数据通信开销。


图1：大型模型训练的单步时间分解。这些模型根据其大小在128到2048个TPU芯片上运行。大型模型的详细信息见第6节。

创新点与核心贡献

本文提出了一种通过分解（Decomposition）来重叠依赖性计算与通信的新技术，旨在解决上述问题。其核心思想是，不再等待所有数据准备就绪后才开始计算，而是将原始的粗粒度通信集合操作（如AllGather）和其依赖的计算操作（如Einsum）分解为一连串更细粒度的操作。通过这种方式，可以创造出更多的重叠机会，并以迭代的方式并行执行这些新生成的细粒度通信和计算。具体贡献如下：

A3 背景知识

2.1 通信原语

MPI风格的集合操作：在单程序多数据（SPMD）范式下，层内模型并行中发生的通信使用MPI风格的集合操作【【21】，MPI: A Message-Passing Interface Standard，1994】来表示。一些常见的集合操作包括：
* CollectivePermute：在每个指定的源-目标对之间进行点对点通信。
* AllGather：一种多对多的通信模式，每个逻辑设备分区从所有分区接收数据，并根据数据分区ID进行拼接。主要用于在已分区的维度上复制数据。
* ReduceScatter：与AllGather的通信模式相反。每个逻辑设备分区对接收到的具有相同分区ID的数据分片执行逐元素归约。
* AllReduce：可以看作是一个ReduceScatter后跟一个AllGather。操作完成后，每个逻辑设备分区都拥有对所有输入数据进行逐元素归约后得到的相同张量值。

2.2 典型分区策略

单维度分区策略：许多大型机器学习模型可以通过仅沿张量的一个维度进行分区来适应内存并获得高系统吞吐量。我们使用的第一种常见分区策略类似于Megatron【【46】，Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism，2019】所使用的策略。但不同之处在于，我们不在层输入和输出上保持数据复制，也不在分区计算后沿归约维度执行AllReduce，而是在每个设备上保留激活值和形状的一个分片，并在执行期间通过集合操作按需构建权重。

图2展示了一个双层多层感知器（MLP）示例中的数据通信模式。在此示例中，每个设备在整个过程中都保留了激活值的一个分片（形状为[B/N, F]和[B/N, H]）。在每次Einsum计算之前，会执行一次AllGather操作来重构原始的多维权重（[F, H]和[H, F]）。在反向传播过程中（图中未显示），AllGather将变为ReduceScatter，以生成用于权重更新的分区梯度。


图2：一个双层MLP在单维度上进行N路分区的前向传播。图中，每个张量由一个包含两个元素的列表表示，B、F和H分别代表批次、特征和隐藏维度的大小。注意，激活值和权重已经被分区，由某个维度上的“/N”表示。

双维度分区策略：当模型规模扩大到非常大时，由于设备内存容量的限制，需要更多的分区数量。在这种情况下，设备被视为形成一个逻辑网格或环面，而不是一个环。张量会沿两个维度进行分区，以避免分区后维度尺寸过小导致在GPU或TPU上计算效率低下。我们采用一种不同的分区策略，以实现低峰值内存使用和高有效计算效率。

图3展示了该双层MLP示例在双维度分区下的新数据通信模式。与图2中的示例类似，激活值的批次维度保持分区状态。在每次Einsum之前，输入激活值和权重会沿不同维度进行AllGather。此外，为了减少内存使用，需要长时间存在的输出激活值必须保持完全分区状态——在第二次Einsum的部分分区结果上，会沿x维度执行一次子组ReduceScatter。


图3：一个双层MLP在两个维度上进行N*M路分区的前向传播。在此示例中，设备网格或环面的形状为[M, N]。张量维度除以M或N意味着张量分别沿x或y维度分区。每个通信集合操作末尾的(x)或(y)注释代表相应的数据传输维度。

3 相关工作

4 概览

核心思想：由于层内模型并行所需的通信集合操作非常耗时且会阻塞后续的依赖计算，本文提出通过将通信与依赖的计算重叠来更好地利用空闲的计算单元。其关键思想是，对每个数据分片分别进行迭代式地部分计算和异步、非阻塞的数据通信，而不是将它们作为一个整体来处理。

分解与重叠过程：为了实现这一点，一个目标的多步、阻塞式集合操作（如AllGather）被分解为一系列单步、非阻塞的集合操作（CollectivePermutes）。该集合操作希望与之重叠的计算操作（如Einsum）也被分解为一系列更细粒度的计算操作，并通过一个简单的累加或拼接操作来组合部分结果。这两个序列在语义上等价于它们各自的原始操作，并且可以构成一个具有相同迭代次数的循环。本文技术不是按顺序执行这两个序列，而是交错执行通信和计算的迭代。由于序列的生成方式保证了在同一次迭代中通信和计算操作之间没有数据依赖，因此每个异步、非阻塞的集合操作都被调度与相应的分解计算并行执行。

AllGather-Einsum重叠示例：图4展示了该技术如何应用于AllGather的情况。在当前系统中，所有设备必须等待AllGather操作完成，获得完整的操作数A（即[A0, A1]）后才能开始计算。而本文提出的重叠方案则不同，它不等待所有数据都准备好。每个设备从异步地将其本地持有的数据分片发送给其他设备开始，同时使用当前可用的数据分片进行部分计算。当第一个部分结果产生时，第二次计算所需的数据分片（分区0上的A1；分区1上的A0）已经接收完毕。因此，第二次部分计算可以立即继续。每个部分结果需要一个额外的DynamicUpdateSlice操作来更新最终结果的相应切片。通过分解和重叠获取A所需的数据通信与相应的计算，该方案能够实现更好的硬件利用率和性能。


图4：一个关于2路层内模型并行的AllGather-Einsum关键思想的说明性示例。上方的“原始”部分说明了操作在当前系统中的执行方式；下方的“重叠”部分说明了所提出技术的工作方式。图中每个操作的延迟并不精确，仅用于说明思想。

Einsum-ReduceScatter重叠示例：类似地，ReduceScatter也可以与相应的计算重叠，如图5所示。在此示例中，通信作用于计算结果。每个设备需要异步传输累积的结果分片，而不是操作数分片。累积结果分片在每个设备上初始化为零。在每次迭代开始时，每个设备异步地将累积结果分片发送给其他设备，并与部分Einsum计算并行执行。当计算完成时，部分结果会加到迭代结束时收到的累积结果分片上。值得注意的是，用于对操作数A进行DynamicSlice的索引并非从与设备分区ID相同的数字开始。这样设计是为了使最终的结果分片ID与其对应的设备分区ID对齐。


图5：一个关于2路层内模型并行的Einsum-ReduceScatter关键思想的说明性示例。

5 方法细节

该技术作为编译器遍（compiler pass）实现，可为采用层内模型并行的大型深度学习模型自动工作。生成细粒度通信和计算操作序列的分解部分作为数据流图转换添加。然后，一个解耦的异步、非阻塞CollectivePermutes生成和调度遍被添加以执行重叠。最后，添加了优化措施和一个基于预估收益自动启用该功能的机制，以进一步提高系统性能。

5.1 Looped CollectiveEinsum

循环重构：遵循第4节的思想，将Einsum计算与相应的通信集合操作分解为具有重复模式的操作序列，并重构为一个循环，我们称之为Looped CollectiveEinsum。循环的每次迭代对不同数据分片执行部分计算和相关的数据传输，数据分片ID根据循环索引变量计算得出。

AllGather-Einsum重叠的三种情况：在图4的AllGather-Einsum示例中，操作数A是沿非收缩维度分区的。然而，实际的分区策略可能会选择不同的逻辑维度以获得更好的系统性能。沿不同逻辑维度分区会导致每次迭代中对部分计算的输入和输出有不同的处理方式。假设输出的分区策略与右侧操作数（RHS）兼容，重叠AllGathers的三种不同情况如下：
* 情况1 - LHS沿非收缩维度分区：这是最简单的情况（与图4所示相同）。每次迭代中，分解后的Einsum直接使用接收到的左侧操作数（LHS）和RHS进行计算，然后需要一个DynamicUpdateSlice来更新最终结果。
* 情况2 - LHS沿收缩维度分区：在这种情况下，RHS需要先进行DynamicSlice才能送入Einsum。DynamicSlice沿RHS的收缩维度执行，切片大小与LHS相应分区维度的大小相同。前一种情况中Einsum之后的DynamicUpdateSlice被替换为Addition以累积部分结果。
* 情况3 - LHS沿批处理维度分区：与情况2类似，RHS也需要DynamicSlice来获取正确的计算数据分片。然而，切片是沿批处理维度进行的，大小与RHS的批处理维度相同。Einsum计算之后是一个沿批处理维度的DynamicUpdateSlice来更新最终结果。

Einsum-ReduceScatter重叠：当LHS或RHS沿收缩维度分区，并且结果沿非收缩维度分区时，执行Einsum-ReduceScatter重叠。检查LHS和RHS，看哪一个具有与输出的分区非收缩维度相对应的非收缩维度。找到的操作数然后沿匹配的非收缩维度进行DynamicSlice，以获得Einsum的正确输入分片。结果更新操作是在每次迭代结束时对相应的输出分片进行Addition。

通用算法与数据传输：算法1展示了Looped CollectiveEinsum对AllGather和ReduceScatter情况的通用伪代码。CollectivePermutes用于在设备分区之间交换数据分片（对于AllGather情况是操作数，对于ReduceScatter情况是Einsum的结果）。为了模块化和更好的代码生成，该技术将异步、非阻塞CollectivePermutes的生成和调度与分解和循环生成解耦。如算法1所示，在循环生成期间使用默认的同步、阻塞CollectivePermutes。后续的编译器遍会将其转换为异步、非阻塞的，并调度它们与相应的Einsum重叠。每次迭代中CollectivePermute的{source, destination}对被构造为{0, N - 1}, {1, 0}, {2, 1}, ..., {N - 1, N - 2}。图6和图7分别说明了在AllGather和ReduceScatter情况下，数据分片如何在每次迭代中在设备分区之间传输。数据分片在两种情况下都是循环左移，但模式略有不同：在AllGather情况下，数据分片ID在循环开始时与设备分区ID对齐；而在ReduceScatter情况下，相同的模式出现在循环结束时。这是因为在AllGather情况下，数据分片既是Einsum的操作数也是CollectivePermute的操作数，但在ReduceScatter情况下，它在CollectivePermute之后用于累积部分结果。


算法1：Looped CollectiveEinsum的通用简化算法。looped_operand输入指代将在设备分区之间传输的已初始化数据分片（例如图4中的Ax）。local_operand输入指代不会被传输的数据（例如图4中的Bx）。Update可以是Addition或DynamicUpdateSlice，具体取决于实际情况。Select并非实现中实际使用的操作，仅用于表示在两个可能的输入中进行选择。

图6：一个在4路层内模型并行下，循环AllGather-Einsum中数据分片传输的说明性示例。

图7：一个在4路层内模型并行下，循环Einsum-ReduceScatter中数据分片传输的说明性示例。

5.2 异步CollectivePermute与调度

异步非阻塞行为的实现：为了实现异步、非阻塞行为，每个CollectivePermute被分解为一对CollectivePermuteStart和CollectivePermuteDone。CollectivePermuteStart仅启动源、目标设备对之间的数据传输，它不阻塞后续指令，并且几乎不花费执行时间。CollectivePermuteDone用于指示数据传输的完成。尽管实际的数据传输时间没有改变，但观察到的执行时间根据数据传输可以与计算重叠的程度而变化。为了最大化重叠性能，异步CollectivePermuteStart和CollectivePermuteDone在指令序列中的调度至关重要。

调度算法：为避免大幅改变变量的生命周期，所提出的调度算法以现有指令调度遍（该遍使用旨在最小化内存使用的算法）调度后的指令作为输入。
* 自底向上调度方法：该方法采用自底向上的方式——指令从数据流图的根节点开始反向调度。其简化伪代码如算法2所示。ready_queue保存准备好调度的指令，并优先处理CollectivePermuteDones及其用户。pending_queue保存所有用户指令都已调度但自身尚未准备好调度的指令。指令的就绪性通过比较其预估就绪时间与当前时间来确定。调度启发式算法会尝试在ready_queue和pending_queue中找到下一个要调度的指令。如果选择的指令是CollectivePermuteDone且进行中的异步指令总数超过预算（受系统中同步标志数量的限制），则会选择队列中的下一个指令。
* 自顶向下调度方法：另一种方法使用自顶向下的方法，其思想简单而有效——在操作数指令被调度后，尽早调度CollectivePermuteStart；在其被使用前，尽可能晚地调度CollectivePermuteDone。为了解决原始输入指令顺序可能显著减少重叠机会的问题，该方案首先根据运行时成本在每个CollectivePermute区间内重新平衡指令。
方法选择：两种方法都进行了评估比较。自底向上的方法比自顶向下的方法更通用，但实现更复杂。由于其性能稍好，最终评估采用了自底向上的方法。


算法2：一种用于隐藏通信开销的自底向上调度算法。为简化起见，忽略了每个辅助函数的细节。

5.3 在XLA编译器中的实现

XLA中的实现：上述的分解和重叠方案都在XLA中实现，XLA是专为线性代数设计的领域特定生产编译器【【3】，XLA: Optimizing Compiler for TensorFlow，2021】。它将机器学习模型转换为后端无关的HLO图，这是一个将操作表示为节点、张量表示为边的数据流图。许多前端框架，如Tensorflow【【7】，TensorFlow: A System for Large-Scale Machine Learning，2016，OSDI '16】、JAX【【11】，JAX: composable transformations of Python+NumPy programs，2018】、PyTorch【【38】，PyTorch: An Imperative Style, High-Performance Deep Learning Library，2019，NeurIPS】和Julia【【10】，Julia: A Fresh Approach to Numerical Computing，2014】，已经具备了获取这种中间表示的逻辑。

普适性：在XLA中实现后，所提出的技术可以自动为上述支持的前端和后端执行图转换和通信重叠。然而，该技术不限于XLA，其思想可以应用于其他使用类似通信原语来运行大规模深度学习模型的机器学习系统和框架。

5.4 性能优化

5.4.1 循环展开（Loop Unrolling）：直接实现算法1中的Looped CollectiveEinsum会在循环中产生额外的、不可忽略的Copy操作，这是由于传输数据的循环携带别名问题。为了消除显式的复制开销，我们添加了展开度为2的循环展开，以提供类似于双缓冲的效果。这对Einsum-ReduceScatter情况有额外的好处。如果没有循环展开，异步的CollectivePermuteDone必须放弃与大的Einsum重叠，因为结果更新的累加操作与Einsum融合，从而在CollectivePermuteDone和融合节点之间增加了数据依赖。图8说明了如何通过循环展开来提高循环Einsum-ReduceScatter的性能。通过将累积链分成两个交错的部分，两条独立的路径使得即使Einsum与累加操作融合，异步CollectivePermuteDone也能成功地与另一路径的Einsum重叠。为了保证正确性，需要添加一个额外的epilogue过程（图中未显示），其开销与循环相比很小。


图8：一个在4路层内模型并行下，展开度为2的循环Einsum-ReduceScatter的说明性示例。与图7类似，图中的Dx块表示在相应设备分区上用于该特定迭代计算的ID为x的数据分片。

5.4.2 双向数据传输（Bidirectional Data Transfer）：随着分区数量的增加，每个分区Einsum内的计算量可能不足以与相应的通信重叠。在这种情况下，采用双向数据传输来缓解这个问题。通过将两个部分Einsum的操作数拼接起来并作为单个操作执行，它有效地使每次迭代的计算量加倍。此外，双向数据传输可以充分利用TPU上提供双向高带宽的专用芯片间互连（ICI）链路【【26】，A Domain-Specific Supercomputer for Training Deep Neural Networks，2020，Commun. ACM】、【【36】，The Design Process for Google’s Training Chips: TPUv2 and TPUv3，2021，IEEE Micro】，这有助于普遍提高所提重叠技术的系统性能。图9和图10分别说明了AllGather-Einsum和Einsum-ReduceScatter重叠情况下的数据通信模式。在AllGather-Einsum情况下，循环前需要一个prologue过程来初始化数据分片。在Einsum-ReduceScatter情况下，则需要一个epilogue来对齐结果分片并执行求和。


图9：一个在4路层内模型并行下，采用双向数据通信的循环AllGather-Einsum中数据分片传输的说明性示例。

图10：一个在4路层内模型并行下，采用双向数据通信的循环Einsum-ReduceScatter中数据分片传输的说明性示例。

5.4.3 融合优化（Fusion Optimization）：为了让更多内存绑定的逐元素操作能够融合到生成的CollectiveEinsum循环中，我们执行了一些局部的图重写使其对融合更友好。例如，将Concatenation(a, b)替换为语义上等价的操作集合Max(PadLow(a), PadHigh(b))。此外，通过交换Reshape和Concatenation或Slice的顺序，可以融合更多的操作。同时，我们还调整了融合启发式策略以避免有害的融合。如图11所示，默认的融合（a）会创建一个从Fusion_0到CollectivePermuteDone的数据依赖，从而阻止重叠。通过改变融合启发式，优先将Addition与依赖CollectivePermuteDone的Einsum融合（b），数据通信就可以成功地与Fusion_0重叠。


图11：具有不同融合决策的HLO图示例。

5.5 基于预估收益启用重叠特性

成本模型与启用机制：在该方案中，数据沿着由设备分区形成的逻辑环异步传输。当原始的AllGather或ReduceScatter在网格或环面上执行时，分解后的CollectivePermutes序列的总运行时间可能比原始的更长，因为它只利用了一半的互连带宽（沿单个维度）。为了防止性能下降，我们增加了一个通用机制，仅在预估有收益时才为每个AllGather-Einsum或Einsum-ReduceScatter启用重叠特性。预估基于以下公式：$original_comp + original_comm >= max(original_comp, ring_comm) + epilogue_cost$。其中，$original_comp$和$original_comm$分别代表原始计算和通信的时间。$ring_comm$代表沿逻辑环的通信时间。$epilogue_cost$是prologue或epilogue的开销。如果上述公式不满足，则保持原始操作不变。

多候选选择：当一个Einsum有两个AllGather（或一个AllGather和一个ReduceScatter）候选可以重叠时，方案会选择能带来更高收益的一个。如果Einsum预估比两个通信集合操作都快，则选择数据分片尺寸较小的那个，因为它在循环外的额外CollectivePermute开销较小。否则，选择预估执行时间较长的集合操作进行分解和重叠。

6 实验环境与结果

实验环境

表1：评估的应用。

表2：评估的GPT模型，参数从320亿扩展到1万亿。

实验结果


图12：评估应用的性能。


图13：弱扩展GPT模型的性能。


图14：循环展开带来的性能改进。

图15：双向数据传输带来的性能改进。

图16：两种调度方法的性能比较。

7 补充细节

7.1 在推理任务中的应用

适用性：尽管该重叠方案主要在训练任务上进行了评估以展示其对系统吞吐量的提升，但它同样可以应用于推理任务，以有效隐藏通信延迟。

示例与未来工作：例如，一个内部的推荐推理模型，采用2路层内模型并行，通过所提出的重叠技术实现了2倍的延迟改进。对不同类型推理工作负载的全面研究将作为未来的工作。

7.2 在其他模型和系统中的应用

模型扩展性：本文评估的模型主要是自然语言处理（NLP）和语音模型。对于需要模型并行来适应加速器内存的新兴的基于多层感知器（MLP）或基于注意力的计算机视觉模型【【54】，Scaling Vision Transformers，2021，arXiv】以及图像生成模型【【42】，Zero-Shot Text-to-Image Generation，2021，arXiv】，该技术同样适用。

系统与框架扩展性：尽管该技术在XLA中实现和评估，但其思想也可以应用于其他编译器（如TVM【【14】，TVM: End-to-End Optimization Stack for Deep Learning，2018，arXiv】、Glow【【44】，Glow: Graph Lowering Compiler Techniques for Neural Networks，2018，arXiv】）或ML框架（如PyTorch【【38】，PyTorch: An Imperative Style, High-Performance Deep Learning Library，2019，NeurIPS】、Mesh-Tensorflow【【45】，Mesh-TensorFlow: Deep Learning for Supercomputers，2018，CoRR】），前提是它们使用类似的通信原语来运行大规模深度学习模型。

硬件扩展性：同样，该重叠思想可以应用于其他硬件ML系统，例如通过高带宽、低延迟NVLink网络互连【【4】，NVIDIA H100 Tensor Core GPU Architecture，2022】连接的GPU集群。对于采用性能较低互连的系统，由于其极长的通信时间无法被并发计算所覆盖，所提技术的收益将会减少。

7.3 其他模型并行模式

与流水线并行的结合：流水线模型并行【【24】，GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism，2018，CoRR】、【【29】，Pipelined Backpropagation at Scale: Training Large Models without Batches，2020，CoRR】、【【33】，PipeDream: Generalized Pipeline Parallelism for DNN Training，2019，SOSP '19】、【【34】，Memory-Efficient Pipeline-Parallel DNN Training，2020，CoRR】、【【53】，PipeMare: Asynchronous Pipeline Parallel DNN Training，2019，CoRR】是另一种为大型深度学习模型提出的扩展技术。一些近期工作（例如，【【35】，Efficient Large-Scale Language Model Training on GPU Clusters，2021，CoRR】）将层内模型并行和流水线模型并行与数据并行相结合，以在GPU集群上实现高效训练。

新的优化机会：在这些情况下，所提出的重叠技术仍然可以应用，以进一步提高系统性能。通过降低层内模型并行引入的通信成本，该技术改变了不同类型并行之间的性能权衡。这也为未来寻找更好的并行组合提供了新的优化机会。

8 结论

在采用层内模型并行运行大型深度学习模型时，当前系统会花费大量时间等待昂贵的通信集合操作为后续的依赖计算准备数据。本文提出了一种新颖的技术来解决这个问题。通过将分解后的非阻塞、点对点通信集合操作与相应的计算操作重叠，该技术可以显著降低通信成本并提高系统吞吐量。尽管性能在TPU系统上进行了评估，但其思想同样可以应用于其他分布式加速器系统（例如GPU集群）。我们得出结论，所提出的技术有潜力在未来的大型深度学习系统中有效提高FLOPS利用率。

方法细节中的引用汇总