当前位置: 首页 > 科技观察

本文带你了解MindSpore支持的万亿级参数超大模型的关键技术!

时间:2023-03-18 19:51:02 科技观察

前言最近一段时间,增加模型的尺寸成为提升模型性能的主要手段。尤其是NLP领域的自监督预训练语言模型越来越大,从GPT3的1750亿参数到SwitchTransformer的1.6万亿参数,又是一个数量级的增长。模型尺寸的数量级增长实现了一定程度的性能提升,甚至产生了一些意想不到的“神奇”效果(比如GPT3),但背后的计算开销却成了最大的问题,比如GPT3训练使用Tens数以千计的GPU和数周的训练时间。如何利用超大规模参数来提升模型表达和性能,同时控制计算量的微小增加,成为最重要的挑战之一。着重介绍了以MoE为代表的动态神经网络技术。大脑是典型的低能高效计算模型,稀疏激活是最重要的特征。除了巨型模型在训练推理,特别是训练方面的计算效率挑战之外,目前巨型模型训练优化算法的另一个更大的挑战(这里不讨论),BP算法是目前最可用的深度网络优化,但更理想的优化算法要求并行性高,优化过程非对称,能够在时空维度上通过局部连续优化完成整体优化。1、在传统的神经网络模型中,前馈时,在输入batch中,对每一个样本的处理都会激活网络中的每一个参数参与计算。2.条件计算最宽泛的定义是指一类只激活网络特定部分的算法。条件计算是指仅激活网络中某些不同部分的一类算法。在特定类型的条件计算实现中,条件选择模式可以根据输入批次中的每个样本独立激活网络的不同部分,也可以根据数据空间的不同部分(如不同区域或通道)激活图像)可能根据输入数据时间的不同部分(例如时间序列的不同幻灯片窗口或视频的不同帧。)而不同,可能根据不同的目标任务独立于每个任务,可能是根据对不同子网络的不可学习的固定随机分配独立计算。3.针对不同的输入(原始或前层),根据一定的条件,选择性地执行网络后续部分的计算。在这种技术下,有一些近似或相关的技术,例如:动态神经网络、条件计算、条件激活、稀疏激活、选择性执行、专家混合(MoE)、动态路由……;条件计算的分类(广义)1.根据路由是否可以学习,可分为:可学习的路由条件计算和不可学习的路由条件计算。2.根据激活是否不进行非激活计算,可以分为:硬条件计算和软条件计算。对于hard-mode条件计算,张量选择、分割等操作,无论是什么条件选择模式,不需要激活的数据根本不会参与非激活网络部分的计算;soft-modeconditionalcalculations可能只需要对相关数据进行Zeroing等,避免了计算效果,但仍然不需要激活网络部分来真正执行计算过程。条件计算的主要优点1.计算有效,降低能耗:通过部分激活部分计算,以per-sample条件激活的条件计算为例,单个样本只需要经过一部分计算整个超网参与计算。2.更大的网络,更强的表达:由于路由从一处到多处,每一(层)的输入路由到不同的子网进行独立计算,每一层不同输入的表达相对独立,没有效果,表达能力更强,网络可以更大,但是表达效率降低了。条件计算的网络和计算形式条件计算的网络和计算形式比较灵活,部分构造形式如下:(这里省略具体模型和论文参考,见:http://intellabs.github.io/dis)1.根据CV等任务的特点,使用多个独立的CNN作为专家网络,根据任务独立路由,尾部合并形成一个大网络。2、使用级联等更复杂的形式,将不同层次的不同专家网络组合起来。3.通过决策树等方法,通过数据转换实现路由。4.通过可学习的网络进行路由。其中,policylearning的loss有多种构建形式:直接使用分类等任务的主要loss,不同专家的重要性和负载构建loss作为辅助loss等。条件计算的路由策略1.Non-learnable/hard-mode,使用一定的确定性策略,比如LSH,来计算路由。2.learnable-mode,通过可学习网络计算路由。网络规模可大可小,简单的可学习路由是单层权重:G(x)=P(X*W),G(x)是路由门函数,X是输入,W是通过损失函数衡量路由权值可以学习到的,P是一个选择函数(比如topk,sort等)。在实际实现中,X*W的输入和权重计算结果可能会作为后续网络的一部分输入信息,不仅要使用G(x)来选择路由,还需要对X*W的结果进行归一化处理.比较典型的形式是:G(x)=P(N(X*W)),其中N是Normalization函数的表达式,比如Softmax。条件计算的冗余策略条件计算的冗余策略分为非冗余条件计算和冗余条件计算:1.非冗余条件计算可以通过P(.)函数实现,如topk(k=1,...)实现;2.冗余条件计算可以通过多种形式实现,例如topk(k=n,...),n>=2,或者通过硬冗余方式,支持输入复制和全网多路计算实现。ConditionalComputing的挑战1.路由算法对模型质量的影响不管输入和路由权值信息(X*W)是只作为路由选择和后续网络单元的输入,还是直接作为后续网络输入的一部分units,路由算法决定了输入信息的处理流程,对模型的整体质量影响很大。2.routing/gate的稳定性routing/gate的权重是随机初始化的,权重本身也在通过训练不断调整;前后层的网络不断地训练和变化,同一个样本在训练的不同阶段会被分配到不同的后续网络单元,这种动态变化过于剧烈,会严重影响网络的稳定性和收敛速度整个网络训练过程。3.专家样本在路由中的重要性和训练阶段的负载均衡,样本batch中每个专家的重要性和样本的相关性,每batch中样本的负载均衡被平均分配给不同的专家,这些两个指标既相关又相互冲突。需要单独构造一个损失函数作为辅助损失来优化这两个指标。在arxiv:1701.06538《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》进行了相关讨论。Aboutconditionalcomputing/dynamicneuralnetwork关于conditionalcomputing/dynamicneuralnetwork,更多信息在《Dynamic Neural Networks: A Survey》arxiv:2102.04906(http://arxiv.org/abs/2102.0490)一文中,作者有广义动态神经网络,各种动态网络相关技术按照实例级别、时间级别、空间级别进行分类。1.Instance-wiseDynamicNN:Instance-wisedynamic,每个样本独立激活不同的网络和参数(MoE就是这个方向)。动态架构:动态深度、动态宽度、动态路由/MoE;DynamicParameter:ParameterAdjustment,ParameterPrediction,DynamicFeature(s)2.Spatial-wiseDynamicNN:Spatial-leveldynamics:不同的空间位置如图像激活后续不同的网络和参数。(CNN等):PixelLevel,RegionLevel,ResolutionLevel3.Temporal-wiseDynamicNN:Time-leveldynamics:时间序列数据根据时间序列维度进行分割,激活后续不同的网络和参数。(video-frames,text-sequence,time-series,stream,...)上面的Text-SequenceVideo-Frames是这篇综述论文中DynamicNN的整体分类。主要从超大规模网络动态网络技术支持、高表现力、低计算成本的角度考虑分类,从两个维度对动态网络技术进行分类:1.根据是否部分激活时feedforwardcalculation:Hard-Dynamic:在前馈过程中,部分网络永远不会被激活参与计算。Soft-Dynamic:在feedforward时,部分网络会在经过softmax等gates/route后将张量元素置零,失去表达能力,但会参与计算。.2.根据动态激活决策算法的输入:Sample-by-samplelevel:(在输入层)根据每个样本的实例决定动态网络的后续激活。子样本级别:(在输入层)在样本内的时间/空间级别激活不同的后续网络单元。通常,深度网络不仅在输入层被选择性激活,在中间层也会被选择性激活。其中,智能平台支持Hard-Dynamicsample-by-sample动态神经网络,可以自然地获得大规模网络结构的稀疏激活,在超大型模型的训练和推理中可以实现高能效。与静态结构的神经网络相比,动态神经网络在性能、表达、泛化、鲁棒性和可解释性等方面做了大量的对比研究。从以尽可能低的计算成本支持超大规模网络以提高模型性能的智能平台来看,Efficiency和Representation是最重要的:1.Efficiency:静态网络“牵一发而动全身”,每个样本输入到整个网络/所有的参数都要有响应,这对于超大网络的模型能耗要达到领先的结果挑战太大了。2.Representation:参数数量更多,表达能力更大;但是在深层网络各层特征的表达上,MoE等结构减少了复用,各参数的表达效率较低。实现策略要实现各种模型动态路由稀疏激活的超大规模参数版本,需要分模型进行研究和实现。以SwitchTransformer为例,其参数扩展到Transformer的FFN部分的一部分。它的MoE扩展,如下图所示:(图片来源:SwitchTransformerpaper)可以看出,MoE的主要变化是在需要Expert子网络前后增加了MoE相关的逻辑。本文主要介绍平台上的实现。动态路由条件的计算主要包括四个步骤:路由计算、数据分发、独立计算、结果合并。1.路由计算-Gate:根据输入(可以是整个网络的输入,也可以是前一个网络单元/层的输出),在routingunit完成计算,在sample-wiserouting在批次中,计算每个样品的所需值。分配后续网络路由(专家混合专家/MoE)。2.数据调度-Dispatch:从输入的整体Tensor中,根据路由计算出的样本-专家关系,收集合并每个专家需要处理的Tensor。如果在设计一个固定的expert-batch时,由于样本输入的随机性,需要平衡每批训练分配给每个专家的样本数和每轮专家训练的最大容量,有难以保证相对均匀的分布。对于有最大容量的batch,固定batch-size的样本需要pad,大于最大容量的样本可以采用延迟重采样等方法。为了保持正确的输入输出关系(Input/X–Label/Y)和训练是反向传播的导数关系,实现需要保持从原始batch到每个专家的子batch的索引关系,以及然后derive和combine合并时使用。3.独立计算-Expert:并发调用每个expert(逻辑上一个接一个)处理对应的sub-batch。这也是智能平台要支持的并发API之一。4.ResultMerge-Combine:将每个专家的结果张量合并到整批的张量中,并根据数据分配索引,交换为原始输入的顺序。在主流的深度学习智能平台中,主要可以采用两种实现策略:Tensorzeroing:对于需要分配到不同后续网络单元(专家网络子网等)的,复制若干张量给需要分配的专家,将不应该进入当前专家流程的数据维度归零。该方法实现简单,全张量运算,在归零计算逻辑正确的情况下,对平台没有特殊要求。适用于算法研究,仅体现条件计算的前序数据动态路由到不同的后续网络单元。分析算法的效果。如果该方法设置为零,则该方法中每个专家处理的张量在batch维度上都是一个完整的batch,无法节省计算量和内存占用。张量排序:对于需要分配到不同后续网络单元(专家网络子网等)的,为需要分配的专家复制若干个张量,不保留不应该输入的数据维度现任加工专家。并维护变换前后样本级指标的对应关系。在分布式友好的实现中,如果将专家子网划分为不同的计算节点,那么专家网络的实现最好继承子网级平台对象,例如:MindSpore中的mindspore.nn.Cell。详细的实现细节见后续技术实现章节。核心代码核心代码:路由计算、数据分发、独立计算、结果合并参考代码由MindSpore实现。(注:importmindsporeasms)MixtureofExperts的核心逻辑,对于输入I,经过routing_network(最简单的*W即可),然后topk(如果variant算法需要gateweight,则需要softmax,否则为notnecessary),然后使用tensor操作(根据batch)选择每个subnetwork/expert的tensor。为了调试方便,使用一个非常小的非随机定值构造输入和路由权值,路由网络使用简单的X*W。1.路由计算当上面输入5行(只有3个类别,希望分配给3个专家)样本,和Gate权值矩阵相乘,就可以清楚的计算出每个样本要分配的专家。可以使用matmul,也可以类似于gates_weighted=einsum('bd,de->be',[data_inputs,gate_weights])第一轮矩阵乘法的结果为:输入与权重相乘,可以使用@在python中,也可以用matmul也可以用爱因斯坦求和简单助记函数einsum。当是简单的矩阵乘法时,使用einsum编译计算图时,实际上会拆分成多个算法,性能不好;但是当input和weight超过2维时,需要固定batch维度进行路由计算,使用einsum可以编程实现,非常简单。2.assignmentassignmentconditioncalculation,主要逻辑是根据路由网络的输出计算每个样本的top-k专家。它的实现可以通过topk函数来实现。由于topselectionscore可以作为后续网络单元的输入信息(包括路由信息),所以一般需要对路由输出进行softmax归一化。按需计算1:all-N专家之间的归一化权重(请参考#2),同gates_weighted,按照dim=-1归一化,输出为:selectTopforeachsampleinthebatch-Kexpertsherebatch中每个专家的权重,可以是softmax-ed的top-k,也可以直接从gates_weighted到top-k;由于此处可能无法进行softmax或delay,因此可以在此处使用gates_weighted。batch中每个专家序号的输出为:下一篇:按需计算2:如何根据assignmentindex张量对每个专家从原始输入中提取top-n专家之间的归一化权重,目前主流的智能平台,没有专门的算子,通过其他算子的组合可以达到类似的效果。在MindSpore中,可以通过底层C++实现operators,也可以通过继承Python的Cell实现bprob,然后将原始门张量根据索引组织成目标输出。这里我们实现一个Dispatch类3,独立计算直接并行调用后续的专家网络。平台可以支持并行部分。它可以通过一个特殊的函数或注释来标识,也可以在编译时由平台优化并行执行。(在非动态路由条件计算的网络模型中,一般不会有类似的优化。)4.合并合并的逻辑比较简单。首先使用cat按照batch维度进行拼接,然后构造出正确的zerostensor并使用index_add对每个索引专家网络的结果进行合并,同时保持输入顺序,作为MoE模块的输出。以上就完成了整个MoE的完整计算过程。代码框架我们将基于上述基本动态路由条件的基于张量运算的逻辑展开为一个完整的训练代码框架::实现路由中的组装逻辑Network类Network(ms.nn.Cell):平台用户定义的大型网络类MSELoss(ms.nn.Cell):实现MSE损失和辅助损失的逻辑类OutputLossGraph(ms.nn.Cell):outputinferandloss,PyNativemodesingle-stepclassDataset:数据集类,只满足输入shape和X-Y合理对应,举个例子deftrain(…):trainingentryconditioncalculationandimplementationtechnicalpoints1.动态路由无法学习路由如使用LSH(localitysensitivehashing)进行路由:在整个可学习网络的前端,使用LSH来分发样本,可以避免LSH偏导问题;如果在网络中间加入LSH模块,则需要通过梯度估计来完成确定性算法的部分梯度传递。学习路由的一种简单方法是将gate_weights定义为可学习的参数。对于二维张量,通过python@或matmul完成权重路由计算;如果是高维张量,需要固定batch维度,einsum('bd*,*de->b*e')完成计算。2.topk和softmax的关系在两类Gates中实现,G_1(x)=softmax(topk(X*W)))和G_2(x)=topk(softmax(X*W))),softmax放在Topk前后,top-k的选择不变;当需要将G_*作为后续网络输入的一部分,即路由权值信息作为后续网络的输入信息时,需要考虑:all-N专家之间的归一化需要权值,softmax放在top-k之前;否则softmax放在top-k之后以计算top-N专家之间的归一化权重。3、批处理中各个专家如何平衡?根据每个样本的路由权重总和,即batch中分配给单个样本的1+出口的重要性和权重总和,求和计算重要性;专家根据每个样本的路由权值和中间非0的总和与负载进行计算得到负载。使用coefficient_of_variation(importance)+coefficient_of_variation(load)作为auxiliary_loss参与优化,平衡importance和load。变异系数(CoefficientofVariation)用来衡量无量纲数据的分散程度。越是离散,这里的平衡就越差,需要优化到更小的尺寸。在Transformer等多层(多)MoE模型中,多组auxiliary_loss合并为auxiliary_loss,加入dominated_loss后。