关于深度学习:阿里云机器学习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端分为两个过程：

编译（Build Execution Plan）：在这个过程中会构建执行打算，包含：主动分布式策略的搜寻，device mesh的调配，TepDist执行引擎的构建——Task Graph的构建以及为Task Graph确定动态调度程序。架构和流程列于下图左图中；
运行（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划分为了三个层级，并对每层应用不同的优化办法进行求解。

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的候选策略。

Segment

尽管Cone构造的粒度更粗，在肯定水平上缩减了问题规模。但对大模型来说，Cone的数量仍然微小，可能须要进一步解决。事实上，大模型在构造方面具备反复重叠的特色。这意味着每个模型可能能够划分为多个相似“layer”的构造，从而将大模型分而治之。为此，咱们剖析了PAI平台上运行的Workload，并总结了一种通用的分图解决办法。通过图剖析辨认图中所有的要害节点（Critical nodes），并以它们为分图距离点，将整体模型划分为近似线性的多个Segment，如上图中标出的三个Segment，每个Segment都蕴含若干Cone构造。这里同样存在一个问题：每个Segment须要保留多少种候选策略？答案是取决于其蕴含的所有Critical nodes的切分策略数量组合。在做法上，通过对每个Critical node枚举切分策略，以其外部的Cone为根本单元，以最小通信代价为指标，应用整数线性规划（ILP）确定该segment的整体切分策略。

对于critical nodes的辨认，能够参考咱们论文中的形容。

整体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的调度程序。

其余

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

分布式初始化

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

通过NcclContext治理简单的通信域

简单的分布式策略可能蕴含汇合通信和点对点通信，而每次通信波及到Device可能齐全不同，这使得通信域的治理变得复杂。TePDist在首次运行开始前会对所有的通信指令和节点做一次收集，而后顺次建设对应的通信域，并存入Cache中，以在适当的机会复用。

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零碎，并一直实现产品化工作。

持续优化现有的执行引擎
反对更多样的并行策略
提供更丰盛的前端反对
自动化显存优化技术

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

心愿各位感兴趣的开发者们退出咱们，一起打造更快更好的主动分布式系统！

TePdist开源我的项目钉群