Press "Enter" to skip to content

安培架构下稀疏特性的实践与挑战

分享嘉宾 :余翀 NVIDIA Senior Architect

 

编辑整理: 赵泽明  零氪科技

 

出品平台:DataFunTalk

 

导读: 稀疏特性是我们在做机器学习时,经常会碰到的一个问题。 利用矩阵的稀疏特性,也是在机器学习过程中常用的一种优化方式。 了解稀疏特性,以及了解如何利用稀疏特性来解决模型部署、提高模型精度中的一些问题,很有必要。 今天的分享主要围绕以下几点展开:

 

稀疏特性的整体介绍

 

安培架构下稀疏特性的挑战及利用

 

从非结构化稀疏到结构化稀疏

 

生成对抗网络的模型压缩

 

01

 

稀疏特性的整体介绍

 

首先,稀疏特性是在机器学习中经常会碰到的一个问题。

 

就一个神经网络模型来讲,训练好的每一层的权重大部分都是符合类似于高斯分布的情况,大量的元素的幅值非常小,甚至接近于零。

 

另外,在神经网络中,会大量使用各种各样的激活函数,很多传统的激活函数,例如ReLU,会使坐标轴小于0的地方均为零,所以会使得模型在训练和部署的过程中产生大量的稀疏矩阵的运算。

 

对于传统的计算机来说,稀疏特性并不友好,因为稀疏矩阵的稀疏度不是很高时,如果专门为它设计稀疏运算或者运算库,最后得到的性能收益非常小。原因在于,计算机最大的开销往往都在对内存的非顺序访问,而稀疏特性会导致在计算机访存中产生很多断续操作,最终就会导致虽然做出了一个可能看上去比较高效的稀疏的加速,但事实上在真正部署的时候,效果并不理想。

 

一般来说,稀疏特性分成两种,如下图所示。

 

 

从左到右为从细粒度稀疏过渡到到粗粒度稀疏。 分别对应相应的细粒度稀疏化(Fine-grained)、向量稀疏化(Vector-level)、核稀疏化(Kernel-level)、滤波器稀疏化(Filter-level)。从图中看到, 越粗粒度的稀疏,对模型来说,最后得到的pattern往往越规则,而越规则的pattern,对于计算机来说越容易对其加速。

 

但是, 粗粒度的稀疏特性也有一个弊端 ,就是在filter-level维度对它做剪枝或稀疏化, 保持精度比较困难 。而相应地,对细粒度的稀疏,因为它的裁剪力度可以到元素或者神经元级别,所以容易保持精度。但是细粒度稀疏又因为其非规则性,使得硬件加速非常困难。

 

所以,如果可以把模型里面的细粒度的稀疏特性和GPU的高性能运算器件的硬件特性相结合,使其能够在保持模型的精度的同时,减少网络模型大小(因为这样对于内存调用的压力会减小),使GPU可以对细粒度的稀疏特性进行高效的执行和运算提速,则是非常有意义的。

 

这就引出我们今天的第二个主题。

 

02

 

安培架构下稀疏特性的挑战及利用

 

在安培架构中,现已经支持细粒度的稀疏特性,细粒度的稀疏特性是N:M的结构化稀疏,以2:4的稀疏举例来说,2:4的含义是:

 

在一个权重矩阵中,对于任何四个连续在一起的元素来说,只要保证其中有两个元素非零,就认为它是一个2:4的细粒度稀疏矩阵,在实际处理过程中,如下图。

 

 

左侧矩阵是一个C×R的稀疏矩阵,实际上在存储时,只需把其中50%非零的元素存下来就可以。同时需要保存一个metadata作为index,指向每四个元素中非零值的两个元素。其实2bit的metadata也会带来一些开销,但是比起32位或者16位的运算,metadata带来的开销相对是比较小的。

 

我们用 两个实例 来说明我们对稀疏特性的应用。

 

例如使用GEMM(通用矩阵乘加)进行矩阵乘法计算,A×B的运算,在硬件维度,会转化为一个个小的tile运算。我们以16×16和16×8的两个tile举例说明,这两个tile最终生成C,一个16×8的tile。如下图左。

 

 

我们把问题的维度变大一点,从刚才A(维度16×16)×B(维度16×8)的矩阵运算,扩大一倍,变成了A(维度16×32)×B(维度32×8)的运算。此时,硬件维度实现这个tile的乘法,需要循环两次,才能完成一个是最后的生成16×8的tile,如上图右。而如果A矩阵符合2:4系数稀疏特性,那幺在只要一半的空间就存储整个A的元素,在运算时,根据A的metadata从B中选择哪些元素需要进行相乘,这样既保证了计算结果,同时只需要一个cycle即可完成运算。如下图左。

 

 

在安培这一代GPU架构里面,对于各种数据类型的运算,都可以提到提供到两倍的运算加速(除了binary格式,其实因为它的metadata的开销过大)。

 

03

 

从非结构化稀疏到结构化稀疏

 

上文中提到的2:4的这种细粒度稀疏,可以认为是结构化稀疏的运算特性。但是稀疏特性更多时候都是非结构化的稀疏。我们可以广义的认为只要不符合这种2:4的稀疏特性的矩阵,都是非结构化稀疏,没有办法直接利用安培GPU的技术特性进行加速。那幺如果强制对一个矩阵进行2:4稀疏方式,也就是四个元素只保留两个,相当于加了一个2:4结构的mask,之后得到的结构化稀疏矩阵,会有什幺影响?

 

我们选取了一些baseline,使用distiller开源框架 , 进行一些尝试,直接使用加上2:4的sparse mask,结论是模型的精度会有相当大的损失 。所以直接强硬的进行从unstructured到structured sparsity的转换,会带来很大的精度损失。我们一般会用一些例如微调、知识蒸馏的方法。知识蒸馏就是一个teacher模型,一个student模型,我们期望可以使student模型从teacher模型中学到知识,在整个Fine Tune的过程中接受teacher模型的指导,最终提升被压缩的student模型的精度。知识蒸馏和传统的Fine Tune不一样的地方是,在最后的softmax之后,会进行hard prediction,与hard label 计算一个Loss。

 

此外,因为student模型还要尽量的去模仿teacher模型,所以 在softmax层会定义一个温度 ,其实就是一个概率分布,例如T=5,就是说在top5分类时,去学teacher模型对于top5的每一个预测的概率,如果student模型和teacher模型最后学出来的概率分布类似,就可以认为是说student模型基本上就是符合teacher模型,在计算topk时,会根据soft prediction计算一个Distillation Loss。

 

最终在训练的过程中是student的两个Loss加权,来做整个模型的Fine Tune,会比直接使用Hard Label去做Fine Tune的运算效果更好。

 

 

类比思考:对unstructured sparsity模型, 如果用一个dense模型直接用知识蒸馏去做Fine Tune,是不是也有效果? 其实在大部分情况之下确实也可以,但是会带来两个问题,第一个问题是,如果是在一个已知的训练集上做训练,label信息是已知的,这样做效果相对较好。但是很多时候,我并不一定可以知道所有的label信息或者可以访问到所有的数据集,这时,这个方法的有效性就会大打折扣。

 

第二个问题是,因为从dense模型到一个structure sparsity模型,硬件上的约束要求比较高,硬件约束要求越高的情况之下,teacher model 和student model 的soft Prediction的Distillation Loss,两者的差距本身就会越来越大,这时两者越来越难以通过Distillation Loss去学到一个共同的行为。

 

在这个情况之下,我们做一个转换,把传统的知识蒸馏改成了如下图的结构。

 

 

Dense Model作为总的Teacher Model,还会放一个Unstructured Sparsity Model作为第二个Teacher Model,我们最终的目标是从Unstructured Sparsity Model转换到最后的2:4的Target Structure Sparsity Model。此时,我们的Loss分为三部分。

 

首先第一部分依然是 Hard Prediction ,用baseline模型和Target Model的预测做一个Prediction Loss。这样做的好处是,即使访问不到数据集,但是任意的一个数据集,只要最后预测出的label和Dense Model一致即可(因为Dense Model是已知的),这样就解决了上文所述的第一个问题,即对数据集的依赖。

 

第二部分,依然像传统的知识蒸馏的Soft Prediction一样, 计算Dense Model 和Target Model之间的Distillation Loss 。

 

第三部分, 计算Unstructured到Structured Model之间的Distillation Loss ,这个Loss 模拟了因施加2:4 sparsity mask后导致的两个模型模拟能力的不同,加上了这个Loss后,解决了上述的第二个问题:teacher model 和student model的Distillation Loss过大。

 

另外,在各种各样的task上面做了一系列实验,在传统的分类问题上,很多模型经过distill最终得到的structured sparsity模型,它的精度都是要远高于unstructured sparsity模型。同样对于一些NLP的模型也做了这样的尝试,比如transformer,也能够回到Dense模型本身的精度。

 

04

 

生成对抗网络的模型压缩

 

生成对抗网络的含义是两个模型之间做博弈,包含一个生成器和一个判别器。生成器可以将一个输入的数据转化成fake data,和一个实际的训练样例一起输入至判别器,去判断哪个数据是真实的,哪个数据是伪造的。

 

当判别器无法判断出两个数据到底谁真谁假,则说明生成器的性能就很强,相反,如果很容易地判断出哪个真哪个假,那就说明生成器效果不佳,此时判别器就会把结果回馈给生成器,促使生成器更好的训练。

 

那幺 为什幺我们会单独的去研究生成对话网络的模型压缩? 因为生成对抗网络和传统的各种网络结构是有很大的区别的,例如classification network,detection network,都是有一个可量化的指标,但是对于生成网络,往往都是很主观的判断目标。另外,在传统的机器学习中,基本上都是根据loss做一个反馈。如果把一个生成器做压缩,用压缩过的生成器和判别器,在后面直接做Fine Tune,loss曲线基本上也是一个比较平滑的收敛曲线,但最后效果却不佳。

 

所以对于生成对抗网络,Loss曲线其实只表示在训练过程中间生成器和判别器的训练是否相互平衡。如果发散,则说明最后生成的效果一定不好,但如果Loss曲线最后收敛,并不能说明结果就一定很好。另外,在对抗训练中,保证生成器和对抗器的平衡性非常重要,假如使用传统的方式对生成器进行模型压缩,那幺训练就会非常的不平衡,Loss就很容易发散,导致训练出来的压缩模型效果不佳。而如果把判别器也做压缩,这样似乎可以保证在训练中的平衡性,但实际效果并不好。

 

主要原因在于,对判别器做裁剪后,判别器的能力被削弱,会导致模式崩塌问题。判别器要让每个不同样本最终能够收敛到其期望的不同的局部最小值,来保证最终结果的多样性,而如果判别器性能被减弱,会导致训练时它依然收敛,但是这些收敛的值都收敛到同一个局部最小值,这就是模式崩塌现象。最后生成的样本都变成单一的样式,失去了多样性的效果。

 

此外,对于前面所述的各种传统的机器学习,例如detection或者classification,最终只要判断top 1或者top k的精度,都是一个信息熵减少的过程。所以怎样在信息熵不损失的情况下对模型所进行压缩,也是目前面临的一个挑战。

 

所以针对以上问题,我们改进了framework,如下图。

 

 

为了保证训练是平衡的,引入了两个生成器,一个是用来做压缩的生成器(Compressed Generator),另外一个是原始的生成器,在训练过程中间逐步对Compressed Generator做压缩,保证在训练过程中两个生成器是相互平衡的。同时考虑可能的模式坍塌,我们对判别器不做任何压缩,之后计算Loss。最终保证这两条支路的Loss尽量一致,来保证整个系统的稳定。

 

如果做粗粒度的稀疏,也就是filter-level的稀疏,虽然可能最后加速的效果会比较好,但很难保持精度稳定,压缩的生成器模型生成的样本会有很严重的失真、变色或者虚影现象。但是如果进行Fine-grained级别的稀疏化来说,在50%的稀疏度情况下,压缩的生成器模型生成的样本和原生成器模型生成的样本几乎一致。

 

今天的分享就到这里,谢谢大家。

 

在文末分享、点赞、在看,给个3连击呗~

 

01 / 分享嘉宾

 

 

余翀

 

NVIDIA Senior Architect

 

目前在英伟达公司担任资深架构师,主要负责在GPU架构上的稀疏模型压缩和低精度量化的算法研究,模型部署以及性能优化。曾在英特尔公司负责深度学习框架和算子在CPU架构上的优化工作。

 

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注