如何基于 MindSpore 实现万亿级参数模型算法？

​​摘要：近来，增大模型规模成为了提升模型性能的主要手段。特别是 NLP 领域的自监督预训练语言模型，规模越来越大，从 GPT3 的 1750 亿参数，到 Switch Transformer 的 16000 亿参数，又是一个数量级的增加。

本文分享自华为云社区《一文带你了解MindSpore支持的万亿级参数超大模型关键技术！》，原文作者：HWCloudAI  。

前言

近来，增大模型规模成为了提升模型性能的主要手段。特别是 NLP 领域的自监督预训练语言模型，规模越来越大，从 GPT3 的 1750 亿参数，到 Switch Transformer 的 16000 亿参数，又是一个数量级的增加。

模型规模的数量级的增大，虽然取得了一定程度的性能提升，甚至产生了某些意想不到的“神奇”效果（如 GPT3），但其背后的计算开销成了最大的问题，比如 GPT3 训练使用了万级的 GPU 和数周的训练时间。如何既能利用超大规模的参数来提升模型表达和性能，又能控制计算量的较小的增加，成为了最主要的挑战之一。以 MoE 为代表的动态神经网络技术被重点引入。大脑是典型的低能耗高效率的计算模式，稀疏激活是最重要的特性。除了巨型模型在训练推理特别是训练时的计算效能挑战外，当前巨型模型的训练优化算法另一个更大的挑战是（不在此处讨论），BP 算法是当前最为可用的深度网络优化，但更理想的优化算法需要高并行、优化过程非对称、并能够在时空维度通过局部持续优化完成整体优化。

1. 传统的神经网络模型，前馈的时候，输入的 batch 中，每一个样本的处理，都将激活网络中的每一个参数参与计算。

2. 条件计算最宽松的定义，指仅激活网络中某些部分的一类算法。Conditional Computation refers to aclass of algorithms that activate only some of the different parts in anetwork. 在具体某类条件计算实现中，条件选择模式，可能按照输入的 batch 中每 sample 独立激活网络不同部分，可能按照输入数据空间上不同的部分（比如 image 不同区域或者 channel），可能按照输入数据时间上不同的部分（比如 time series 的不同 slide window 或者 video 的不同的 frame。），可能按照目标任务的不同每 task 独立的，可能按照非可学习的固定的随机分配不同的子网独立计算。

3. 对不同的输入（原始或者前层），按照一定条件，选择性的执行后续部分网络的计算，这个技术下，有一些近似或相关的技术，如：dynamic neuralnetwork(s), conditional computing, conditional activation, sparse activating,selective execution, mixture of experts (MoE), dynamic routing, …；强相关的一些模型比如 Switch Transformer 等。

条件计算的分类（广义）

1. 按照 routing 是否可学习可以分为：learnable routingconditional computation 和 unlearnable routing conditionalcomputation.

2. 按照 activation 是否不执行 non-activation 计算，可以分为：hard conditionalcomputation 和 soft conditional computation。对于 hard-mode 的条件计算，通过 tensor 挑选切分等操作，无论何种条件选择模式，不需要激活的数据将完全不参与不激活的网络部分的计算；soft-mode 的条件计算，可能仅采取将相关数据置零等方式来避免产生计算效果，但还是和不需要激活网路部分实际执行计算过程。

条件计算的主要优势

1. 计算有效，降低能耗：通过部分激活部分计算，以每样本条件激活的条件计算为例，单个样本只需要经过整个 SuperNet 的一部分参与计算。

2. 更大网络，表达更强：由于一处到多处的 Route，各处（层）的 Input 被路由到不同的子网独立计算，不同的输入的相互在各层的表达相对独立没有影响，表达能力更强，网络可以更大，但表达效率降低了。

条件计算的网络和计算形式

条件计算的网络和计算形式比较灵活，部分构建形式如：（此处省略具体模型和论文引用，参见: intellabs.github.io/dis）

1. 按照 CV 等 task 的特点，用多个独立的 CNN 作为 expert 网络，按照 task 来独立路由，尾部组合后给一个大网络。

2. 使用更复杂的 cascading 等形式组合不同层级的不同的 expert 网络。

3. 通过决策树等方法做数据变换实现路由。

4. 通过可学习的网络来选择路由。其中策略学习的损失有多种构建形式：直接使用分类等任务的主损失，对不同专家的重要性和负载构建损失作为辅助损失等等。

条件计算的路由策略

1.non-learnable/hard-mode，通过某种确定性策略，如 LSH 等方式计算路由。

2. learnable-mode，通过可学习网络计算路由。网络规模可大可小，简单的可学习路由为单层权重：G(x) = P(X*W)，G(x)为路由 Gate 函数，X 为输入， W 为通损失函数来度量的可学习路由权重，P 为某种挑选函数（如 topk, sort 等），在实际实现中，X*W 的输入与权重计算结果可能作为后续网络的输入信息的一部分，不仅仅利用 G(x)来选择路由，则需要对 X*W 的结果做归一化，更典型的形式则为：G(x)=P(N(X*W))，其中 N 为表达 Normalization 函数，如 Softmax。

条件计算的冗余策略

条件计算的冗余策略，可分为无冗余条件计算和冗余条件计算：

1. 无冗余条件计算可通过 P(.)函数的实现如 topk(k=1,…)来实现；

2. 冗余条件计算，可以多种实现形式，可以通过 P(.)函数的实现如 topk(k=n,…)，n>=2 来实现，也可以通过硬冗余模式，整个网络中支持输入的复制和多路计算实现。

条件计算的挑战

1. 路由算法对模型质量的影响

无论输入和路由权重作用的信息（X*W），是仅作为路由选择并作为后续网络单元的输入，还是直接作为后续网络单元的输入的一部分，路由算法决定了输入信息的处理流向，对模型的整体质量都有很大影响。

2. 路由(routing)/门(gate)的稳定性

随机初始化的路由/门的权重，权重自身在不断被训练调整；在前后层的网络持续训练变化，同一样本在训练的不同阶段会被分派到不同的后续网络单元中，这种动态变化过于剧烈，将严重影响整个网络训练过程的稳定性和收敛速度。

3、路由的专家样本重要性和负载的平衡性

训练阶段，每专家和样本批次中样本的关联度重要性，和每批次中样本被均衡分派到不同专家的负载平衡性，这两个指标既相关又冲突。需要分别构建损失函数作为辅助损失，来优化这两个指标。在 arxiv:1701.06538《Outrageously Large Neural Networks: The Sparsely-GatedMixture-of-Experts Layer》做了相关讨论。

关于条件计算/动态神经网络

关于条件计算/动态神经网络，更多的信息在《DynamicNeural Networks: A Survey》arxiv:2102.04906(arxiv.org/abs/2102.0490)一文中，作者对广义的动态神经网络，将各种动态网络相关的技术按照实例级、时间级、空间级做了分类。

1. Instance-wise Dynamic NN：逐实例动态，每样本独立激活不同的网络和参数（MoE 为这个方向）。Dynamic Architecture：Dynamic Depth、Dynamic Width、Dynamic Routing/MoE；Dynamic Parameter：Parameter Adjustment、Parameter Prediction、Dynamic Feature(s)

2. Spatial-wise Dynamic NN：空间级动态：图像等不同空间位置激活后续不同网络和参数。(CNN 等)：Pixel Level、Region Level、Resolution Level

3.Temporal-wise Dynamic NN：时间级动态：时序数据按时序维切分激活后续不同网络和参数。（video-frames,text-sequence, time-series, stream, ...)Text-SequenceVideo-Frames

上述为该综述论文对 Dynamic NN 的总体分类。

从超大规模网络动态网络技术支撑角度，高表达能力，低计算代价为主的来考虑分类，从两个维度对动态网络技术分类:

1. 按照在前馈计算时是否部分激活:

Hard-Dynamic：在前馈的时候，部分网络绝对不激活参与计算

Soft-Dynamic：在前馈的时候，部分网络经过 softmax 等 gate/route 后，通过张量元素置零等方式，失去表达能力，但会参与计算。

2. 按照动态激活判定算法的输入：

逐样本级：（在输入层）按照每样本的实例来决定动态网络的后续激活。

亚样本级：（在输入层）样本内时间/空间级激活不同的后续网络单元。一般深度网络，不仅在输入层会被选择性激活执行，在中间层也类似。

其中，智能平台支持 Hard-Dynamic 逐样本级的动态神经网络，能比较自然的获得网络结构大颗粒的稀疏激活，在超大模型中能实现训练和推理的高能效。

动态神经网络相比与静态结构的神经网络，在相关研究中，从效能，表达，泛化、鲁棒，可解释等方面做了大量对比研究。从智能平台通过计算成本尽量低的支持超大规模网络来提升模型性能的角度看，Efficiency 和 Representation 最为重要：

1、Efficiency：静态网络“牵一发而动全身”，每一个样本输入整个网络/所有参数都要响应，这对超大网络来取得领先效果的模型能耗挑战太大。

2、Representation: 参数量更大，表达容量更大；但 MoE 等结构在深度网络的各层特征的表达上，复用降低，每参数的表达效率更低。

实现策略

实现各种模型的带有动态路由稀疏激活的超大规模参数版本，需要分模型研究和实现。

以 Switch Transformer 为例，其参数扩展到部分在 Transformer 的 FFN 部分。其 MoE 化扩展，如下图：

​(图片来源：Switch Transformer 论文)

可见，MoE 化主要变化在需要 Expert 子网络前后增加 MoE 相关的逻辑。本文主要介绍平台上的实现。动态路由条件计算，主要包括四个步骤：路由计算、数据分派、独立计算，结果合并。

​1. 路由计算-Gate：根据输入（可以为整个网络的输入，或者前面网络单元/层的输出），在路由单元完成计算，在以 batch 内 sample-wise 的路由中，计算出每个样本要分派的后续网络路由（Mixture-of-Experts/MoE 中的专家）。

2. 数据分派-Dispatch：从输入的整体的 Tensor 中，按照路由计算的样本-专家关系，收集合并出每个专家需要处理的 Tensor。如果在固定 expert-batch 的设计中，要平衡每批训练中，分派到每个专家的样本数和专家每轮训练最大容量，由于样本输入的随机性，很难保证较为均匀的分派，对于低于最大容量的批次，对固定 batch-size 的要做 pad，对于高于最大容量的样本，可以采用延后重采样等方式。为了维护正确的输入输出关系（Input/X – Label/Y）和训练是反向传播的求导关系，实现中需要维护原始 batch 到每专家的 sub-batch 的 index 关系，在后来求导和结合合并时使用。

3. 独立计算-Expert：并发（逻辑上可以先后）调用各个专家处理对应的 sub-batch。这也是智能平台要支持的并发 API 之一。

4. 结果合并-Combine：合并每专家的结果 tensor 到整个 batch 的 tensor，并按照数据分派索引，交换到原始输入的顺序。

在主流的深度学习智能平台中，可以采用两类主要的实现策略：

张量置零：对需要分派到不同的后续网络单元（专家网络子网等），对需要分派的专家拷贝若干份 tensor，对于不应输入当前专家处理的数据维度置零。该方式在保证置零计算逻辑正确的情况下，实现简单，全张量操作，对平台无特殊要求，适用于算法研究，仅体现条件计算前序数据被动态路由到不同的后续网络单元，分析算法的效果。如果通过置零方式，该方法每个专家处理的 tensor 在 batch 维度大小是全 batch，不能节省计算量和内存使用量。

张量整理：对需要分派到不同的后续网络单元（专家网络子网等），对需要分派的专家拷贝若干份 tensor，对于不应输入当前专家处理的数据维度不保留。并维护好 sample 级的 index 在变换前后的对应关系。在分布式友好的实现中，如果专家子网为单位被划分到不同的计算节点，那么专家网络的实现最好从子网级的平台对象继承后实现，比如：MindSpore 中的 mindspore.nn.Cell。详细实现细节参见后续技术实现章节。

核心代码

核心代码：路由计算、数据分派、独立计算，结果合并

参考代码采用 MindSpore 示意实现。(注：importmindspore as ms)

Mixture of Experts 的核心逻辑，对输入 I，经过 routing_network(最简单*W 即可)，然后 topk(若变种算法需要 gate 权重则需要 softmax，否则可不)，然后用 tensor 的操作(可按照 batch)选择出每个 subnetwork/expert 的张量。

为方便调试，采用了规模极小的非随机的确定数值构造输入和路由权重，路由网络采用简单的 X*W。

1、路由

​当上述输入 5 行（仅 3 类，希望分派给 3 个专家）样本，和 Gate 权重做矩阵乘后，可以明确算出每个样本要分派的专家。可以用 matmul，也可以类似 gates_weighted = einsum('bd,de->be', [data_inputs, gate_weights])第一轮矩阵乘的结果为：

输入和权重乘法，在 python 中可以采用 @，也可以采用 matmul，也可以采用爱因斯坦求和简记忆法函数 einsum。当是简单的矩阵乘的时候，采用 einsum 在计算图编译的时候实际会拆分成多个算法，性能并不好；但当输入和权重超过 2 维，需要以 batch 维固定做路由计算的时候，使用 einsum 可以编程实现很简单。

2、分派

条件计算的分派，主要逻辑是根据路由网络的输出，为每个样本计算出 top-k 的专家。其实现可以通过 topk 函数实现。由于 top 选择 score 可作为后续网络单元的输入信息（含路由的信息），所以一般要对路由输出做 softmax 做归一化。

按需计算 1：all-N 专家之间的归一化权重(please refer to #2) ，gates_weighted 一样，按照 dim=-1 做了归一化而已其输出为：

​为 batch 中每个 sample 选择 Top-K 个专家 这里为 batch 中每个的专家权重，可以从 softmax-ed 来 top-k，也可以直接从 gates_weighted 来 top-k；由于这里可能不做 softmax 或者延后，所以可 gates_weighted，这里为 batch 中每个的专家序号

​其输出为：

接着：

按需计算 2: top-n 专家之间的归一化权重

如何根据分派索引，从原始的输入中，为每个专家提取出属于该专家处理的 tensor，在当前的主流智能平台，都没有专门的算子，可以通过其他算子的组合来实现类似的效果。在 MindSpore 中，可以通过底层的 C++实现算子，也可以通过 Python 中继承 Cell 并实现 bprob，然后将原始 gate tensor 中按照 index 组织到目标输出中。这里我们实现一个 Dispatch 类

​3、独立计算

直接并行调用后续的专家网络。并行部分可以通过平台来支持。可以通过特殊的函数或者 annotation 等标识，也可以由平台编译时优化为并行执行。（在非动态路由条件计算的网络模型中，一般不存在类似的优化。）

4、合并

合并的逻辑相对简单，先通过 cat 按照 batch 维度做拼接，然后构造正确的 zeros tensor 用 index_add 按照索引将各个专家网络的结果在保持 input 序合并到一起，做为该 MoE 模块的输出。

​上述完成了整个 MoE 的完整计算过程。

代码框架

我们按照上述基本动态路由条件计算的张量操作为主的逻辑，扩展到一个完整的训练代码框架中：class Dispatch(ms.nn.Cell): 实现路由中的分派逻辑 classCombine(ms.nn.Cell): 实现路由中的组装逻辑 class Route(ms.nn.Cell): 完成整个动态路由逻辑，可以实现为相对通用的类 class Expert(ms.nn.Cell): 平台用户自定义的专家网络 classNetwork(ms.nn.Cell): 平台用户自定义的大网络 class MSELoss(ms.nn.Cell)：实现 MSE 损失，实现辅助损失的逻辑 classOutputLossGraph(ms.nn.Cell)：输出 infer 和 loss，PyNative 模式单步 class Dataset: 数据集类，仅满足输入 shape 和 X-Y 合理对应关系，仅仅示例 def train( …): 训练入口完整的代码在 mindspore 官网：https://gitee.com/mindspore_ci/mindspore

条件计算实现技术点

1、动态路由

• 不可学习路由

如使用 LSH (locality sensitive hashing)做路由：在整个可学习网络的前端，使用 LSH 来分派样本，这样可以避免 LSH 部分求导问题；如果在网络中间增加 LSH 模块，需要通过梯度估计完成确定性算法部分梯度传递。

• 可学习路由

简单的做法，定义 gate_weights 为可学习 Parameter，对于二维的张量，通过 python@或者 matmul 等完成权重路由计算；如果是更高维度的张量，且需固定 batch 维，einsum('bd*,*de->b*e')的形式完成计算。

2、topk 和 softmax 的前后关系

在 G_1(x)=softmax(topk(X*W)))和 G_2(x)=topk(softmax(X*W)))两类 Gate 实现中，将 softmax 置于 Topk 前后，对 top-k 的选择不变；当需要将 G_*作为后序网络输入的一部分，即将路由权重信息作为后续网络输入信息，则需要考虑：需要 all-N 专家之间的归一化权重，则 softmax 置于 top-k 之前；否则 softmax 置于 top-k 之后，来计算 top-N 专家之间的归一化权重。

3、如何每专家在批次处理中平衡

按照每样本的路由权重求和，即对 batch 单个样本被分配的 1+个 export 的重要性和权重求和，计算出 importance；按照每样本的路由权重中非 0 的求和，计算出有负载的专家来求得 load。将 coefficient_of_variation(importance) +coefficient_of_variation(load)作为 auxiliary_loss 参与优化，来平衡 importance 和 load。变异系数(Coefficientof Variation)是用于无量纲度量数据的离散程度，越离散在此处表示均衡性越差，需要向更小优化。

在 Transformer 等多层（多处）MoE 的模型中，将多组 auxiliary_loss 联合作为 auxiliary_loss,在加 dominated_loss 之后即可。

点击关注，第一时间了解华为云新鲜技术~