在 3.25 日的 MegEngine Meetup 中，旷视研究院周亦庄讲师分享了《利用 MegEngine 分布式通信算子实现简单的并行训练》。
直播回放链接：利用 MegEngine 的分布式通信算子实现简单的并行训练 – MegEngine Meetup No.2_哔哩哔哩 (゜ - ゜)つロ 干杯~-bilibili
分享内容次要分为四个局部：
1. 介绍 MegEngine 的分布式通信算子；
2. 简略参数并行，用于相熟模型并行的一些基本概念；
3. 层内模型并行；
4. 层间模型并行和流水线并行，同时介绍了如何实现一个简略的 GPipe。
以下为该分享的文字实录，Enjoy~
一、背景
并行训练是发展深度学习钻研和业务十分重要的一环，很多根底钻研都须要大规模的计算集群甚至是超级计算机来实现。比方，像咱们晓得的 DeepMind 下围棋的 AlphaGo，还有 OpenAI 的 1750 亿 (175 billion) 参数的超大语言模型 GPT-3，最近 OpenAI 还搞了一个 CLIP 和 DALL-E，他们都是用十分大的集群来进行分布式训练的。而因为旷视研究院有 Brain++ 这个分布式的计算平台，所以咱们也有很多优良的成绩。大模型在各类视觉和语言工作上相比于小模型都有显著劣势，所以最近的一种趋势是模型规模、数据规模越大越好，“大即正义”，因而更须要大规模的并行训练。
并行训练，一方面能够调动上百甚至上千块 GPU（图形处理器，又称”显卡”，简称”卡”，是深度学习最常见的计算设施）进行训练，第二局部也是依据业务或模型的特点，咱们能够设计出最高效的并行模式。这是我明天讲的并行训练的一个现实意义。
先来讲一下深度学习当中有三种比拟罕用的并行模式，三种并行模式的关系用上面这张图就能够表白分明。
第一种（层内模型并行），是利用矩阵乘法人造的并行个性，把每层（比方全连贯层或卷积层）外部的矩阵乘法计算给拆开，体现为沿着输出 / 输入 通道（channel）拆开进行分组计算，这就叫层内模型并行。
第二种（层间模型并行），是利用神经网络串行执行的个性，把网络依照执行程序拆开，别离放到不同的设施上进行计算，比如说咱们一个 ResNet18，它有 17 层卷积层加上最初一层全连贯层，如果咱们把前九层的和后九层的计算放到两块卡（即 GPU/ 显卡）上，它就是叫层间模型并行。层间与层内这两种模型并行形式是“正交”的，互不影响，能够同时存在。
以上说的两种并行，它的模型参数都是拆开来的，每个计算节点（计算节点是底层计算设施的一种形象，它能够是一张卡，也能够是一台或者一组 8 卡机，即装载 8 块 GPU 的计算机）只负责管理整个网络的一部分参数以及这部分参数参加的相应计算。
最初一种就是咱们最罕用的数据并行，它又是另外一个维度，在数据并行维度上，模型参数都是共享的，然而接管的数据是不一样的。通过减少计算设施，咱们能够近似线性地减少单次迭代的 batch size（批量，即训练图片的数量），从而节俭训练模型的工夫。
这三种并行维度是两两正交的，意思是在理论训练中咱们既会用到两种模型并行也会用到数据并行。小模型可能数据并行就足够了，但大模型因为参数特地多、计算量十分大，计算难以用单个 GPU 实现，这时候就要将计算拆解到不同 GPU 上，此即模型并行。
二、MegEngine 的通信算子
接下来，进入到明天要讲的正题。先说通信算子。
人类的历史它其实就是一个信息交互的历史，也就是一个通信的历史——人与人之间谈话就是通信，我明天做直播，它其实也是通信，我把信息播送给大家，这也是通信，电视和播送当然也是通信。
对于深度学习框架来说，通信是最重要的性能之一，否则数据并行和模型并行难以实现。简略来说就是我有很多个计算设施（GPU），我须要让信息在所有计算设施之间进行交互，那就须要汇合通信——汇合通信是一个求导齐备的一套通信规定。
表中列了有 8 种汇合通信算子和 2 种点对点通信算子，这就是 MegEngine 全副的通信算子。8 种汇合通信算子，形成一套求导齐备的通信的规定，它们相互各自为导数。MegEngine 提供了对通信算子的主动求导，所以和其它所有用于计算的算子（如卷积、ReLU、转置等）一样，咱们能够自在地把通信算子退出前向计算图，框架将负责对其求导。
思考到有些同学没有背景常识，咱们一一介绍一下汇合通信算子的性能。
Broadcast
Broadcast 即播送。
它示意的是数据的一个同步的过程，将一张 GPU 上的信息同步给其它所有 GPU。这在数据并行中十分有用，因为数据并行的话，每张卡下面的参数应该确保都是一样的，因而在初始化时咱们会通过 Broadcast 进行参数同步，咱们也会周期性同步一些缓存信息（buffer，比方 BatchNorm 的统计量）。
ReduceSum
第二个是 ReduceSum。ReduceSum 叫做求和或者归约，将所有 GPU 上的数据收集到一个 GPU 上并相加。
咱们方才讲的 Broadcast 和 ReduceSum 这两个通信算子是形成参数服务器 Parameter Server 的一个基石，它是核心式的，在这外面 GPU0 就起到一个核心的作用，我先把核心参数通过 Broadcast 同步给各张卡进行前传，反传后通过 Reduce 收集各张卡的梯度，进行参数更新。Broadcast 和 ReduceSum，互为导数的，ReduceSum 的导数就是 Broadcast，Broadcast 的导数是 ReduceSum。
AllReduce
咱们再介绍 AllReduce，原本 Reduce 是归约到一张卡上，AllReduce 则是归约到每一张卡上。它即能够了解为 Reduce Broadcast 的组合，即我先 Reduce 到一张卡，而后再 Rroadcast 到所有卡；也能够了解为每张卡都同时调用了 Reduce，AllReduce 它的导数就是 AllReduce 自身。
只管只用 Reduce 和 Broadcast 就能够实现 AllReduce，然而 AllReduce 的高效实现（即 Ring-AllReduce）才是形成古代 分布式深度学习框架的基石，它的通信工夫根本不随 GPU 数量的减少而减少，因而能够高效地实现分布式训练的规模化。在数据并行中，咱们用 AllReduce 将所有梯度求和，并用于模型参数更新。
Gather
Gather 简略来说就是把每张卡上不同的信息都给收集过去，并沿着第一维相连（Concatenate）。
AllGather
AllGather 就是全收集，和 AllReduce 相似，它能够了解为 Gather 后接 Broadcast。
AllGather 是咱们层内模型并行当中的一个很重要的操作，因为你的参数在不同的卡上，你的数据也在不同的卡上，在进行模型并行的时候，我须要把数据或者参数都收集起来放到一张卡上能力进行接下来的计算，这就是 AllGather 的作用。
AllToAll
AllToAll 也是层内模型并行中常常用到的一个操作，特地是在模型并行和数据并行进行切换的时候，它实质上对一个矩阵进行了转置，咱们前面在具体利用中会进一步阐明。AllToAll 的导数是它自身。
Scatter 和 ReduceScatter
最初 Scatter 和 ReduceScatter 合起来讲，Scatter 就是散发，它将一张卡上的数据拆分给各张卡，它和 Gather 互为导数。
ReduceScatter 能够了解为在散发之前先进行了求和，它和 AllGather 互为导数。
三、简略参数并行
介绍完 MegEngine 的通信算子，咱们来理解它们如何应用。首先，让咱们从简略参数并行开始，它只波及 AllGather 这一通信算子。
简略参数并行是怎么一回事？咱们先用一个简略的全连贯层（即矩阵乘法）来回顾一下数据并行——数据并行中，W 是咱们的模型（即咱们的权重 weight，每张卡领有一份同样的拷贝），x 是数据。数据并行要求咱们将数据均匀拆分到每张卡上，2 卡拆 2 份，即 x0 和 x1，4 卡则拆成 4 份，依此类推，各张卡别离进行矩阵乘法计算，失去对应的后果 y。
简略参数并行实质是数据并行的优化？咱们不用在每张卡上都放残缺的模型，而是只放局部模型，只有在咱们须要（即前传）的时候，把扩散在各张卡上的参数收集（AllGather）起来参加计算。
如何实现？咱们在做矩阵乘法操作之前，先对参数进行 AllGather，从各个节点上收集被咱们拆开的参数，AllGather 当前每张卡都有全副的权重了，计算就变得和数据并行截然不同的。所以，简略参数并行的外围操作就是 AllGather，实质用通信来节俭显存。
为什么能节俭显存呢？咱们当初把整个求导过程也画进去了，咱们晓得在训练一份参数的时候，它其实是会占掉三份显存——参数一份，梯度一份，优化器的 momentum 一份，所以一个参数量 1 million 的模型，如果咱们应用数据并行，会占用 3 * 4G = 12G（1 million fp32 类型的数据占用 4G）的显存，那咱们一张 2080ti 就齐全没有显存能够用于训练了。
咱们再来钻研一下这张图，咱们在前传的时候做了一次 AllGather，在反传的时候，咱们晓得 AllGather 的导数是 ReduceScatter，所以，它反传的时候会进行一次 ReduceScatter。这和数据并行不一样，数据并行前传不须要通信，反传须要进行 AllReduce 这是他们的区别。
咱们用 MegEngine 写了一套数据并行和简略参数并行的代码，它们有三个不同：
- 一个不同是它们的前传是不一样的——左边（简略参数并行）就是要做一次 AllGather；
- 还有一个不同就是他们在参数初始化的时候。在数据并行中咱们须要参数同步，所以咱们要 Broadcast，然而在简略参数并行外面，咱们须要的是参数散发，所以用 Scatter，就把它们给散发进来。
- 最初一个不同就是在求导的时候，求导的时候在数据并行当中咱们须要进行 AllReduce（MegEngine 应用 AllReduce callback 来反对数据并行），然而在简略参数并行外面不须要进行 AllReduce，主动微分器会负责反传时正确调用 ReduceScatter。
四、层内模型并行
层内模型并行在原理上更加简单。咱们方才讲的参数并行，它其实是一种层内模型并行的一种特例，因为它十分的简略，只须要对参数进行 AllGather。实际上咱们的层内模型并行还有多种不一样的实现。
上图给出了残缺的矩阵乘法、数据并行和两种模型并行的实现。
咱们晓得矩阵乘和卷积神经网络中的卷积层（卷积层能够视为对 channel 维度进行的矩阵乘），都人造具备并行的个性。咱们在数学意义上的矩阵乘法，每一行每一列的运算都能够独立进行，数据并行就充沛的利用了这个个性，咱们把数据进行均匀切分，各自放在不同的设施上各自做矩阵乘法，最初能够合并起来失去残缺后果。
在层内模型并行当中，咱们是把每层（全连贯 / 卷积层）的参数矩阵 W 进行切分。一种形式是按输入维度进行切分（纵切）。第二种品种是按输出维度进行切分（横切）。前者在每张卡上失去局部输入维度的对应后果；后者利用了矩阵的低秩个性（Low Rank），每张卡的后果是最终后果的低秩重量，后续须通过 AllReduce 或者 ReduceScatter 将其求和。
接下来咱们在多层神经网络中应用层内模型并行——咱们实现纯正的层内模型并行，或者和数据并行搭配应用，实现混合并行。
上图第一行是纯数据并行。数据在一开始就被切分到各张卡上，之后不须要进行替换或信息交换，因而数据并行后接数据并行不须要进行非凡操作。
第二种纯层内模型并行。首先你须要残缺样本数（batch）的输出特色“X”，最初矩阵乘进去它是残缺样本数但局部输入通道数（channel）的特色“Y”，为了后续持续进行模型并行的矩阵乘法，我必须做一次 AllGather，把“Y”沿着通道（channel）收集起来，把它再变成样本数和通道数皆残缺的“Y”，再与模型并行的“V”相乘。如果网络持续加深，那么每次矩阵乘完结都要进行 AllGather 操作。
第三种混合并行混合了数据并行与层内模型并行。咱们还是以模型并行开始，模型并行的全连贯层输入一个纵切的“Y”（即沿输入通道切分的特色 Tensor），然而咱们数据并行要的是横切的“Y”（即沿样本数维度切分的特色 Tensor），应该怎么操作？在介绍 MegEngine 通信算子的时候咱们提到一个转置操作叫 AllToAll，它能够间接把这个纵切的“Y”变成了横切的“Y”。接下来咱们就能够复原数据并行了，进行一次数据并行的矩阵乘法后，咱们还想进行一次模型并行的矩阵乘法，那就再做一次 AllGather，失去全副样本数且全副通道数的残缺特色 Tensor。把握了利用 AllToAll 和 AllGather 实现的“切换”当前，你就能够本人设计与训练混合并行的模型。
接下来咱们举例两个利用场景。
场景一：全连贯的层内模型并行
咱们来进入一个具体场景，在人脸识别工作中利用全连贯的层内模型并行。
在人脸识别工作当中，可能有百万、千万的 ID（Identity，同一个人为一个 ID），相当于要去做一个输入维度为百万 / 千万的分类工作，所以，最初这一层，分类的这一层 FC 层（全连贯层）它可能参数特地大，比如说咱们有一百万（1 million）的 ID，提取的人脸特色是一个 1024 维的向量，它们乘起来就会占用 4 个 G 显存，咱们方才提到 4G 参数的模型在理论训练中会固定占用 3 倍显存，就是 12G，个别的显卡装不下。我只能把这个全连贯给放到各张卡上，如果咱们有 8 张卡，每张卡就只会分到 1.5G，那么还是能够承受的。这个场景的特点是什么？就是人脸特色维度相比于我的参数矩阵其实十分小的，所以咱们对数据进行通信（AllGather），它的代价要比对权重进行通信（AllReduce）它的代价小得多，所以在这个场景下特地适宜做模型并行。
在模型并行下分类器 W 输入的后果 Y 的具体含意是什么？咱们晓得 Y 是竖着切分的，竖着这一维是样本（batch）维，就是它有多少个训练的样本，横着的这一维其实是 ID 维度，就是类别维，示意样本属于各个 ID 的概率，而模型并行下它只输入了一部分标签的概率。求损失函数的时候咱们往往用穿插熵（CrossEntropy），穿插熵须要全副的类别概率。没错，利用之前咱们介绍的 AllToAll 算子，咱们把输入的模型并行的概率矩阵给进行 AllToAll 转置，它就变回了数据并行的格局。（讲师注：实际上你并不需要进行 AllToAll，在分类工作的非凡场景下，你并不需要 AllToAll，因为通信代价很大，你能够籍由两次极低代价的通信来实现穿插熵的计算，然而这个超纲了，但不是很艰难，留给大家当思考题。）
咱们间接上代码。
整个过程中有三步，第一步是 AllGather，第二步进行矩阵乘，第三步进行 AllToAll。
那么上图框起来的这段代码是什么货色呢？咱们做了这么多 reshape，什么 transpose——这叫数据重排布，咱们再花 5 分钟的工夫来讲一下数据重排布是什么。
咱们 AllToAll 做完当前，失去的其实并不是咱们想要的局部数据加上全副分类的一个后果，它其实在底层的数据排布（layout）下面它不是咱们冀望的。上图是 1 个简化版本的例子，它的分类从 0-7 总共有 8 类，它的样本是 4 张人脸图片。通过模型并行，在卡 0 下面咱们失去的输入是 0-3 类的后果，卡 1 下面失去的是 4-7 类的后果。咱们做完 AllToAll 当前它变成的矩阵（0,1,2,3,10,11,12,13）并不是咱们想要的，咱们最初想要的就是 0,1,2,3,4,5,6,7，上面是 10-17，所以的话咱们必须先做一次 reshape，沿着这个方向是最外面维 0,1,2,3 数据是间断的，咱们把这里面两维（0,10,4,14）个给进行一次转置，就是转过来，最初 reshape 为想要的后果。为了当前使用方便，我简略进行了以下两个封装，下面封装叫 mp2dp，就是从模型并行变成数据并行（Data Parallelism）的一个封装，上面这个是 dp2mp，有了这两个封装当前，咱们下面的前传代码就变得简略了。
场景二：组卷积模型并行
讲完了全连贯，接下来咱们再讲组卷积（Group Convolution），
Group Convolution 在咱们的挪动端模型下面特地常见，组卷积和一般卷积它的区别就在于组卷积相当于 K 个一般卷积。比如说你有三组，就相当于三个一般卷积，然而每个一般卷积都比本人的小，你们也能够发现这个是人造并行的，上图红色的、绿色的、黄色其实能够各自做，在不同的设施上做。
下图用之前二维的示意形象一下卷积和组卷积的不同——组卷积的模型，它和卷积不一样，组卷积相当于一个稠密的矩阵乘法，它不是一个浓密的的矩阵（dense matrix）。
数据并行状况下和一般卷积一样，咱们把数据进行切分；模型并行咱们能够间接按色彩把这三个组分开，咱们第一块卡上做第一个组，第二块卡上做第二个组，第三块卡上做第三个组，对于每块卡来说，本来的组卷积计算都变成了一般的卷积操作。
如果咱们后面是一般卷积，两头要插入一组模型并行的组卷积，咱们应该怎么样从这两种数据排布之间切换？
很简略，咱们就做一次数据重排布（即 AllToAll），因为是数据并行到模型并行，所以咱们调用 transpose_dp2mp。
如果咱们有多个组卷积，他们连在一起，实际上咱们并不需要重复地在数据和模型并行间切换，咱们只须要关注头和尾。所以，咱们的组卷积在前传函数外面有一个叫 is_head 和 is_tail，咱们 is_head 的时候，咱们做一次通信，is_tail 的时候再做一次通信，咱们两头就齐全不须要通信了。
五、层间模型并行
咱们进入层间模型并行，方才的层内模型并行咱们介绍了相干原理和利用（全连贯和组卷积）。层间模型并行和层内模型并行很不一样，次要就是简略模型并行和流水线并行。层间模型并行简略来说就是把网络的前半部分、两头局部和后半局部离开（甚至分成更多份），就像一条鱼，鱼头、鱼中和鱼尾。
咱们简略来看一下数据并行和层间模型并行的比照示意图。
数据并行就是把数据切开，层间模型并行不切数据，而是把模型的前半部分和后半部分给拆分到不同的 GPU 上，这边就波及到一个问题，怎么把“Y”第一块 GPU 的输入后果，给“放”到第二块 GPU 上，这外面就须要 send 操作。MegEngine 提供了八个汇合通信算子，加上两个点对点通信算子——一个就是 send，一个就是 receive。这两个算子组成了层间模型并行的外围操作，接下来次要讲 send receive。
如果层间模型并行，咱们用一个图表来形象的话（如上图下半局部），横轴是计算工夫，随着计算推动，纵轴是咱们的计算设施（GPU），咱们发现工作之间存在依赖关系，所以 GPU 0 算完后必须做 send 操作，同时卡 1 做 receive 接管卡 0 的后果，而后进行本人的计算，算完再 send，卡 2 receive……这样能力做完一个流程。
为了不便起见，咱们这边又做了一次封装，第一个函数是把咱们进去的计算结果给发到下一个 GPU，这个函数是下一块 CPU 调用的，就是它从上一个 GPU 去给它拿出去，MegEngine 自带的 recv 不带主动的形态和类型推导（讲师注：在 MegEngine 的下个版本行将反对），因而封装的时候我也简略实现了一下。
简略模型并行
咱们间接看代码，在一般的数据并行外面，这是一个简略的 ResNet 18 的模型，它总共有 17 层卷积加上一层全连贯，在简略模型并行外面，如果它是第 1 块 GPU，它就负责第一局部的 5 层卷积，第 2 第 3 块各负责 4 层卷积，最初一块 GPU 负责 4 层卷积和最初的一层全链接。
在前传的时候先进行判断——当咱们如果不是第 1 块 GPU 的话，咱们就从后面一块卡拿数据。之后进行本人负责的卷积计算。失去后果后再次进行判断——如果不是最初一块 GPU，咱们要把我的数据给送到下一块 GPU 上，如果是最初一块，就间接 return。
咱们能够用代码来展现简略模型并行的推理和训练的后果：
在推理过程中，输出一张组（32 张）224 分 辨率的图片，前三块 GPU 输入的都是网络的两头特色，最初的 GPU 输入的是网络的预测值。在训练当中值得一提的：第一，因为是模型并行，所以咱们不须要进行 AllReduce；第二，前三块 GPU 在调用 gm.backward 时传入了一个 None，其实咱们在设计 API 的时候，backward 任何货色都能够，backward None 在这里会产生什么？因为前传有一个 send，所以主动微分的时候就会插入一个 recv，它会先期待来自上游的梯度，而后进行失常的反传。
流水线并行
咱们接下来讲流水线并行。简略的模型并行须要算完同一批次的全副的数据再给下一个批次的数据，实际上每一张卡都会有很长时间的闲暇期，它要么在等上一块卡跑完，要么实现了本人这一批的工作，在期待下一批次的数据。
如果咱们把一个批次的数据给分成很多小份的话，咱们能够让第 0 块卡先算一小份，算完当前立马送给下一块卡，而后再计算下一小份，这样子的话这个时刻卡 0 和卡 1 能够同时算，空置率就上来了。
这就是流水线并行的一个核心思想，咱们看一下它代码怎么实现。
比方在这个外面，咱们想要把一份数据给拆成 4 份，咱们用 F.split 将它拆成 4 分，而后遍历一遍这 4 份数据，如果它是第一块卡，它就拿那个数据，不然的话它会等，等着接管前一块卡的计算结果。不论怎么样拿到数据当前的事件就是进行计算，计算完当前咱们要解决计算结果——和简略模型并行一样，如果他不是最初一块 GPU，我要把它送到下一块，如果它是最初一块 GPU 的话，就间接进去返回后果。
这就是流水线并行。当然到理论场景中流水线并行的代码须要思考执行效率，没有这么简略，比如说会引入异步 send/recv，以升高等待时间。
咱们不光要推理，咱们还要训练，训练的话就波及到一个反传，在一般的模型并行当中，咱们的反传和前传时间轴是如下图所示：
咱们先前传完，再顺次反传。然而在咱们流水线并行外面，其实反传也是一个流水线的过程。然而这外面有个非凡的中央，留神一下从新前传（或重算）。如果咱们不从新前传的话，意味着咱们后面的这些两头后果都要保留着期待反传完结后能力抛弃 / 开释，这意味着我的贵重的显存又要被节约了，这样子的话咱们还不如算完就全副扔掉，因为我曾经把后果交给下一块 GPU 了，临时就不须要了。而反传时咱们还须要两头后果的时候，我大不了再重算一次（换句话说每张卡只有保留本人的输出就能够了）。重算后咱们能够失常做反传，失去对于输出的梯度，而后把这份梯度传给上一张卡。上一张卡同样执行重算、反传和发送梯度，直到所有卡都实现了梯度计算。
从新前传的操作叫做 checkpoint 或 sublinear，在 PyTorch 外面有 checkpoint，在 MegEngine 里也有 sublinear，咱们目前实现的是十分粗粒度的 sublinear，它不是两头保留几个后果重算局部就能够了，它其实是全副都重算了，这就是 GPipe。
前传还是一样的代码，如上图左侧给大家做一个参考。
反传是精妙的中央，咱们拿到 label，loss 当前看一下，第一就是咱们 GradManager，这是 MegEngine 一个十分重要的个性，就是 GradManager 能够对两头的 feature（就是两头后果）进行求导，所以咱们能够在计算过程中对两头变量进行 attach，在 GPipe 的场景下，咱们须要的是对输出的导数，所以咱们在一开始就 attach 输出数据 x，而后进行前传（或者称为重算）。如果它是最初一张卡的话，咱们就计算相应的损失，并把梯度算进去。通过 grad_to_prev_gpu，咱们把对于输出的梯度传给了上一张 GPU。后一块卡对于输出的梯度即前一块卡输入的梯度 dy。咱们通过 gm.backward(dy=grad）手动指定梯度，从而实现两头 GPU 的求导过程。这就是一个简略的 GPipe。
如果大家想试着玩一下这个 GPipe 的话，在 GitHub 下面 MegEngine Parallel Tutorial 是我写的，大家能够去跑一下玩一下。
欢送小伙伴退出咱们 MegEngine 旷视天元开发者交换 QQ 群：1029741705
框架应用相干交换或反馈，欢送拜访论坛：http://discuss.megengine.org.cn；
GitHub 我的项目地址：http://github.com/MegEngine/M…