技术干货 | 如何基于MindSpore实现万亿级参数模型算法？

本文是Switch Transformer的动态路由条件计算的模型分析的第二篇 - 算法实现。

动态路由条件计算的原理介绍可以参见↓

技术干货 | 一文带你了解MindSpore支持的万亿级参数超大模型关键技术！

实现策略

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

以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、路由

data_inputs = ms.Tensor([
               [0.1,0.9],
               [0.8,0.8],
               [0.9,0.1],
               [0.1,0.9],
               [0.9,0.1],
            ])  #假设输入为5个样本，每个2维，当然可以扩展到高维 (batch,dimension) = (5,2)  
    gate_weights = ms.Parameter(ms.Tensor([
               [0.1,0.5,0.9],
               [0.9,0.5,0.1],
            ] , ms.float32) , 
            requires_grad=True)  #假设路由门权重，3个专家，每个2维和输入一样 (dimension,experts) = (2,3)

当上述输入5行（仅3类，希望分派给3个专家）样本，和Gate权重做矩阵乘后，可以明确算出每个样本要分派的专家。

可以用matmul，也可以类似gates_weighted = einsum('bd,de->be', [data_inputs, gate_weights])

第一轮矩阵乘的结果为：

    gates_weighted= [[0.8200, 0.5000, 0.1800],
            [0.8000, 0.8000, 0.8000],
            [0.1800, 0.5000, 0.8200],
            [0.8200, 0.5000, 0.1800],
            [0.1800, 0.5000, 0.8200]]

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

2、分派

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

gates_softmax = softmax(input=gates_weighted, axis=-1)

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

其输出为：

    gates_softmax= [[0.4438, 0.3222, 0.2340],
            [0.3333, 0.3333, 0.3333],
            [0.2340, 0.3222, 0.4438],
            [0.4438, 0.3222, 0.2340],
            [0.2340, 0.3222, 0.4438]]

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

gates_topk_value, gates_topk_index =topk(gates_softmax, 1)

其输出为：

    gates_topk_value= [[0.4438],
            [0.3333],
            [0.4438],
            [0.4438],
            [0.4438]]
    gates_topk_index= [[0],
            [1],
            [2],
            [0],
            [2]])

接着：

gates_topk_softmax = softmax(gates_topk_value)

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

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

    class Dispatch(ms.nn.Cell):
    def __init__(self, expert_number):
        super().__init__()
        self.expert_number = expert_number
        self.reshape = ms.ops.Reshape()
        self.concat = ms.ops.Concat()
        self.zeros = ms.ops.Zeros()
        self.add = ms.ops.AddN()

    def set_indices_in(self, indices_in): #可以作为construct的参数
        self.indices_in = indices_in

    def get_indices_out(self): #可以用construct的返回值返回
        return self.indices_out

    def construct(self, data):
        dispatch = []
        indices_out = []
        for _ in range(self.expert_number):
            dispatch.append([])
            indices_out.append([])
        for uid,(idx,dat) in enumerate(zip(self.indices_in, data)):
            dat = self.reshape(dat, (1, dat.shape[0]))
            if len(dispatch[idx]) == 0:
                dispatch[idx] = dat
                indices_out[idx] = [uid]
            else:
                dispatch[idx] = self.concat((dispatch[idx], dat))
                indices_out[idx] = indices_out[idx]+[uid]
        self.indices_out = [y for x in indices_out for y in x]
        return dispatch

    def bprop(self, data, out, dout):   #反向梯度计算
        dall = None
        for one in dout:
            if dall == None:
                dall = one
            else:
                dall = self.concat((dall, one))
        do = self.zeros(dall.shape, ms.float32)
        for idx_target, idx_source in enumerate(self.indices_out):
            do[idx_target] = self.add((do[idx_target], dall[int(idx_source)]))    
        return do

3、独立计算

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

4、合并

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

class Combine(ms.nn.Cell):
    def __init__(self):
        super().__init__()
        self.zeros = ms.ops.Zeros()
        self.add = ms.ops.AddN()

    def set_indices(self, indices):  #可以作为construct的参数
        self.indices = indices

    def construct(self, data):
        O = self.zeros(data.shape, ms.float32)
        for idx_target, idx_source in enumerate(self.indices):
            O[idx_target] = self.add((O[idx_target], data[int(idx_source)]))
        return O

    def bprop(self, data, out, dout):  #反向梯度计算
        do = self.zeros(dout.shape, ms.float32) 
        for idx_target, idx_source in enumerate(self.indices):
            do[idx_target] = self.add((do[idx_target], dout[int(idx_source)])) 
        return do

上述完成了整个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( …): 训练入口

完整的代码在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。变异系数(Coefficient of Variation)是用于无量纲度量数据的离散程度，越离散在此处表示均衡性越差，需要向更小优化。

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

MindSpore官方资料

官方QQ群 : 486831414

官网：https://www.mindspore.cn/

Gitee : https : //gitee.com/mindspore/mindspore

GitHub : https://github.com/mindspore-ai/mindspore

论坛：https://bbs.huaweicloud.com/forum/forum-1076-1.html