关于深度学习:阿里云机器学习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 开源我的项目钉群