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

举报
HWCloudAI 发表于 2021/06/22 10:10:28 2021/06/22
【摘要】 本文是Switch Transformer的动态路由条件计算的模型分析的第二篇 - 算法实现。

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

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

实现策略

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

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

1.png

(图片来源:Switch Transformer论文)

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

2.png

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_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专家之间的归一化权重

如何根据分派索引,从原始的输入中,为每个专家提取出属于该专家处理的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的关键技术是不是很心动呢!赶紧【点击链接】并【立即报名】,即可在 ModelArts 平台学习到一个经典案例掌握基于MindSpore的深度学习!

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

实现实现策略策略现策略


【版权声明】本文为华为云社区用户原创内容,转载时必须标注文章的来源(华为云社区)、文章链接、文章作者等基本信息, 否则作者和本社区有权追究责任。如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容,举报邮箱: cloudbbs@huaweicloud.com
  • 点赞
  • 收藏
  • 关注作者

评论(0

0/1000
抱歉,系统识别当前为高风险访问,暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称,即可参与社区互动!

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。