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

直接并行调用后续的专家网络。并行部分可以通过平台来支持。可以通过特殊的函数或者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的完整计算过程。

代码框架

条件计算实现技术点

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)是用于无量纲度量数据的离散程度，越离散在此处表示均衡性越差，需要向更小优化。