本文为稀土技能社区首发签约文章，30天内制止转载，30天后未获授权制止转载，侵权必究！

近年来，跟着Transformer、MOE架构的提出，使得深度学习模型轻松打破上万亿规划参数，传统的单机单卡形式现已无法满足超大模型进行练习的要求。因而，咱们需求依据单机多卡、乃至是多机多卡进行散布式大模型的练习。

而运用AI集群，使深度学习算法更好地从很多数据中高效地练习出性能优秀的大模型是散布式机器学习的首要目标。为了完结该目标，一般需求依据硬件资源与数据/模型规划的匹配状况，考虑对核算使命、练习数据和模型进行区分，然后进行散布式存储和散布式练习。因而，散布式练习相关技能值得咱们进行深入分析其背后的机理。

下面首要对大模型进行散布式练习的并行技能进行解说，本系列大体分九篇文章进行解说。

  • 大模型散布式练习并行技能（一）-概述
  • 大模型散布式练习并行技能（二）-数据并行
  • 大模型散布式练习并行技能（三）-流水线并行
  • 大模型散布式练习并行技能（四）-张量并行
  • 大模型散布式练习并行技能（五）-序列并行
  • 大模型散布式练习并行技能（六）-多维混合并行
  • 大模型散布式练习并行技能（七）-主动并行
  • 大模型散布式练习并行技能（八）-MOE并行
  • 大模型散布式练习并行技能（九）-总结

之前的文章现已对多种并行技能进行了解说，如：数据并行、张量并行、流水线并行以及多种并行办法组合运用的多维混合并行。可是大模型的散布式练习是一个非常杂乱的问题，现在的绝大多数的散布式练习体系，都依靠用户人工重复尝试以及体系专家经验来进行布置，造成严峻的资源运用功率低下的问题。因而，主动并行计划应运而生。

简述

主动并行的目标便是用户给定一个模型和所运用的机器资源后，可以主动地帮用户选择一个比较好或者最优的并行战略来高效履行。可以说，主动并行是散布式并行的终极目标，它可以解放工程师去手动设置散布式并行战略。

而主动并行可以分为全主动并行和半主动并行形式。

  • 半主动形式下用户可以依据自己需求指定某些tensor和operator的切分办法。如：Mesh-TensorFlow、GShard、GSPMD 等说到的主动并行切分计划。
  • 全主动形式下一切 tensor 和 operator 都由结构自适应选择最优切分战略。如：OptCNN、Flexflow、Unity、Alpa等说到的全主动并行切分计划。

现在，很多的通用AI结构（如：PaddlePaddle、OneFlow、PyTorch、MindSpore、TensorFlow、JAX等）都对主动并行(全主动或半主动)进行了完结。

下面将共享一些典型的散布式练习主动并行计划。

Mesh-TensorFlow

布景

在深度学习中，由于数据量和核算量的庞大，往往会运用到散布式核算。而最常用的散布式形式是SPMD(Single-Program-Multiple-Data)，即数据并行，这种形式相当于在数据的batch维去做拆分；然后，进行并行。Mesh-Tensorflow对这种形式做了泛化，即除了batch维外的其他维度也可做并行。

SPMD 的 batch 切分

首要，回顾下之前的数据并行，每个设备上都有悉数模型参数的备份，在每一次迭代中，数据首要被切分分发到各个设备上；然后，各个设备分别进行核算，得到的梯度再经过AllReduce进行聚合，然后再更新参数。

Mesh-tensorflow 的切分

散布式依靠的是数据分发和聚合，这点上面解说的batch切分也是，但 Mesh-tensorflow 做了更泛化的笼统。

  • 让Tensor的每一个维度都有姓名。比方：假如每个样本都是一个向量，那么每次练习的输入x的维度便是[batch, d_io]。
  • 类似的，把处理器集群也表明成一个矩阵，比方：一个二维的结构，表明成[rows, cols]。
  • 界说一个computation layout，这个layout是从tensor维度到集群维度的一个二分图映射。例如，上面的batch切分可以表达为[("batch", "all_processors")]。

Mesh-tensorflow 完结

每个操作都经过并行核算和 collective communication 来完结，这儿，咱们介绍几个 Mesh-Tensorflow 中比较重要的操作。

  • Component-wise Operations: 所谓的component-wise，便是指输入和输出的维度相同。这一类的操作可以直接散布式的进行。
  • Reduction(reduce_sum, reduce_max, etc): Reduction操作是指会消减维度的操作，这一类操作可以先在每个切片上操作，然后用MPI-allreduce来聚合。
  • Einstin Summation(max multiplication, etc): Einstin操作是一组矩阵核算的统称，在 TensorFlow 中被完结成了一个可以装备的API，装备的办法便是用维度的姓名来表达核算，这点其实和 Mesh-Tensorflow 殊途同归，所以可以很便利的完结。相同的，完结的办法便是先本地核算，然后再 MPI-AllReduce 。
  • Reshape: Reshape尽管简略，可是在散布式环境下却需求网络通讯才能完结，不同的reshape需求的操作不同，涉及到的MPI通讯包含MPI-allgather，MPI-alltoall等。

小结

Mesh-Tensorflow 界说了一套DSL语法，用于描绘模型的维度和布局，你用它重写你的整个Model后，它主动帮你把模型和数据切割到多个TPU上。

另外，Mesh-Tensorflow 没有完结并行的卷积操作，因而，只适合 Language Model 这个领域。

除此之外，需求用 Mesh-Tensorflow 的语法重写你的整个模型，细心考虑维度，不只作业量大，一起对代码侵入性强。

不同的 layout 会带来不同的性能，因而，可以考虑主动查找最优的layout，但 Mesh-Tensorflow不支撑。

GSPMD

经过扩展模型可以进步模型精度，扩展模型的运用范围。但这些模型往往需求在多个device上练习，产生了一些并行练习需求，如：数据并行（切割练习数据）、流水线并行（切割核算图），张量模型并行（切割每个模型层的权重和核算）。而 GSPMD 提出了一种依据tensor sharding annotations的体系，以一种一致的办法去表明不同的并行战略，包含上面说到的办法以及一些新的并行办法，如： image spatial partitioning（一种沿空间维度切割图画输入数据的技能，它有助于在内存容量有限的设备上拟合大型图画数据）和 weight-update/optimizer-state sharding（对数据并行的一种增强）。

GSPMD 简介

上面说到 GSPMD 依据tensor sharding annotations的体系，以一种一致的办法去表明不同的并行战略。

尽管流水线并行对图进行了区分，而不是对单个运算符/张量进行了区分，但 GSPMD 依然可以在一个简略的包装库的协助下完结，该包装库将流水线区分简化为一个张量/运算符区分问题。

GSPMD 有满足的灵活性来表达这些办法的组合，例如：不同的层可以用不同的办法进行分区，不同的办法可以在同一层中进行组合。

GSPMD 别离了机器学习模型编程和并行的问题。它答应用户用巨大的张量编写程序，就像有一个单一的巨大设备相同。然后，用户可以在一些地方刺进注解，指定张量怎么在设备间散布；GSPMD将在编译器pass履行，在整个核算图上完结分片标准，并将其转化为数学上等价的并行核算，在每个设备上运转。

这使得用户可以专心于模型的树立，而不是分片的完结，而且可以轻松地将现有的单设备程序移植到更大的规划上运转。为了实验不同的分片战略，只需注解重新装备即可。

GSPMD 解决了将主动分区运用于出产模型时的几个实际问题：

  • 为每个分区生成一个程序会大大添加编译时刻，所以 GSPMD 为一切分区生成一个程序。这一特性被称为单程序多数据(SPMD)，关于扩展到数以千计的分区至关重要。
  • GSPMD 支撑不均匀切割的维度，使任何张量都可以在恣意设备网格上进行切割。为了便利开发，加速器在编译时要求静态已知的形状，这一般是一个实际的束缚。尽管支撑不均匀的分片，GSPMD 与这种束缚是兼容的。
  • GSPMD 作为 Production ML 编译器 XLA 的一个扩展来完结。该完结涵盖了 XLA 中的悉数运算符，包含那些具有杂乱语义的运算符，如卷积。XLA 是对多个结构(TensorFlow，Jax，Pytorch和Julia)和硬件平台(CPU，GPU和TPU)的一致笼统，使 GSPMD 可以重复运用。
  • GSPMD支撑嵌套的并行形式；在per-operator层面，这意味着不同类型的维度可以在正交的device mesh中进行区分。GSPMD 现已为这种嵌套形式开发了一种递归办法，最大限度地进步了 GSPMD 的通用性，而不需求过多的手写分片规则.

GSPMD 张量分片和主动完结

GSPMD 为张量分片界说了一套直观且通用的表明。遵从别离规划的理念，GSPMD 有两个独立的编译器转化：sharding completion 和 per-operator partitioning。

GSPMD 具有一种机制，答应高档用户经过在子图中输入手动分区形式来准确操控子图的分区办法。 在这个子图中，用户用分片巨细的形状编写程序； 在子图之外，程序依然由编译器主动分区，而且有专门的转化节点在形式之间进行切换。

为了让 GSPMD 依然可以对其他维度进行分区以完结数据或层内模型并行，GSPMD 扩展了手动形式以支撑类似于部分复制的子组，即子组内的设备手动分区，而子组之间的设备主动分区。 在这种状况下，用作流水线阶段（stages）的设备组是手动子组。

GSPMD 依据有限的用户注解主动完结每个张量的分片。它是作为 XLA 中的编译器pass完结的。

GSPMD SPMD 分片

在完结 Partitioner 时有两个选项：

  • 为每个Partitioner创立自界说程序（多个程序多份数据，MPMD）
  • 创立一个程序适用于一切Partitioner（单个程序多份数据，SPMD）

GSPMD 选择 SPMD 是由于咱们的目标是扩展到数千个 Partitioner，而在 MPMD 中，编译程序会变得非常慢。编译时刻是一个重要的可用性问题，由于现代ML结构一般包含JIT优化和编译，特别是关于那些针对自界说加速器的结构。并行化编译或许不简略，由于不同程序中的操作符或许需求全局调度以维护正确的通讯次序。

但在 SPMD 中完结Partitioner相同会给出产ML编译器带来了独特的应战。因而，GSPMD针对SPMD分区所面临的应战提出了一系列解决这些问题的技能。

小结

总归，GSPMD 提出了一种依据编译器的、主动的、通用机器学习并行体系。它是一种半主动并行，用户手动装备部分的并行操作，然后它会对并行战略进行传播得到完结的并行战略。

Flexflow

布景

现有的深度神经网络练习一般需求运用数据并行或模型并行。可是这些战略在并行程度上一般无法到达最优。因而，本文界说了一个 DNN 并行战略查找空间（SOAP），其中，包含在Sample、Operator、Attribute和Parameter维度中并行 DNN 的战略；一起，本文还提出了 FlexFlow，这是一种深度学习结构，它运用 SOAP 空间的引导随机查找来寻觅针对特定的并行机器的快速的并行战略。

为了加速这种查找，FlexFlow 引入了一种新颖的履行模仿器（execution simulator），它可以准确猜测并行战略的性能，而且比之前直接履行每个战略的办法快三个数量级。

SOAP 查找空间

下面来看看 DNN 并行战略的 SOAP 查找空间。为了跨设备并行化 DNN 算子，咱们要求每个设备核算operation输出张量的不相交子集。 因而，咱们经过界说 oi 的输出张量怎么分区来对 operation oi 的并行进行建模。

下图展现了一些算子样例的并行维度：

大模型分布式训练并行技术（七）-自动并行

下图展现了1D卷积并行装备的样例：

大模型分布式训练并行技术（七）-自动并行

下图展现了一个矩阵乘法运算的并行装备示例：

大模型分布式训练并行技术（七）-自动并行

总归，SOAP 维度的切分，是针对op的output tensor来切分的，选择了output tensor的多个维度：

  • Sample：表明 input 的 batch 维。
  • Attribute：表明 tensor 的属性维，例如：height/width。
  • Parameter：表明 tensor 的 param 维，例如：in-channel/out-channel。
  • Operator：表明 op 之间的切分维度。

尽管把 tensor 分红了多个维度，实际上都是归于 tensor 本身的维度。

FlexFlow 全体结构

FlexFlow 依据核算图和设备拓扑主动寻觅并行战略。与现有结构相比，FlexFlow有两个优势：

  • 可编程性。 关于在具有深度设备拓扑的集群上运转的具有杂乱核算图的 DNN 运用程序，运用程序开发人员乃至领域专家都很难手动规划高效的operation分配。 FlexFlow 担任寻觅高效的并行战略，并供给更高效的编程接口。
  • 可移植性。 针对一个集群进行微调的并行战略或许在其他集群上体现不佳。 FlexFlow 的查找办法会主动为每个硬件装备选择有效的战略，而无需更改运用程序。

FlexFlow 的全体结构如下图所示，其中：

  • Operator Graph：核算图的描绘。包含op作为node，tensor作为edge。
  • Device topology：描绘实际设备的topo联系，device作为node，connection作为edge。
  • Execution Optimizer：FlexFlow的中心部件，用于查找最优的split计划，下方是一个运转时（Distributed Runtime），用于履行split计划。

大模型分布式训练并行技术（七）-自动并行

履行模仿器（Execution Simulator）

履行模仿器是FlexFLow中比较中心的部分，担任对提出的战略做评估，得到候选者的性能数据。

这儿为了进步评估的速度，没有运用直接履行的办法，而是用模仿履行。仍是正常去构建履行timelines，可是需求在device上履行时，直接从上一次履行相同input-size的数据中取得履行时刻，这样降低了全体的履行时刻。这儿是假定op针对相同input-size的履行时刻基本不变，而且跟input-data无关。在大多数模型中，这个假定都是成立的。

  • 输入：算子核算图G，设备拓扑结构D，并行战略S
  • 输出：履行时刻
  • simulator的重要假定：
    • 1）每个task的履行时刻都是可猜测的，波动小，与input tensor的内容无关。
    • 2）不同设备之间传输数据的时刻为数据巨细/带宽。
    • 3）每个设备按照FIFO的次序履行使命（GPU便是这样的）。
    • 4）每个设备在完结一个使命后，只要下一个使命的数据准备就绪就立刻开始履行下一个使命，overhead可忽略不计。

为了模仿一次履行，模仿器首要树立一个Task Graph，然后运转模仿算法。

使命图（Task Graph）：

构建使命图时，每个op对应的split都会变成一个normal task。task之间的数据通讯作为communication task。

graph的edge表明的是task之间的依靠联系，即核算先后联系，而不是数据流方向。

在构建使命图的时候，就把每个task的execTime填入了。normal task 的 execTime 是在 device 上多次履行的平均耗时，这儿 cache 之后，会一直运用。communication task 的 execTime 是用 tensor size / bandwidth 得到。

模仿算法类型：

  • 全模仿算法 ：首要用 Dijkstra 算法遍历，一切使命都被放到一个队列里，出队列的次序是按照ready time 的增序。该算法终究回来一切使命中最慢的一个履行完所需时刻。
  • Delta 模仿算法：运用一种 MCMC 查找算法，每次只改动一个 op 的区分办法。这种状况下，前后两个战略的时刻一般没有改动。Delta 模仿算法只重新模仿改动终究成果的 op。

关于相同的使命图，full和delta的模仿算法会给出相同的成果。

履行优化器（Execution Optimizer）

履行优化器以运算符图和设备拓扑作为输入，并主动找到有效的并行化战略。

  • 输入：算子核算图G，设备拓扑结构D
  • 输出：最有效的并行战略

问题笼统为最小化总履行时刻，这个办法避免了平衡履行时刻和通讯时刻二者的问题。

FlexFlow 运用模仿器作为预言机，将并行优化问题转化为cost最小化问题，即最小化猜测履行时刻。 这种办法的首要长处是，它避免了显式地编码相互依靠的优化之间的权衡（例如：削减数据传输与平衡作业负载散布），而仅仅专心于最小化运用程序的全体履行时刻。

经过从最小全体履行时刻找到最佳并行化战略是 NP-hard 问题。 或许的战略数量与运算符图中的op数量成指数联系，这使得穷举查找空间变得困难。

为了找到低本钱战略，FlexFlow 运用本钱最小化查找程序来启发式探究空间并回来发现的最佳战略。

FlexFlow 运转时环境

现有的深度学习体系（例如 TensorFlow 、PyTorch 、Caffe2 和 MXNet ）仅支撑经过数据并行在batch维度中并行操作，在这些体系中，并行其他维度或多个维度组合的操作并非易事。

为了支撑运用并行空间中界说的任何战略并行 DNN 模型，本文在 Legion（论文：Legion: Expressing locality and independence with logical regions） 中完结了 FlexFlow 散布式运转时，这是一种用于散布式异构架构的高性能并行运转时，并运用 cuDNN 和 cuBLAS 作为处理 DNN 算子的底层库。

本文运用 Legion 高维分区接口来支撑可并行维度的恣意组合的并行操作，并运用 Legion 的细粒度操控机制来操控每个算子粒度的并行。

FlexFlow 运转时与现有体系之间的首要区别在于，FlexFlow 支撑以可并行维度的恣意组合并行算子，并以单个算子的粒度操控并行。

小结

总归，FlexFlow 最中心作业便是提出了 execution simulator 来完善 cost model 。

Alpa

布景

现有的一些计划要么被束缚在单个并行办法 (PipeDream)，要么依靠于对模型和集群标准的强假定 (DAPPLE，Tofu)。一起，主动混合并行的查找空间较杂乱，多并行战略的完结不够灵活。除此之外，不同的并行技能是有不同的带宽要求的。

因而，Alpa采用在不同的体系层次运用不同的并行技能，提出了的算子间和算子内并行主动并行计划。

Alpa 技能原理

Alpa提出的算子间、算子内并行区分办法，经过”是否切分了tensor的维度”来区分不同的并行。

  • 算子内并行（intra-op）：切分了tensor维度的并行办法，包含数据并行和算子并行（即张量模型并行）。
  • 算子间并行（inter-op ）：不切分tensor，仅仅把子图进行不同的摆放散布，包含流水线并行。

算子内并行可充分运用带宽，切分带来的通讯基本归于高效的调集通讯。而算子间并行若切点寻觅的适宜，则通讯较小，但同步版本的战略无可避免的会引来 Bubble。所以，可以运用集群的非对称特性，将算子内并行映射到高带宽互联的Devices上；将算子间并行映射到低带宽互联的Devices上。如此组合，就能释放更大的算力，Alpa会主动探究这些战略及组合状况。

Alpa 先经过动态规划（DP）来决议模型怎么切分红 stage，每个 stage 能分到哪些卡。然后在每个 stage 内部，再经过整数线性规划（ILP）的办法来决议每个 op 是怎么切分到这个 stage 的多个卡上，这是一个主动优化的过程。

大模型分布式训练并行技术（七）-自动并行

在 Alpa 开源仓库中，也供给了依据 OPT 大模型进行主动并行的微调事例。

Alpa 的履行过程

Alpa 高度依靠 JAX，它魔改了 XLA （JAX 底层经过 XLA 履行）中的 GSPMD，拿到 XLA 的核算图后，主动对 op 进行切分，生成对应的程序，在每个 worker 上履行。

Alpa 的创新之处

旧有的计划往往焦点在 inter-op，intra-op 和主动并行战略查找的一个或者两个点，而 Alpa 兼顾了一切；比方：在 GShard 中提出了 intra-op 的办法，GPipe 提出 inter-op 的办法，Megatron-LM v2 则经过结合 inter-op 和 intra-op 的办法，经过人工指定的并行战略来支撑散布式练习 GPT 模型。微软 DeepSpeed 提出的 ZeRO 技能企图经过主动的战略，经过多个层级步骤，来优化数据并行中的显存运用。而 Alpa 首要做 inter-op 的主动切分，然后用 intra-op 的层级调度办法，然后到达兼顾一切的优化战略。可以说，Alpa 是当今停止主动并行的集大成者，后续作业要想打破它相当困难。

大模型分布式训练并行技术（七）-自动并行

总结

本文共享了一些典型的散布式练习主动并行计划，可是，在一些业界闻名大模型的预练习过程中，这些计划现在还用得很少。可见，当下的主动并行由于其相对杂乱还不是很老练。

码字不易，假如觉得我的文章可以可以给您带来协助，期待您的点赞收藏加重视~~