Transformer量化笔记(一)

作者:安平博,Xilinx高级工程师;来源:AI加速微信公众号

前言

AI芯片(这里只谈FPGA芯片用于神经网络加速)的优化主要有三个方面:算法优化,编译器优化以及硬件优化。算法优化减少的是神经网络的算力,它确定了神经网络部署实现效率的上限。编译器优化和硬件优化在确定了算力的基础上,尽量最大化硬件的计算和带宽性能。经历了一年多的理论学习,开始第一次神经网络算法优化的尝试。之所以从一个FPGA开发者转向算法的学习,有几个原因:

第一是神经网络在AI芯片上的部署离不开算法的优化。一个浮点数的计算(加法或者乘法)和定点数的计算消耗的资源差距很大,对于FPGA这样逻辑资源有限的芯片而言,定点计算更加友好,而且能够提升几倍于浮点计算的性能。

第二是神经网络量化压缩需要密切的结合FPGA硬件的特点,需要考虑到FPGA的存储资源,计算符号是否能够被FPGA友好的实现等。在AI加速器项目中,算法和FPGA都有各自的开发者,FPGA会对算法组提出要求,比如激活函数量化,normalization如何做等,然后算法组在这些特定要求下去进行算法优化。如果一个人对FPGA和算法都比较熟悉的话,那么就会更容易发现算法优化的点。

第三是FPGA开发方式的趋势是多样化。使用RTL语言仍然是主要的开发方法,需要一个人有一定的数字电路基础。这种开发方式最底层,所以最灵活,可以更好的去调优。但是同时,FPGA一直渴望去突破固有的开发方式,让一个不懂得硬件的软件开发人员也可以很容易的上手,同时能够缩短开发周期,比如HLS。我相信,随着HLS的发展和FPGA芯片的演进,使用这种方式的开发者会越来越多。在那些算法复杂,更新较快的项目中,HLS更有优势,而在一些对资源,时序,功耗要求更高的项目中,RTL更有优势。当硬件平台逐渐软件化后,必然会对FPGA开发者的算法能力提出更高的要求。

Transformer网络结构

Google在《Attention is all your need》的文章中,提出了使用全attention结构替代LSTM的transformer模型,在翻译任务上取得了更好的成绩。这个网络结构计算量大,计算符号相对简单,有一定的应用,所以适合用于网络加速的展示。结构整体模型如下:

1 embedding

包含了input和output的embedding层,完成词汇到网络输入向量的转化,embedding的矩阵大小取决于词汇量的多少,对于翻译来讲,通常都是巨大的,所以其不适合放在FPGA上进行加速,没有量化的必要。Input和output以及softmax前的linear层都共享相同的参数,这样做的目的,是因为共享input和output权重能够降低word level perplexity,当然也降低了参数存储量。最后的linear使用embedding的权重是为了将网络向量转化为词语出现的logits。

2 positional encoding

Transformer是没有循环网络,为了获取词汇位置关系信息,对词汇进行位置编码。其实就是给每个词汇加上位置偏移,位置偏移函数选择了sin和cos函数:

Pos是词汇位置,i是词汇向量的维度位置。

3 encoder

由多层的multi-head attention和linear组成,multi-headattention和linear之间由norm和add,add是一个residual连接。

Multi-head attention结构如下:

Q,K,V分别是query,key和value,这是attention机制中抽象出来的三个重要变量,通过计算q和k的相似度,得到每个k对应的v的权重系数,然后对value进行加权求和就得到了attention值。这个是attention机制的本质思想。Transformer中使用softmax函数来描述相似度,当然还有很多其它方法来描述。

这里添加了一个scale1/squart(dk),这其实是一个参数的调节,防止矩阵乘法得到结果太大而导致softmax函数的梯度太小。

这里还要注意transformer网络没有对Q,K,V直接进行单一的attention计算,而是对这三个变量进行了拆分,平行计算拆分后的变量,得到的attention值最后在拼接在一起。

4 decoder

Decoder和encoder也有类似的结构,不同的是,在decoder中由三层:mask-multi-head attention,multi-head attention以及FC构成。带mask的multi-head是为了屏蔽target句子词之后的词,因为对句子的翻译应该是由前向后进行的,后边的词语不应该出现在前边词语的预测之中。

量化方法

量化实际是一个仿射变换:

其中s是scale,q是量化后的数据,z是偏移,如果采用对称变换,令z为0,那么就有:

去除中心z,可以消除矩阵计算中的交叉项。接下来就是如何获得q和s。q和s通过如下方式获得:

Clip操作是在最小值n和最大值p之间获得x/s的向下整数值,如果x/s向下整数值超过n或者p就取n和p。

S的值通过训练获得,为了保证能够很好的在FPGA上计算,s的值最好可以取得2的幂次。

由于s和x都是需要训练的参数,所以我们需要求得他们的梯度值,梯度值比较简单,对q(x, s)的x和x进行求导,有:

对x的梯度使用的是hinton提出的strait-through estimator,这样做是因为可以消除量化引入的噪声,更快的训练。

实践

transformer中有dense,matmul等操作,需要量化的数据有dense中的权重,matmul中的Q,V,K变量。第一次没有什么经验,还是一点点来。首先选择其中一个dense进行量化。从github上下载了一个transformer的实现源码https://github.com/Kyubyong/transformer,这个代码写的很简洁,容易看懂。官方的实现代码比较复杂,需要安装的库较多,曾经也尝试过,因为某些库无法安装成功,所以放弃了。在使用Kyubyong的transformer的时候,也遇到了一个问题,训练可以完成,但是在eval的时候,报了维度的错误,后来找到是在positional encoding的embedding中,经过查找,源码中存在一个bug,就是eval的数据集的maxlen是设置了10000,但是在embedding中传入的查找表维度是从hparams传入的,两者不相同。不知道作者为什么会有这个bug。经过改正可以正常完成eval了。

量化第一步是需要将量化插入到tensorflow的图结构中,即在要量化的权重数据之后。这需要重新定义op和梯度,tensorflow中提供了tf.custom_gradient装饰函数来对梯度和op进行定义,所以我定义了如下梯度:

其中STE_clip中的y计算了对x的量化值,grad函数是对x和s进行梯度计算。X和s分别是传入的(d, d)权重和scale。dy是传入的上一个节点的梯度,所以完成和STE_cllip节点梯度的乘积,这是由函数梯度计算的传递性质决定的。这里需要注意的是,s是一个标量,q(x,s)对s梯度是一个矩阵向量,需要和dy进行点积和。

在tensorflow图构建中,将这个节点插入如下:

这里还添加了tf.print用于打印量化后的数据。

语法错误修正:

1 定义的custom_gradient函数中报NoneType object is not iterable,因为函数没有返回值,默认返回none。

2 TypeError: Input 'e' of 'Select' Op has type float32 that does not match type int32 of argument 't'. 因为使用tf.greater(x, y)x和y应该有相同数据类型。

3 ValueError: Shapes must be equal rank,tf.greater中数据必须具有相同的rank,即维度。

4 ValueError: Shape must be rank 1 but is rank 2,tf.tile(x, axis)中x必须是具有维度的,不能够是0维。

5 ValueError: Shape must be rank 2 but is rank 3,tf.matmul中两个矩阵维度必须相同。

6 TypeError: Failed to convert object of type to Tensor. Contents: [None]. Consider casting elements to a supported type。使用tf.tile的时候,shape必须为tensor量。

7 TypeError: Expected int32, got None of type '_Message' instead. 这是因为输入为[N, T, d_model],其中N开始是none的,所以当使用tf.constant([N,1,1])的时候就会出现错误,因为N是none类型。

8 Incompatible shapes between op input and calculated input gradient。输入的数据和对该输入数据的梯度维度不一致。

9 使用tf.print无法打印出数据。这是因为print是tensorflow中的一个节点,需要将这个节点加入图中,然后才能输出。而且只有计算流经这个print节点,其才会发挥作用。形象的描述应该是:

功能问题:

1 首先就是发现在训练过程中scale和量化数据都没有更新,一直保持不变,而且量化值和权重数据以及scale计算的数据不相同。目前还在查找当中。

引用文献

1 Learning Accurate Integer Transformer Machine-Translation Models,Ephrem Wu

2 Trained uniform quantization for accurate and efficient neural network inference on fixedpoint hardware,Sambhav R. Jain, Albert Gural, Michael Wu, and Chris Dick

3 Attention Is All You Need,Ashish Vaswani,Noam Shazeer,Niki Parmar

4 Quantized Neural Networks: Training Neural Networks with Low Precision Weights and Activations,Itay Hubara,Matthieu Courbariaux,Daniel Soudry

最新文章

最新文章