关于参数:如何基于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 a class of algorithms that activate only some of the different parts in a network. 在具体某类条件计算实现中，条件抉择模式，可能依照输出的 batch 中每 sample 独立激活网络不同局部，可能依照输出数据空间上不同的局部（比方 image 不同区域或者 channel），可能依照输出数据工夫上不同的局部（比方 time series 的不同 slide window 或者 video 的不同的 frame。），可能依照指标工作的不同每 task 独立的，可能依照非可学习的固定的随机调配不同的子网独立计算。

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

条件计算的分类（狭义）

1. 依照 routing 是否可学习能够分为：learnable routing conditional computation 和 unlearnable routing conditional computation.

2. 依照 activation 是否不执行 non-activation 计算，能够分为：hard conditional computation 和 soft conditional computation。对于 hard-mode 的条件计算，通过 tensor 筛选切分等操作，无论何种条件抉择模式，不须要激活的数据将齐全不参加不激活的网络局部的计算；soft-mode 的条件计算，可能仅采取将相干数据置零等形式来防止产生计算成果，但还是和不须要激活网路局部理论执行计算过程。

条件计算的次要劣势

1. 计算无效，升高能耗：通过局部激活局部计算，以每样本条件激活的条件计算为例，单个样本只须要通过整个 SuperNet 的一部分参加计算。

2. 更大网络，表白更强：因为一处到多处的 Route，各处（层）的 Input 被路由到不同的子网独立计算，不同的输出的互相在各层的表白绝对独立没有影响，表达能力更强，网络能够更大，但表白效率升高了。

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

条件计算的网络和计算模式比拟灵便，局部构建模式如：（此处省略具体模型和论文援用，参见: http://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(XW)，G(x)为路由 Gate 函数，X 为输出，W 为通损失函数来度量的可学习路由权重，P 为某种筛选函数（如 topk, sort 等），在理论实现中，X W 的输出与权重计算结果可能作为后续网络的输出信息的一部分，不仅仅利用 G(x)来抉择路由，则须要对 X W 的后果做归一化，更典型的模式则为：G(x)=P(N(XW))，其中 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-Gated Mixture-of-Experts Layer》做了相干探讨。

对于条件计算 / 动静神经网络

对于条件计算 / 动静神经网络，更多的信息在《Dynamic Neural Networks: A Survey》arxiv:2102.04906 (http://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 示意实现。(注：import mindspore 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): 实现路由中的分派逻辑
• class Combine(ms.nn.Cell): 实现路由中的组装逻辑
• class Route(ms.nn.Cell): 实现整个动静路由逻辑，能够实现为绝对通用的类
• class Expert(ms.nn.Cell): 平台用户自定义的专家网络
• class Network(ms.nn.Cell): 平台用户自定义的大网络
• class MSELoss(ms.nn.Cell)：实现 MSE 损失，实现辅助损失的逻辑
• class OutputLossGraph(ms.nn.Cell)：输入 infer 和 loss，PyNative 模式单步
• class Dataset: 数据集类，仅满足输出 shape 和 X - Y 正当对应关系，仅仅示例 def train(…): 训练入口

条件计算实现技术点

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(XW)))和 G_2(x)=topk(softmax(XW)))两类 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。变异系数 (Coefficient of Variation) 是用于无穷纲度量数据的离散水平，越离散在此处示意均衡性越差，须要向更小优化。

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

点击关注，第一工夫理解华为云陈腐技术~