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

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

动态路由条件计算的原理介绍可以参见上一篇《一文带你了解MindSpore支持的万亿级参数超大模型关键技术！》

实现策略

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

(图片来源：Switch Transformer论文)

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

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

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

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

核心代码

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

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

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)

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]]

gates_softmax = softmax(input=gates_weighted, axis=-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_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]]

接着：

gates_topk_value, gates_topk_index =topk(gates_softmax, 1)

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

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

了解完MindSpore的关键技术是不是很心动呢！赶紧【点击链接】并【立即报名】，即可在 ModelArts 平台学习到一个经典案例掌握基于MindSpore的深度学习！

想要了解更多关于大模型的知识，请点击：专家解惑 | 关于华为云盘古大模型，你想问的都在这里~

实实现实现策略策略现策略

MindSpore 专家

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。