摘要:近来,增大模型规模成为了晋升模型性能的次要伎俩。特地是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等),在理论实现中,XW的输出与权重计算结果可能作为后续网络的输出信息的一部分,不仅仅利用G(x)来抉择路由,则须要对XW的后果做归一化,更典型的模式则为: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之后即可。

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