阿里云机器学习 PAI 发布基于 HLO 的全自动分布式系统 TePDist，并宣布开源！

作者：品曲，宗雁，佀畅，侠雕，伟林

导读

一直以来，大模型在模型效果上被证明具有显著优势。而 ChatGPT 的出现，证明了其在工业生产工具方面具有巨大潜力。毫无疑问，大模型的训练需要巨大的算力，这对分布式并行框架是一项考验。现如今，业界越来越多的团队纷纷转向分布式框架的研究与开发之中，既有针对特定场景特定模型的极致手工优化，也包括面向通用模型通用场景的自动分布式工作。然而，它们在实际生产场景仍然有一些挑战。一方面，对于手工优化的框架，虽然其特定情况下的性能较为极致，但通用性存在限制。另一方面，现有的自动分布式框架可能在通用性和性能方面表现较好，但通常需要忍受较长的策略搜索时间，进而对业务落地产生一定的影响。为了应对以上问题，PAI 团队推出并开源了 TePDist（Tensor Program Distributed System），它通过在 HLO 上做分布式策略搜索，实现与用户模型语言解耦。并且在保持通用性的同时，在可接受的策略搜索时间内，追求高性能分布式策略。

TePDist 背后的技术框架如何设计？未来有哪些规划？今天一起来深入了解。

TePDist 是什么？

TePDist 是阿里云 PAI 团队自研的基于 HLO IR 层的全自动分布式深度学习系统，它不仅仅是一个分布式 Compiler，还拥有自己的分布式 Runtime，有效地解决了深度学习模型并行策略的自动搜索与分布式策略实施问题。

在架构方面，TePDist 采用 Client/Server 模式，实现分布式策略与模型描述的解耦。Server 端是 TePDist 最重要部分，它以 HLO IR 作为输入，自动探索并实施分布式并行策略；Client 端以用户描述的模型为输入，将其转换成 HLO IR。因此，任何具有转换 HLO IR 能力的 Client，都可经过适配后接入 Server 端。

在功能方面，TePDist 分为两个部分。一是在 HLO IR 上进行 SPMD（Data Parallel 和 Sharding）和 Pipeline 并行的策略搜索。并以此构建编译基于 Task Graph 的执行计划。二是高效运行执行计划的分布式执行引擎。

在使用方面，TePDist 提供了不同优化级别，高优化级别更加追求分布式策略质量，低优化级别会额外采取一些 Heuristic，以较为微小策略质量牺牲，换取更快地搜索时间，以此满足落地需求。

项目开源地址：https://github.com/alibaba/TePDist

TePDist 的主要特性如下：

• 一套完整的系统：采用 Client/Server 模式。Client 可以是任何能够生成 XLA HLO 的前端。Server 负责进行分布式策略规划，以及自动分布式任务拉起。Client 与 Server 解耦的动机是期望将来更方便地对接不同的前端框架。

• 以 HLO IR 作为 Server 端输入：HLO 的粒度刚刚好，目前看到的超大模型基本上有上万条 HLO 指令。在 HLO 这个层级做分布式策略规划，可以避免依赖灵活多变的模型高层次抽象，以不变应万变。

• 全自动探索分布式策略：TePDist 支持在没有任何 annotation 的情况下做策略搜索。当然，用户也可以选择通过 annotation 进行一定程度的干预。

• 定义不同优化级别：用户可以选择使用 O2 和 O3 两个优化级别，用以控制策略搜索时间。它们均为 Cost based 策略搜索方法。其中 O2 级别采用了一定的 Heuristic，会以牺牲轻微的并行策略质量为代价，换取更快地搜索时间，这对超大规模模型非常有用。O3 级别对策略的质量有一定的保证。

• 合理拆解策略搜索问题：TePDist 采用了多种手段，把策略探索问题拆解成优化子问题，运用多种算法，分别对子问题进行求解，有效管理了问题的复杂性。

• 具有特色的流水线并行规划：无需把 DAG 排成拓扑线性序，我们把 stage 划分建模成整数线性规划问题（ILP），使用 ILP Solver 自动寻找通讯量最小的切分方案。

TePDist 架构

TePDist 采用 Client/Server 分离的架构，将前端模型构建与后端策略搜索实施解耦。Server 端分为两个过程：

1. 编译（Build Execution Plan）：在这个过程中会构建执行计划，包括：自动分布式策略的搜索，device mesh 的分配，TepDist 执行引擎的构建——Task Graph 的构建以及为 Task Graph 确定静态调度顺序。架构和流程列于下图左图中；

2. 运行（Execute Plan）：编译阶段完成后，Client 下达 Input data 供给指令，即可触发运行该 Execution Plan 的运行。为了追求高性能和高可控性，我们完全自研了 TePDist 的运行时，具体包含无损分布式初始化，分布式 checkpoint，NCCL 复杂通信域的管理，多 worker 协同管理，以及 Task Graph 级别的 Memory reuse 等。架构和流程列于下图右图中；

自动分布式策略搜索

自动分布式是 TePDist 核心功能之一，我们对 SPMD 策略的搜索空间做了合理的分层拆解，在不同的层级使用不同的算法。同时对用户提供不同优化级别，用以控制策略搜索的时间。对 Pipeline stage 划分则无需对 DAG 按照线性排序，然后使用 ILP 对其 stage 的划分进行建模，同时也为用户提供了用于控制搜索时间的剪枝参数。以下对这部分内容做简单的描述，技术的具体细节请参考我们的论文：（TODO）

SPMD Strategy

由于 HLO 指令数较多，直接使用 ILP 建模求解可能会因问题规模过大，导致搜索时间过长。显然，通过缩图可以减小问题求解的规模。为此，TePDist 将 DAG 划分为了三个层级，并对每层使用不同的优化方法进行求解。

1. Cone 结构

HLO DAG 中存在大量的 Cone 结构，它对我们的 Formulation 非常重要。什么是 Cone 结构？它是包含两种节点的子图：Root 节点和非 Root 节点。其中，Root 节点为具有多个输入或特征为计算密集型的节点；非 Root 节点为其他节点。我们可以在 HLO 的 DAG 中识别所有的 Cone 结构。大多数 Cone 结构都呈现出“倒三角”的形态，如上图中灰色三角区域标出的部分。显然，Cone 代表更粗粒度的节点，我们的算法需要为每个 Cone 确定切分策略。每个 Cone 应该有多少种可能得切分策略？答案是取决于其 Cone Root 有多少种切分策略。具体做法：在 Cone 内，通过对 Cone Root 节点枚举的每种切分策略，我们都可以以通信代价最小为目标，通过贪心或动态规划方法求解 Cone 内剩余节点的切分策略。由此，我们可以得到该 Cone 的多种不同切分策略，这些策略都是每个 Cone 的候选策略。

2. Segment

虽然 Cone 结构的粒度更粗，在一定程度上缩减了问题规模。但对大模型来说，Cone 的数量依然巨大，可能需要进一步处理。事实上，大模型在结构方面具有重复堆叠的特征。这意味着每个模型可能可以划分为多个类似“layer”的结构，从而将大模型分而治之。为此，我们分析了 PAI 平台上运行的 Workload，并总结了一种通用的分图处理方法。通过图分析识别图中所有的关键节点（Critical nodes），并以它们为分图间隔点，将整体模型划分为近似线性的多个 Segment，如上图中标出的三个 Segment，每个 Segment 都包含若干 Cone 结构。这里同样存在一个问题：每个 Segment 需要保留多少种候选策略？答案是取决于其包含的所有 Critical nodes 的切分策略数量组合。在做法上，通过对每个 Critical node 枚举切分策略，以其内部的 Cone 为基本单元，以最小通信代价为目标，使用整数线性规划（ILP）确定该 segment 的整体切分策略。

关于 critical nodes 的识别，可以参考我们论文中的描述。

3. 整体 Graph

Segment 之间的近线性拓扑，天然适合动态规划求解。上述为每个 Segment 确定候选切分策略后，TePDist 以最小通信代价为目标，使用动态规划（DP）对整体 DAG 完成确定唯一策略。

Pipeline Strategy

TePDist 在划分 Pipeline stage 时的一大特点是不需要将 DAG 排成线性序列。并且，TePDist 将 Stage 划分建模成整数线性规划（ILP）问题求解。划分的基本原则是尽量保证各个 stage 计算量均匀的情况下，追求最小的通信切面。因为，DAG 中的每个节点都可以找到它的祖先节点和后继节点，这种前驱后继的依赖关系可以被描述成 ILP 问题的线性约束，如下图中蓝色和红色所示。在策略搜索时间的压缩方面，TePDist 向用户提供了用于控制 stage 计算均匀性比率的控制接口。Pipeline 建模的具体 formulation 可以参考我们的论文。

分布式执行引擎

TePDist 自己定制了执行引擎，通过将 HLO computation 实例化并组织成 Task Graph，且采用静态调度执行。

Task Graph

为了能够组合各种各样的并行模式，我们抽象了 Task Graph。当将并行策略应用到 HLO 上时，被拆分或切分的多个 HLO computation 需要组合成完整的执行逻辑，为此我们抽象出了 Task Graph，它由一系列具有连接关系的 Task Node 组成。在 Task Graph 中，我们约定 Source 和 Sink 为起始和终止节点，其他 Task Node 均为各种 HLO computation 的实例。举例说明，对含有 Gradient Accumulation 的训练任务，原始的 HLO computation 被切分为 Compute Gradients、Gradient Accumulation 和 Apply Gradients 三个部分。那么由此创建三种不同的 Task Node，然后组合成完整的执行逻辑。下图展现了三种切分策略时，Task Graph 的具体形态。

静态调度计划

在构建执行计划期间，TePDist 在 Task Graph 上制定静态调度计划。相比于执行期通过动态调度执行 Task Graph 来说，静态调度具有更好的性能。一方面，静态调度在编译期间事先把调度顺序确定好，而动态调度十分依赖执行期的管控节点，从而可能因存在中心管理节点而成为性能瓶颈。但静态调度一经确定，每个 worker 就可以按部就班执行，对中心管控没有依赖。另一方面，静态调度会让显存使用呈现稳定状态。当执行 Pipeline 并行策略时，1F1B 的调度策略能够使显存及时释放，从而降低峰值显存用量。而动态调度就不能保证完全呈现 1F1B 的调度顺序。

其他

在执行引擎方面，我们还做了如下工作：

1. 分布式初始化

TePDist 直接对 Sharding Tensor 的各个分片做初始化。为了保证切分后的初始化结果与切分前完全一致，应该给予不同分片于相同的初始化种子，但不同的随机数起始生成状态。在 TePDist 中，拥有 C++层实现的 Sharding Initializer，它可以在不同 Tensor 分片上，令随机数生成器 Skip 到正确的起始状态进行初始化，并且通过多线程为初始化过程并行加速，这对超大模型十分有用。

1. 通过 NcclContext 管理复杂的通信域

复杂的分布式策略可能包含集合通信和点对点通信，而每次通信涉及到 Device 可能完全不同，这使得通信域的管理变得复杂。TePDist 在首次运行开始前会对所有的通信指令和节点做一次收集，然后依次建立对应的通信域，并存入 Cache 中，以在适当的时机复用。

1. Task Graph 的执行

因为 Task Graph 是我们提出的新抽象，因此需要对 Task Graph 运行时做全面的管理，包括：

• Task Node 的多线程异步 Launch

• Input output alias

• Task Graph 级别的垃圾回收机制

• 协调多机共同执行 Task Graph

性能实验

我们在以下两个商用平台上对 TePDist 做了性能实验，所有实验均采用 FP32 进行。

• M8 平台：8 * V100-SMX2-32GB GPU w/NVLink2, 2 * Xeon (Skylake) CPU 48C 2.5GHz, 768GB DDR4-2666, 1 * 100G RoCE interconnect.

• S1 平台：1 * V100S-PCIE-32GB GPU, 2 * Xeon (Cascade Lake) 52C 2.5GHz, 512GB

DDR4-2666, 1 * 100G RoCE interconnect.

模型扩展实验

（表格格式根据平台需要修改，也可直接用图片）

我们在 GPT 和 MoE 模型上做了 SPMD+Pipeline 混合策略的模型扩展性实验，并与 Alpa 和 Megatron/DeepSpeed 进行了对比。

下面两个表格列出了 GPT 和 MoE 的不同版本配置，参数量均逐行递增。

将统计的计算量（TFLOPS）作为扩展性衡量指标。在单机单卡到两机 16 卡的算力资源上逐步扩展模型，从实验表现看，GPT 和 MoE 在扩展性上表现较好。单精度计算能力方面，V100 的理想计算峰值为 15.6TFLOPS。TePDist 能够使 GPT 和 MoE 分别达到峰值能力的 62%和 58%，和自动分布式的框架 Alpa 相比，TePDist 能够提供基本相当的性能，在某些情况下还会有性能提升。

通用性实验

我们还提供了其他模型的 benchmark，来证明 TePDist 在自动化方面的通用表现。对 VGG-19，DNABert 和 UNet 模型做数据弱扩展实验。其中，将 VGG-19 的分类器扩展到百万分类级别，实验配置如下。

对 Wide-ResNet 模型做模型弱扩展实验，实验配置如下。

在 M8 平台上，TePDist 均表现出接近理想线性加速比的性能。在 VGG-19 模型实验中，TePDist 找到了将最后一层大规模分类器做模型并行的策略。

不同优化级别

针对 GPT-3 模型，我们对比了 SPMD 两种优化级别下的策略搜索完成时间。O2 级别表示带有 Heuristic 的三层搜索算法，O3 表示不分层的搜索算法。实验表明，在小模型上，O2 级别的优化由于三层的划分以及对每个层级多次使用 ILP 求解，在搜索效率上并不占优势。但随着模型的增大，其搜索效率显著提高。在超大模型上，O2 级别的搜索表现出很大的优势。

RoadMap

后续我们计划定期优化 TePDist 系统，并不断完成产品化工作。

1. 继续优化现有的执行引擎

2. 支持更多样的并行策略

3. 提供更丰富的前端支持

4. 自动化显存优化技术

开源地址：https://github.com/alibaba/TePDist

希望各位感兴趣的开发者们加入我们，一起打造更快更好的自动分布式系统！

TePdist 开源项目钉群