Unity: Accelerating DNN Training Through Joint Optimization of Algebraic Transformations and Parallelization

文章标题: Unity: 通过代数变换和并行化的联合优化加速DNN训练
作者/机构: Colin Unger†♠, Zhihao Jia‡♭♠, Wei Wu∗⋄, Sina Lin§, Mandeep Baines♭, Carlos Efrain Quintero Narvaez♭, Vinay Ramakrishnaiah∗, Nirmal Prajapatii∗, Pat McCormick∗, Jamaludin Mohd-Yusof∗, Xi Luo♯, Dheevatsa Mudigere♭, Jongsoo Park♭, Misha Smelyanskiy♭, Alex Aiken†
所属机构: Stanford University†, Carnegie Mellon University‡, Los Alamos National Lab∗, NVIDIA⋄, Microsoft§, Meta♭, SLAC National Accelerator Laboratory♯

---

A1 主要贡献

本文介绍了Unity，这是第一个在分布式DNN训练中联合优化代数变换和并行化的系统。现有的DNN训练优化方法主要分为两类：代数变换和并行化，但它们通常是独立应用的，这会错失重要的性能提升机会。

核心问题：独立地、顺序地应用代数变换和并行化优化是次优的。例如，一个代数优化器可能会融合两个算子以提高计算效率，但这可能会阻止后续并行化优化器应用更高效的并行策略（如减少通信量的缩减并行）。这种相互依赖性表明，为了达到最佳性能，必须对这两种优化进行联合考虑。然而，联合优化面临两大挑战：
1. 表示问题：需要一种能够同时表达计算、并行化和通信的统一表示，以便代数变换可以感知并适应并行化策略。
2. 可扩展性问题：联合搜索空间比单独优化大得多，需要高效的搜索算法来应对模型和设备数量的增长。

研究目标与创新点：为了解决上述问题，Unity提出了一个新颖的框架，其核心思想是将代数变换和并行化都表示为对统一并行计算图（PCG）的图替换，并使用分层搜索算法高效地识别最佳的替换组合。Unity的主要贡献如下：

1. 统一的图表示（并行计算图，PCG）：引入了PCG，这是一种能够同时表达DNN分布式训练中的计算、并行化和数据移动的统一表示。所有现有的并行化策略都可以表示为特定的PCG，而代数变换和并行化策略都可以统一视为图替换序列。这解决了表示问题，使得联合优化成为可能。

2. 变换的自动生成与验证：Unity能够基于算子规约列表自动生成并验证图替换。与以往工作分别生成并行策略或代数变换不同，Unity利用PCG，能够用单一方法生成这两种变换，甚至包括之前自动化方法中缺失的代数-并行化混合优化。这大大减少了支持新算子或并行维度所需的人工工程量。

3. 可扩展的联合优化算法：Unity采用一种新颖的三级分层搜索算法，在保持可扩展性的同时，共同优化PCG替换和设备放置。该算法的成本模型同时考虑了计算和通信时间，并能处理定制的网络拓扑和异构计算设备。尽管搜索空间指数级增大，Unity的搜索时间仍能保持在20分钟以内，远低于实际的模型训练时间。

通过这些技术，Unity提供了一种通用且可扩展的方法来优化分布式DNN训练。实验评估表明，在多达192个GPU上运行七个真实世界的DNN时，Unity的性能比现有DNN训练框架高出最多3.6倍。

A3 背景知识

并行化：DNN训练中的张量代数具有大规模并行的特性，这为并行化创造了许多机会。本文识别了六种主要的并行化形式：
1. 数据并行（Data parallelism）：这是现有框架【6，Martín Abadi等人，Tensorflow: A system for large-scale machine learning，2016，OSDI】、【9，Tianqi Chen等人，MXNet: A flexible and efficient machine learning library for heterogeneous distributed systems，2015，CoRR】、【42，PyTorch，2017】中最常用的方法。数据并行在每个设备上保留整个DNN模型的副本，并为每个设备分配一部分训练数据。
2. 模型并行（Model parallelism）：将一个DNN模型分割成不相交的子模型，并在专用设备上训练每个子模型。
3. 空间并行（Spatial parallelism）：将张量的空间维度（例如，图像的高度和宽度）分割成多个分区，每个分区分配给一个特定的设备【27，Zhihao Jia等人，Beyond data and model parallelism for deep neural networks，2019，SysML】。空间并行通常需要在设备之间同步共享元素（例如，不同子图像边界上的共享像素）。
4. 缩减并行（Reduction parallelism）：利用张量代数算子的线性特性。对于矩阵乘法$C = A \times B$，缩减并行将A沿其列分割，B沿其行分割如下：$A = [A_1, . . . , A_n]$，$B = [B_1^T , . . . , B_n^T ]^T$。矩阵乘法分布在n个设备上，第i个设备计算$C_i = A_i \times B_i$。之后通过一个额外的归约操作恢复原始结果：$C = \sum_i C_i$。
5. 流水线并行（Pipeline parallelism）：利用在不同训练迭代之间进行并行化的机会【39，Deepak Narayanan等人，Pipedream: Generalized pipeline parallelism for dnn training，2019，SOSP】。
6. 算子特定并行（Operator-specific parallelism）：新DNN算子的引入提供了算子特定的并行化机会。例如，Transformer【54，Ashish Vaswani等人，Attention is all you need，2017，CoRR】中使用的批处理矩阵乘法公式为：$output(s, h, o) = \sum_i input(s, h, i) \times weight(h, o, i)$。这与典型的矩阵乘法不同，因为它为每个输入样本应用不同的权重。因此，这些跨注意力头（即h维度）的批处理矩阵乘法可以并行运行，而无需任何张量复制或同步。

并行化的权衡：大多数并行化并非纯粹的性能优化，而是在不同成本指标之间的权衡。例如，应用数据并行会降低每个设备的计算时间，但代价是增加了存储和同步模型参数的内存使用和数据移动。因此，DNN算子通常需要结合多种并行形式才能实现最佳性能。

代数变换：代数变换非常多样，不易像并行化形式那样进行分类，因此这里提供一些例子。
* 算子融合（Operator fusion）：如图4a所示，这是最基本的代数变换。在未融合的情况下，设备需要两次从内存加载和存储激活值，即在每个算子前后各一次。如果两个算子被融合，组合后的核可以在将MatMul的输出存回内存时计算ReLU操作。
* 更复杂的变换：如图4b所示，通过利用DepthwiseConv的线性特性，一个先前需要两次DepthwiseConv操作的计算现在只需要一次加上一个额外的Add操作，从而有效地将计算量减半。

变换的组合效应：小型代数变换可以组合起来产生巨大的变化。如图5所示的变换序列，虽然每个单独的变换相对较小，但最终的输出与原始计算图有根本的不同。此外，并非所有的性能增益都能在单次变换中实现：例如，从图2到图4通过减少执行的DepthwiseConv操作数量来减少计算量，但必须先经过性能比图2或图4都差的图3。

中间表示：大多数现有的优化框架将DNN架构表示为一个计算图：节点是数学张量算子（如矩阵乘法等），边是算子之间传递的张量（即n维数组），如图6a所示。代数变换通过迭代应用图替换来执行，而模型通过为每个节点分配一组并行化注释来进行并行化。

现有表示的局限性：这种表示方式存在两个局限性。
1. 阻碍联合优化：虽然为代数变换（即图替换）和并行化（即节点注释）使用不同的表示很方便，但这阻碍了联合优化。关键问题是代数变换可以添加或替换图中的节点，而并行化将计算图视为静态的，因此无法处理这些新创建的、未注释的节点。这使得两种搜索算法无法交错进行，因为任何时候一个替换都可能将一个有效的并行化方案变成无效的。
2. 通信成本不明确：计算图没有明确地捕捉与并行化相关的通信成本。这种缺失使得代数变换难以推断其对最终模型性能的影响。

A2 方法细节

3 并行计算图 (Parallel Computation Graph)

PCG的引入：为了解决第2.3节中描述的现有模型表示的缺点，我们引入了并行计算图（PCG）作为分布式DNN训练的统一表示，它能够同时表达计算、并行性和通信。PCG允许Unity将代数变换和并行化都视为对一个公共图的图替换。虽然PCG不是第一个将计算和并行化合并到单个图中的表示，但PCG是为优化量身定制的，因此在关键方面与先前的统一图表示有所不同，我们将在第3.4节中讨论这些差异。

PCG的扩展：PCG扩展了现有的计算图表示，允许节点除了表示数学张量操作外，还可以表示并行化的变化；边除了表示数据依赖外，还可以表示张量数据的分布式移动。增加了一组并行化算子，使PCG能够表达所有现有的并行化策略，并为训练过程中的数据移动及其相关成本提供了明确的表示。此外，PCG中的每个算子都与一个机器映射相关联，表示该算子的执行如何映射到并行机器中的各个处理器。图6b展示了一个PCG的例子。

3.1 张量表示

张量维度：Unity将张量建模为一组数据维度，每个维度有两个字段：一个大小（size）和一个度（degree）。度字段指定了张量在该维度上被划分的分区数量。每个张量还包括一个特殊的副本维度（replica dimension），表示该张量数据的副本数量。

3.2 机器映射

定义与功能：PCG中的每个算子都关联一个机器映射，这是一个n维的设备/处理器数组，指定了算子计算的每个部分在哪个设备上运行。更正式地说，给定一个算子和一组n个适用的并行维度，其度分别为$d_1, . . . , d_n$，Unity将该算子划分为$d_1 \times d_2 \times . . . \times d_n$个并行任务，我们用形如$(i_1, . . . , i_n)$的元组索引来引用这些任务，其中$0 \leq i_k < d_k$。机器映射是一个从任务索引$(i_1, . . . , i_n)$到将用于运行该并行任务的单个GPU的映射。为方便起见，我们还将整个PCG的机器映射定义为其所有组成算子的机器映射的集合。

示例：图7展示了我们评估中使用的Summit计算节点【55，Sudharshan S Vazhkudai等人，The design, deployment, and evaluation of the coral pre-exascale systems，2018，SC】的一些机器映射示例。硬件架构如图7a所示。图7b显示了用于数据并行的基本一维机器映射，而图7c显示了结合数据和模型并行的混合并行策略的二维机器映射，其中模型并行应用于同一计算节点内的GPU，数据并行应用于不同计算节点之间。图7d显示了一个三维机器映射，其中我们在连接到同一CPU的GPU之间应用模型并行，在连接到不同CPU但在同一计算节点上的GPU之间应用缩减并行，并在不同计算节点之间应用数据并行。

自定义与优势：Unity包含一套全面的机器映射，能够捕捉并行机器的有效使用方式。此外，开发者可以注册针对特定硬件架构定制的机器映射。例如，当集群中的节点对具有不同的网络带宽和延迟时，可以向现有机器映射添加一个额外的维度来表示节点级别的局部性。机器映射提供了两个关键的理想属性：表达性和可扩展性。PCG中所有有效的并行任务分布都可以用少数几个机器映射来捕捉，并且可以通过添加额外的机器映射轻松利用机器硬件架构的复杂特性。机器映射还允许Unity在保持设备数量线性关系的同时捕捉所有有效的设备分配，并通过从考虑中移除低效分配来帮助Unity的搜索算法。

3.3 并行化算子

六种核心算子：Unity使用六个并行化算子来捕捉与不同并行化策略相关的计算和通信成本。这六个算子进一步分为三对，其中一个算子是另一个的“反向传播”（例如，当在Partition上进行反向传播时，它在语义上等同于Combine，反之亦然）。这三对是：
1. Partition 和 Combine：Partition和Combine改变张量的并行度。具体来说，Partition通过将一个维度分割成多个大小相等的分区来增加张量维度的并行度，如图8a所示。Combine执行相反的操作：通过将多个分区连接成一个来减少张量的并行度。
2. Replicate 和 Reduce：Replicate和Reduce通过复制和求和张量来控制副本维度的并行度，如图8b所示。参数同步被自然地捕捉为应用于权重张量的Replicate操作的反向传播。
3. Pipeline 和 Batch：Pipeline将一个张量维度分割成大小相等的分区，并一次处理一个分区，而Batch则跨迭代聚合张量（见图8c）。请注意，Pipeline不修改张量维度的并行度，而是减小其大小。

表达能力：作为PCG表达能力的基本展示，图9说明了Unity的六个并行化算子如何表示第2.1节中的一些并行化策略示例。这些并行化算子也可以组合起来创建混合并行。图8d展示了一个在同一张量维度上应用Replicate和Partition的例子，即复制张量并对每个副本进行分区。为了提高效率，Unity在运行时会将特定的并行化算子序列替换为融合版本（例如，一个Reduce后跟一个Replicate可以实现为一个AllReduce）。

3.4 讨论与比较

PCG与注释计算图的对比：Unity决定使用PCG而不是带注释的计算图，是基于这些表示对联合搜索的便利性，而非带注释计算图的根本局限性。理论上，存在与PCG同构的注释语言，但设计这种语言的尝试很快会遇到一些困难。
* 注释数量问题：首先，因为每个算子可以使用不同形式的并行，包括算子特定的并行形式，注释的数量很快会变得过大。哪些算子支持哪些注释，以及它们的语义和组合，都必须被嵌入到表示本身中。相比之下，Unity的PCG将这些知识移入PCG替换中，而这些替换是自动生成的。这种关注点分离使得Unity的核心更简单，更易于维护。
* 通信模式重构困难：其次，不明确表示通信会迫使其源和目标节点注释来重构数据流边上的通信模式，由于Unity考虑的并行形式表达力强，这变得很困难。具体来说，支持n个并行维度需要考虑多达$2^n$个不同的这些维度的子集，因此在算子之间存在$2^n \times 2^n = 4^n$种潜在的通信模式。Unity在整个搜索过程中明确表示通信模式，避免了需要复杂的分析来重构它们。通过一小组并行化算子来表示这些模式，也使得Unity能够轻松识别和优化常见的通信模式，例如将一对Reduce-Replicate算子作为AllReduce执行。在带注释的计算图中，这些优化与重构通信模式的代码纠缠在一起，增加了显著的复杂性和实现工作量。
* 联合优化挑战：最后，联合优化带注释的计算图是具有挑战性的，因为代数变换可以引入新的算子，由于这些算子尚未并行化，它们缺少注释。因此，内部表示变得不完整，成本也无法定义。可以添加额外的机制来“填补”这些缺失的注释，例如插入一个随机注释、一个固定值、一个来自相邻节点的值（当相邻节点有不同的并行化时这变得具有挑战性），或者评估有效的并行化并选择最佳的一个。然而，Unity的PCG通过将新算子的每种并行化策略表示为一个或多个PCG替换，避免了这种额外的复杂性，为联合优化提供了一种高效且统一的方法。

与pONNX的比较：Unity不是第一个将计算和并行性集成到单个图中的系统：pONNX 【57，Fei Wang等人，Parallel Training via Computation Graph Transformation，2019，IEEE International Conference on Big Data】 提出使用Split、Concat和自定义算子Send和Recv来实现这一点。然而，Unity专注于优化，而pONNX被设计为一种序列化格式，这导致了关键差异。
* 算子表示：首先，pONNX中一个并行度为n的算子被复制n次，要求优化器从多个节点重构该算子。Unity只是添加一个并行化算子，因此该算子在PCG中仍然是单个节点。
* 通信表示：其次，pONNX为每个通信分配其自己的Send/Recv节点，这极大地增加了图的大小。由于DNN训练中的通信模式高度规律，Unity避免了实例化每个通信，而是优化通信模式（例如，Reduce、Replicate等），这使得Unity能够表示通信成本而无需推理单个通信。
* 设备放置：最后，pONNX将设备放置作为算子的一部分，而Unity将其作为机器映射分开表示。这使得Unity的搜索可以独立优化设备分配，并忽略由大量具有相同能力的计算设备所产生的对称性。

其他统一表示：还提出了其他统一表示【47，Keshav Santhanam等人，DistIR: An Intermediate Representation and Simulator for Efficient Neural Network Distribution，2021，arXiv】、【48，Michael Schaarschmidt等人，Automap: Towards Ergonomic Automated Parallelism for ML Models，2021，arXiv】，这些将在第7节中讨论。

4 图替换 (Graph Substitutions)

替换作为统一表示：由于Unity将代数变换和并行化都表示为图替换，联合优化的有效性依赖于拥有一组合适的图替换。潜在替换的数量随大小呈指数增长，因此Unity将大型复杂的代数变换和并行化策略表示为小型PCG替换的组合。例如，图10展示了Megatron-LM 【50，Mohammad Shoeybi等人，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，CoRR】中使用的手动调整并行化策略的替换序列。

替换生成：为了减少支持新并行化策略的工程量，Unity自动生成并形式化验证所有达到固定大小的有效PCG替换，作为搜索算法构建复杂优化的“基础集”。这也使Unity不仅能自动发现代数变换和并行化策略，还能找到先前方法遗漏的两者的新型混合体。为此，Unity采用了TASO的超优化方法【25，Zhihao Jia等人，Taso: Optimizing deep learning computation with automatic generation of graph substitutions，2019，SOSP】。

候选替换识别：与TASO一样，Unity分两步发现替换：首先使用快速启发式方法识别候选替换，然后使用更昂贵的正式验证来确保正确性。为了找到候选替换，Unity枚举所有达到固定大小的可能PCG。请注意，这个固定大小并不限制Unity可以应用的变换的大小，因为许多更大的替换是较小替换的组合。对于每个生成的PCG，Unity计算一个指纹：通过在一些固定输入张量上评估PCG生成的其输出张量的哈希值。为了让Unity能够考虑并行化，我们扩展了TASO【25】中的指纹函数，以包括每个张量维度的并行度。如果一对PCG具有相同的指纹，则它们被视为候选替换。并行化的加入使Unity在TASO先前识别的743个候选替换之外，发现了651个新的候选替换。

替换验证：与TASO类似，Unity使用自动定理证明器（在我们的实现中是Z3【12，Leonardo De Moura和Nikolaj Bjørner，Z3: An efficient smt solver，2008，TACAS/ETAPS】）来形式化验证新的替换。算子规范以一阶逻辑提供，其中算子表示为其输入和配置参数的函数。例如，$Reduce(d, x)$定义了一个输入为x、并行度为d的Reduce算子。Reduce与矩阵乘法可交换的事实由以下算子属性捕获（其中$Replicate(d, y)$表示输入为y、并行度为d的Replicate）：

$$\forall d, x, y. \mathrm{Matmul}(\mathrm{Reduce}(d,x),y) = \mathrm{Reduce}(d,\mathrm{Matmul}(x,\mathrm{Replicate}(d,y)))$$

操作符属性开发：我们遵循TASO开发算子并行化属性的方法：我们尝试使用Z3形式化验证所有候选替换，当一个替换无法被验证但是正确的，我们添加缺失的算子属性。这个过程重复进行，直到Unity发现的所有651个新替换都被验证。总的来说，我们在TASO的43个属性【25，表2】之外引入了33个算子属性，以验证所有的PCG替换。整个替换生成和验证过程总共耗时30分钟。由于可用的替换仅在添加新算子或并行形式时才会改变，这个过程可以完全离线运行，从而不影响Unity联合搜索算法的执行时间。

替换示例：Unity生成的大多数新替换只是陈述了某个算子有效的并行性。例如，图11a中的替换表明ReLU在行维度上支持空间并行。然而，结合代数变换和并行化也产生了新颖的混合体，如图11b所示的例子，其中Unity识别出Add算子等价于一个Concat后跟一个Reduce。在左边的PCG中，当输入1和输入2位于不同的设备上，且输出需要在这些设备上复制时，输入1和输入2必须被发送到单个设备上进行相加，然后再发送回来。通过应用这个变换，Unity能够通过一个Concat（不移动数据）将输入张量合并成一个单一的分布式张量，并用一个Reduce后跟一个Replicate（实现为AllReduce）来替代通信。

5 联合优化 (Joint Optimization)

核心问题：本节描述了Unity用于联合优化代数变换和并行化的搜索算法。核心问题如下：给定一个PCG（第3节）、一组算子级别的机器映射（第3.2节）和一组PCG替换（第4节），找到（1）一个PCG替换序列和（2）结果PCG的一个机器映射，以最小化每次迭代的训练时间。一个关键挑战是统一代数变换和并行化所产生的指数级增大的搜索空间。搜索还必须能够扩展到复杂的DNN（即大的输入PCG）和大量的计算设备（即大量的算子级机器映射）。

三层分层搜索算法：Unity使用一个三层分层搜索算法，如图12所示。简而言之，Unity将一个输入的PCG分解成子图，为每个子图确定一个优化的替换序列（这需要为每个候选项确定优化的机器映射），然后组合这些子解决方案以产生最终输出。这使得Unity能够扩展到超过300个算子的DNN和192个GPU的机器，同时将搜索时间保持在20分钟以下，这与训练现代DNN所需的数小时或数天相比可以忽略不计。

5.1 替换选择

基于成本的回溯搜索：Unity使用来自TASO【25，Zhihao Jia等人，Taso: Optimizing deep learning computation with automatic generation of graph substitutions，2019，SOSP】的基于成本的回溯搜索算法来识别一个替换序列，该序列最小化输入PCG的执行时间。Unity维护一个按执行时间排序的候选PCG队列，直到队列为空或超过固定预算，Unity会迭代地从队列中移除最佳候选项，并通过在PCG中每个适用的位置应用每个可用的替换来生成新的候选项。执行时间比目前所见的最佳候选PCG差一个阈值因子的候选PCG会被剪枝，其余的则插入队列。该阈值因子允许用户在搜索时间和探索量之间进行平衡。在我们的实验中，我们使用的阈值因子为1.05。

成本估算器需求：该算法允许Unity探索任意的替换序列，但需要一个准确的成本估算器来评估每个候选PCG的执行时间。由于PCG只包含每个算子的并行化方式，而不包含其分配到的设备（即机器映射），因此这个成本估算器必须首先确定一个优化的机器映射。必须使用高效的算法来识别这个映射，因为成本估算器会对每个候选PCG被调用。第5.2节介绍了我们寻找优化机器映射的算法。

5.2 寻找优化的机器映射

核心观察与分解策略：Unity搜索算法的最低层是为候选PCG识别优化的机器映射。这一层的关键观察是，大多数现代DNN架构由独立的并行计算链的线性链组成。例如，ResNeXt【19，Kaiming He等人，Deep residual learning for image recognition，2016，CVPR】是围绕两个并行的卷积链构建的（见图13），这些卷积链重复构成最终模型。Unity利用这种结构，通过分别为这些线性链和并行链进行顺序和并行图分割，将它们递归地分解为独立的子图。图13演示了如何迭代地应用顺序和并行图分割来将ResNeXt模块分解为可以通过动态规划解决的递归子问题。

顺序图分割：顺序图分割通过找到一个后支配节点n来划分输入PCG G，使得从G的输入到输出的所有路径都经过n。这个后支配节点将G分割成两个不相交的子图G1和G2。由于G2的所有部分都依赖于n，而n依赖于G1的所有部分，所以G1中的每个算子必须在G2中的任何算子开始之前完成。这将为G寻找优化机器映射的任务简化为为G1、G2和n优化机器映射。例如，对于图13中的分割1，假设没有其他分割（如0）已经被应用，n将是Add节点，G1将是所有从Input到但不包括Add的节点，G2将是Add之后到Output的所有节点。

并行图分割：并行图分割将一个PCG G划分为可以并行执行计算的独立子图。在这种情况下，Unity考虑两种运行G1和G2的方式：顺序运行（可以访问全部机器资源）或并行运行（每一方获得不相交的可用资源份额），并选择更快的一种。Unity不允许串行和并行执行的组合，即分支部分并行运行、部分串行运行。虽然这排除了一些策略，但考虑它们会显著降低Unity的可扩展性，因为它需要分析指数级多的算子交错。正如第6节的结果所证明的，这些策略对于实现良好性能并非必需。为了确定在并行运行时如何划分可用资源，Unity遍历了可以分配给每一方的所有可能的资源数量。通过考虑资源数量，Unity忽略了那些仅仅是分配的GPU不同而GPU数量和位置相同的冗余划分，将一个指数级的设备子集搜索替换为一个二次级的资源数量搜索。

缓存优化：作为一项额外的优化，Unity为所有子图维护一个已选择的机器映射的缓存。由于替换选择为每个替换生成一个新的候选项，且每个替换只修改PCG的一小部分，许多候选PCG的大部分子图与其他候选项是相同的。这使得Unity可以跳过计算除被考虑的替换所修改的PCG部分之外的所有部分的成本和机器映射。

5.3 扩展到大图

可扩展性问题：即使有动态规划算法和跨调用缓存，目前描述的搜索算法也无法扩展到大型模型。为了理解原因，我们研究了替换选择中候选PCG的数量如何随输入PCG的大小而扩展。如第5.1节所述，在每次迭代中，Unity为每个替换的每个可能应用生成一个候选PCG。在最坏的情况下，这需要检查$O(2^{gs})$个候选项，其中g是PCG中的节点数，s是Unity考虑的替换数。实际上，s对搜索时间的影响有限，因为Unity考虑的替换中只有一小部分可以应用于任何一个模型，但对于大型模型，g的指数行为成为问题。

解决方案：图分割：为了解决这个问题，我们借鉴了第5.1节的方法，将PCG分解为独立的顺序子图。然而，这种方法阻止了在这些分割点上应用替换，这是一个问题，因为Unity使用替换来表示并行化。因此，简单的图分割会将所有分割点上的并行度降低到1，从而排除了许多常见且重要的并行化策略，例如在整个模型上使用数据并行。

跨分割点并行化搜索：Unity通过显式搜索每个分割位置上的最优并行化来解决这个问题。更具体地说，对于跨分割点通信的张量的每一种可能的分区方式，Unity都会在以下条件下优化产生的两个子图：第一个子图的输出和第二个子图的输入都必须与正在考虑的分区方式相匹配。当任一子图不满足此条件时，会插入并行化算子，以确保因分区方式改变而产生的任何通信成本都被计入。

处理Replicate和Reduce的挑战：这个方法对Partition和Combine有效，但对Replicate和Reduce遇到了问题。例如，考虑跨分割位置的张量，其副本度被搜索算法固定为2。为了强制第一个子图输出这种格式的张量，搜索算法可以在其最后插入一个Replicate操作，同样地，算法可以在第二个子图的开头插入一个Reduce操作。然而，这将错误地将张量缩放2倍！核心问题是，与Partition和Combine不同，Replicate和Reduce不是彼此的逆操作。幸运的是，由于跨越计算图多个节点的缩减并行在实践中很少有用，我们将跨分割点的分区限制为仅那些副本度为1的分区。

选择分割点：为了减少这些分割阻止的代数变换数量，Unity遵循MetaFlow【26，Zhihao Jia等人，Optimizing dnn computation with relaxed graph substitutions，2019，SysML】的方法，选择那些在保持最小子图大小k的同时，干扰最少替换的分割位置。因此，图分割将最坏情况下的候选PCG数量从g的指数级降低到g的线性级，具体是从$O(2^{gs})$降低到$O(\frac{g}{p}k \times 2^{ks})$，其中p是有效张量分区的数量。

成本估算：为了估算算子运行时间和通信成本，我们使用了与先前工作【24, Zhihao Jia等人，Exploring hidden dimensions in accelerating convolutional neural networks，2018，ICML】、【27, Zhihao Jia等人，Beyond data and model parallelism for deep neural networks，2019，SysML】类似的方法。虽然可能有更精确的成本模型【47, Keshav Santhanam等人，DistIR: An Intermediate Representation and Simulator for Efficient Neural Network Distribution，2021，arXiv】，但我们没有注意到我们的模型不准确性导致任何问题。

流水线并行：在考虑流水线并行时，Unity采用了1F1B调度（即在每个设备上交错进行前向和后向微批次）和PipeDream-2BW【40，Deepak Narayanan等人，Memory-efficient pipelineparallel dnn training，2020】的权重更新语义，这实现了高训练吞吐量和低内存占用。为了减少搜索空间，Unity只考虑将流水线并行应用于PCG中所有算子的策略，因为PCG中的非流水线并行算子会使流水线并行的好处失效。此外，与先前的工作【20, Yanping Huang等人，Gpipe: Efficient training of giant neural networks using pipeline parallelism，2018】、【39, Deepak Narayanan等人，Pipedream: Generalized pipeline parallelism for dnn training，2019，SOSP】、【65, Lianmin Zheng等人，Alpa: Automating inter- and intraoperator parallelism for distributed deep learning，2022，CoRR】类似，Unity只考虑顺序流水线并行，其中每个阶段只与流水线中的单个下一阶段通信（除了最后一个阶段，它在前向处理后直接进行反向传播）。Unity还遵循先前的工作【16, Shiqing Fan等人，Dapple: A pipelined data parallel approach for training large models，2021，PPoPP】、【20】、【53, Jakub M Tarnawski等人，Piper: Multidimensional planner for dnn parallelization，2021，NeurIPS】，假设一个小批次中的微批次数量远大于流水线阶段的数量，因此可以忽略流水线初始化引入的额外延迟。这些约束使得Unity能够在保持合理搜索时间的同时，探索一个包含现有流水线并行策略的全面搜索空间。搜索算法也稍作修改：我们不使用每次迭代的运行时间作为吞吐量的代理，而是直接最大化吞吐量。

A4 实验环境

• 实现与软件：Unity基于FlexFlow【27，Zhihao Jia等人，Beyond data and model parallelism for deep neural networks，2019，SysML】实现，修改其以使用PCG表示模型，并增加了对新并行形式的支持。替换生成器基于TASO【25，Zhihao Jia等人，Taso: Optimizing deep learning computation with automatic generation of graph substitutions，2019，SOSP】实现，并扩展了其指纹函数以考虑并行化。
• 硬件配置：所有实验均在Summit超级计算机【2，Summit supercomputer，2018】、【56，Sudharshan S. Vazhkudai等人，The design, deployment, and evaluation of the coral pre-exascale systems，2018，SC】上进行。每个计算节点配备：
  • CPU: 2个IBM POWER9 CPU。
  • 内存: 512 GB 主内存。
  • GPU: 6个NVIDIA Volta V100 GPU。同一节点内，3个GPU连接到同一个CPU，并通过NVLink互连。
  • 网络: 节点间通过Mellanox EDR 100Gb InfiniBand连接。
• 模型与数据集：评估使用了七个DNN模型，具体信息如下表所示。
  • 模型: ResNeXt-50, Inception-v3, BERT-Large, DLRM, XDL, CANDLE-Uno, MLP。
  • 数据集: ImageNet【46】, WikiText-2【35】, Criteo Kaggle【4】, Dose response data【1】, 以及合成数据。
  • 任务: 图像分类、语言模型、推荐系统、精准医疗、回归。
  • 模型参数: 遵循先前工作【3, 14, 38, 41, 60】设置超参数。MLP模型包含16个密集层，隐藏维度为8192。
• 训练配置：
  • 批大小（Per-GPU）: ResNeXt-50和Inception-v3为64，BERT-Large为4，DLRM和XDL为1024，CANDLE-Uno和MLP为256。全局批大小为 n_gpus * per_gpu_batch_size。
  • 优化器: BERT-Large使用Adam【29】，学习率为0.0001；其他模型使用SGD【18】，学习率为0.01。
• 搜索时间：Unity的搜索时间非常短，对于除Inception-v3之外的所有模型，在最多192个GPU上，搜索时间都在10分钟以内。对于最复杂的Inception-v3（323个算子），搜索时间也在20分钟内完成，与数小时到数天的训练时间相比可以忽略不计。

任务	架构	数据集
图像分类	ResNeXt-50 [60]	ImageNet [46]
	Inception-v3 [51]	ImageNet [46]
语言模型	BERT-Large [14]	WikiText-2 [35]
推荐系统	DLRM [41]	Criteo Kaggle [4]
	XDL [28]	Criteo Kaggle [4]
精准医疗	CANDLE-Uno [3]	Dose response data [1]
回归	MLP [17]	Synthetic data

Table 1: 评估的七个DNN概述。

A4 实验结果

端到端性能评估

该实验比较了Unity与现有框架（Megatron【50】, DeepSpeed【43】）以及顺序优化方法（TASO【25】+FlexFlow【27】）的端到端训练吞吐量。结果如图14所示。

• BERT-Large: Unity的性能与专家手动优化的Megatron相当，并优于DeepSpeed和FlexFlow。Unity自动发现的策略与Megatron的专家策略几乎相同，证明了其在高度优化模型上自动发现高性能策略的能力。
• DLRM 和 CANDLE-Uno: 这两个模型内存需求超出单GPU容量，无法使用纯数据并行。与专家设计的混合并行策略【38】和TASO+FlexFlow相比，Unity在DLRM上实现了高达3.6倍的加速，在CANDLE-Uno上实现了1.6倍的加速。
• 其他模型: 与数据并行和TASO+FlexFlow相比，Unity在ResNeXt-50上性能持平（因为数据并行已是其最优策略），在Inception-v3上提升1.3倍，MLP上提升2.0倍，XDL上提升1.9倍。

结论: 性能提升主要来自两个方面：(1) 支持算子特定的并行化策略；(2) 联合优化代数变换和并行化。

并行维度影响分析

该实验通过在BERT-Large上进行消融研究，评估了不同并行维度对训练性能的提升。

• 实验内容: 从数据+模型并行开始，逐步增加缩减并行、注意力头并行和流水线并行，并测量训练吞吐量。
• 实验结果 (图15a):
  • 仅添加缩减并行对性能无提升。
  • 同时结合缩减并行和注意力头并行，性能提升高达1.2倍。
  • 启用流水线并行后，总体加速达到1.4倍。
• 结论: 混合并行策略和算子特定的并行维度对于最大化DNN训练性能至关重要。

联合优化效果分析

该实验旨在量化联合优化相对于顺序优化的优势。为了隔离变量，对比实验使用了Unity的全部并行维度和搜索算法，仅改变优化方式（联合 vs. 顺序）。

• 实验结果 (图15b): 联合优化比顺序优化带来了高达1.4倍的加速。
• 具体优化案例 (图16):
  1. MLP (图16a): 联合优化选择不融合第一个MatMul和ReLU，从而能够应用通信效率更高的缩减并行（通过AllReduce实现），显著减少了通信开销。
  2. DLRM/XDL (图16b): 联合优化通过将Concat后的MatMul替换为多个独立的、与Embedding算子使用相同模型并行的MatMul，消除了Concat操作带来的all-to-all通信瓶颈。
  3. EmbeddingBag (图16c): Unity将EmbeddingBag算子转换为普通的Embedding算子，从而解锁了更多的并行化机会。
• 结论: 联合优化能够发现顺序优化无法找到的、跨越代数变换和并行化边界的复杂优化机会，从而带来显著的性能提升。

搜索算法可扩展性分析

该实验通过在ResNeXt-50上进行消融研究，评估了三种搜索优化技术（图分割、跨调用缓存、动态规划）对搜索时间的影响。

• 实验结果 (表2):
  • 所有优化启用: 搜索时间从6 GPU到48 GPU大致呈线性增长，证明了算法的可扩展性。
  • 禁用图分割: 搜索时间增加了4.3至8.8倍，且呈非线性增长。
  • 禁用跨调用缓存: 搜索时间额外增加了8.9倍。
  • 禁用动态规划: 即使是最小的案例也会超时。
• 结论: 图分割、跨调用缓存和动态规划这三种技术对于Unity搜索算法的性能和可扩展性至关重要。

	All	w/o Split	w/o Cache+Split
	Time	Scaled	Time
6 GPUs (1 node)	57s	1×	4m 01s
12 GPUs (2 nodes)	1m 47s	1.9×	11m 15s
24 GPUs (4 nodes)	3m 00s	3.1×	> 1h
48 GPUs (8 nodes)	5m 55s	6.1×	> 1h

Table 2: 搜索算法消融研究。“Scaled”数字是相对于启用所有优化时2 GPU时间的倍数。

A7 补充细节

7 相关工作

手动设计的并行化策略：大多数现有的DNN框架【5, ONNX Runtime】、【6, TensorFlow】、【42, PyTorch】、【43, DeepSpeed】、【49, Mesh-TensorFlow】使用手动设计的并行化策略。例如，Neo【38】为DLRM优化，对计算密集型算子使用数据并行，对通信密集型算子使用模型并行。Megatron-LM【50】为大型语言模型提出了一种结合数据、缩减和注意力头并行的模型特定策略。这些策略通常不具通用性。Unity能够自动发现性能更优的策略。

自动DNN并行化：近期工作提出了自动化方法。例如，ColocRL【36, 37】和Placeto【7】使用强化学习寻找设备放置。Baechi【22】使用内存受限算法快速进行设备放置。FlexFlow【27】使用随机搜索优化数据、模型和空间并行。GSPMD【61】基于用户提示寻找并行策略。PipeDream【39】使用动态规划优化流水线和数据并行。Tofu【59】使用递归搜索自动发现并行维度。Tarnawski等人【52, 53】和Alpa【65】使用动态规划和整数线性规划结合优化。Whale【23】支持异构设备。TensorOpt【8】考虑多目标优化。AutoSync【63】学习优化同步策略。然而，除Tofu外，这些方法支持的并行维度有限，且没有一个能联合优化代数变换和并行化。Unity支持所有现有并行维度，可扩展，并进行联合优化。

自动代数变换：TASO【25】能自动发现DNN的代数变换，但不支持并行化。Unity采纳了TASO的超优化思想来生成和验证PCG替换。然而，Unity处理的搜索空间远大于TASO，并需处理设备分配等额外任务。为应对此挑战，Unity引入了三个新颖的搜索技术：用于寻找优化机器映射的动态规划算法、利用图替换局部性的子图缓存，以及实现复杂模型可扩展性的并行兼容分治方法。

自动DNN代码生成：TVM【10, 11】和Ansor【64】等工作专注于为DNN算子生成硬件特定的优化代码。Unity在比这些方法更高的层次上优化DNN计算。因此，Unity的优化是正交的，可以与现有代码生成技术结合。将代码生成集成到Unity中是未来的工作。

DNN并行化的中间表示：TensorFlow【?】、MLIR【31, 32】、Relay【45】和ONNX【33】使用基于图的IR表示DNN计算。它们通过注释算子来表示分布式训练，将代数变换和并行化分开并顺序优化，从而错失了联合优化的机会。pONNX【57】、automap【48】和DistIR【47】提出了能同时表示计算和通信的IR，但它们的抽象层次太低，不适用于Unity风格的联合优化（详见3.4节）。Unity使用更适合优化的更高层表示——PCG，并将并行化和代数变换都表示为PCG上的图替换。

8 局限性与未来工作

对DNN结构的假设：为了扩展到大型DNN和机器，Unity的搜索算法利用了现代DNN的顺序和并行结构（见5.2节）。然而，存在一些违反此结构的DNN架构（如NASNet【66】）。将Unity扩展到这些DNN会提高通用性，但可能会以降低可扩展性为代价。

未考虑的优化：虽然Unity成功优化了两种最重要的优化类别（即代数变换和并行化），但还有多种其他优化目前未被考虑，例如张量卸载和重计算【21, 30, 44】。PCG可以扩展以表示这些优化，但第5节中介绍的搜索算法没有考虑内存使用，因此可能生成违反内存约束的并行化策略。虽然这些无效策略可以通过后续应用必要的张量卸载和重计算来变得有效，但不将这些优化纳入Unity的联合搜索可能会导致次优性能。因此，将内存优化集成到Unity的搜索算法中是未来研究的一个有前景的方向。

流水线并行的局限性：Unity对流水线并行的支持也存在局限性。虽然PCG能够表示在PCG中交错使用流水线并行和非流水线并行算子的策略，但我们的搜索算法为了减少搜索空间排除了这些情况。此外，Unity的搜索算法不考虑非顺序的流水线并行策略，即一个阶段可以有多个前驱/后继阶段。

A5 结论

本文介绍了Unity，这是第一个在分布式DNN训练中联合优化代数变换和并行化的系统。Unity将并行化和代数变换都表示为对一个统一图表示的替换，使用一种新颖的分层搜索算法来识别优化的替换序列，并能扩展到大量GPU和复杂的DNN。

我们在多达192个GPU上对七个真实世界的DNN基准进行了评估，结果表明，Unity的性能比最先进的并行化方法高出3.6倍，同时优化时间保持在20分钟以内。由于近一半的加速完全归因于使用联合优化而非顺序优化，Unity证明了联合优化是实用的，并且未来的系统需要包含它，否则将错失显著的性能增益。