【白话模型量化系列一】矩阵乘法量化

模型量化是模型加速方向一个很重要的方法,主要思想就是用int8数据格式来存储和进行计算。这样做有两点好处:

  1. 可以减小模型存储的体积。原本float32存储需要4个字节,现在int8存储只需要1个字节,体积是原来的1/4。
  2. 可以加快计算速度。这主要是因为int8数据的读写更快,并且int8矩阵乘法一般来说会更快一点。

以现在最常用的Transformer模型来举例,在使用CUDA推理加速库(例如LightSeq)之后,矩阵乘法的占比高达将近90%。所以优化非矩阵乘法的速度意义不是很大了,占比不高,你算得再快对整体的提速也很小,因此可以尝试优化矩阵乘法。

那么我们考虑浮点数矩阵乘法$C = AB$,如何将它转化为整数矩阵乘法,并且得到几乎相同的乘法结果呢?

用整数矩阵来表示浮点数矩阵

首先我们需要将一个浮点数矩阵$A$用整数矩阵$A_I$来表示。我们假设$A$的数值范围在$[-c_1, c_1]$之间,其实这个假设是合理的,例如一般深度学习模型参数初始化都是正态分布,那么数值范围就在$[-1, 1]$之间。然后整数矩阵$A_I$的数值范围其实就是有符号整数的表示范围$[-128, 127]$,为了实现的简单,我们只量化到$[-127, 127]$,这样就和$A$一样关于零点左右对称了。我们令$s = 127$,用来表示int8的数值范围,如果$s = 15$,那就是int4的范围了。

接着整数矩阵$A_I$就可以表示为$A_I = int(\frac{s}{c_1}A)$,也就是将浮点数区间$[-c_1, c_1]$里的数字等比例映射到整数区间$[-s, s]$,然后向最近的整数取整。同理,整数矩阵$B_I$可以表示为$B_I = int(\frac{s}{c_2}B)$。

这样我们就可以得到两个浮点数矩阵的整数表示,接下来就可以利用他们来进行整数矩阵乘法的转换。

转化为整数矩阵乘法

整数矩阵$A_I$还原为浮点数很简单,只需要$A = \frac{c_1}{s}A_I$即可。但是注意$A_I$是取过整的,所以还原回去的$A$并不完全等于原始的$A$,是有误差的。举个通俗的例子,两个浮点数0.1和0.101经过量化都变成了整数13,但是还原回浮点数后全都变成了0.102,再也没法区分两个浮点数有什么不同了。

所以回到原始的问题,浮点数矩阵乘法$C = AB$可以改写为$C = \frac{c_1}{s}A_I \frac{c_2}{s}B_I$,也就是$C = \frac{c_1c_2}{s^2}A_I B_I$。

那么就可以先计算整数矩阵乘法$A_I B_I$,然后得到整数的输出矩阵之后,乘上系数$\frac{c_1c_2}{s^2}$,还原为浮点数矩阵。

注意输入矩阵$A_I$和$B_I$都是int8的,但是乘法结果$A_I B_I$一定是int32的。

总结一下流程

  1. 输入两个浮点数矩阵$A$和$B$,先分别转化为各自的整数矩阵$A_I = int(\frac{s}{c_1}A)$和$B_I = int(\frac{s}{c_2}B)$。
  2. 然后计算整数矩阵乘法结果$A_I B_I$。
  3. 最后乘上系数还原为浮点数的乘法结果$C = \frac{c_1c_2}{s^2}A_I B_I$。

进阶(relu激活函数)

熟悉Transformer的同学应该知道,FFN第二层输入分别是relu的结果$A$和参数$B$。那么这里就存在一个问题,relu结果的数值范围是$[0, c_1]$,而不可能是$[-c_1, c_1]$。

如果我们强行还按照$[-c_1, c_1]$的范围来量化relu结果$A$的话会怎么样呢?这样会导致整数区间$[-127, 0)$永远不会有数字,因为根本没有负数浮点数的存在。这样就白白浪费了127个整数,就会导致量化的精度大大受损。

那按照$[0, c_1]$来量化的话,怎么计算整数矩阵乘法的结果呢?

稍稍推导一下就可以得出,$A$可以表示为$A = \frac{c_1}{2s}(A_I + s \textbf{1})$,其中$\textbf{1}$表示和$A$相同形状的全1矩阵。而$B$的话依然表示为$B = \frac{c_2}{s}B_I$。

这样矩阵乘法可以改写为$C = \frac{c_1c_2}{2s^2}A_I B_I + \frac{c_1c_2}{2s}\textbf{1}B_I$。其中第二项因子可以用$B = \frac{c_2}{s}B_I$来进一步简化,最终得到$C = \frac{c_1c_2}{2s^2}A_I B_I + \frac{c_1}{2}\textbf{1}B$。

第一项因子和之前一样,先算整数矩阵乘法$A_I B_I$,再乘上系数,只不过系数变成了$\frac{c_1c_2}{2s^2}$。

第二项因子$\textbf{1}B$的维度和$C$相同,并且它的矩阵元素等于$B$中同一列的元素之和。那么问题就很简单了,我们只需要提前计算出矩阵$B$每一列的元素和,再乘上系数$\frac{c_1}{2}$,结果存下来。最后在计算完$\frac{c_1c_2}{2s^2}A_I B_I$整数矩阵乘法结果之后,加上这个列元素之和就行了,你可以将其理解为残差项。

总结

如果矩阵乘法两个输入的范围都是关于零点对称的,那么计算公式为:
量化:
$$A_I = int(\frac{s}{c_1}A), B_I = int(\frac{s}{c_2}B)$$
反量化:
$$C = \frac{c_1c_2}{s^2}A_I B_I$$

如果矩阵乘法其中一个输入是relu的结果,那么计算公式为:
量化:
$$A_I = int(\frac{2s}{c_1}(A - s \textbf{1})), B_I = int(\frac{s}{c_2}B)$$
反量化:
$$C = \frac{c_1c_2}{2s^2}A_I B_I + \frac{c_1}{2}\textbf{1}B$$

当然还有很多其他情况,例如softmax的输出范围一定是$[0, 1]$,那么attention中的矩阵乘法公式还得改写。

此外为了减小量化的损失,还需要在模型结构中插入伪量化节点,然后进行量化感知训练(QAT)。接着还需要将finetune后的模型存储为int8格式。然后还需要开发加载int8模型的推理加速库代码。最后就是本文讲到的整数矩阵乘法了。整个流程比较繁琐,这部分内容今后我会慢慢给大家分享。网上关于量化的优秀教程非常多,我不会讲太多理论上的量化知识,只会从实践的角度来白话一下我们在Transformer模型量化过程中做的一些尝试。


   转载规则


《【白话模型量化系列一】矩阵乘法量化》 韦阳 采用 知识共享署名 4.0 国际许可协议 进行许可。
 本篇
【白话模型量化系列一】矩阵乘法量化 【白话模型量化系列一】矩阵乘法量化
模型量化是模型加速方向一个很重要的方法,主要思想就是用int8数据格式来存储和进行计算。这样做有两点好处: 可以减小模型存储的体积。原本float32存储需要4个字节,现在int8存储只需要1个字节,体积是原来的1/4。 可以加快计算速度
2021-11-25
下一篇 
昨晚学妹参加了B站秋招笔试,还想考考我? 昨晚学妹参加了B站秋招笔试,还想考考我?
学妹昨晚参加了B站的2022届秋招算法笔试,做完给我发来了一道题,想考考我,说挺难的。 我看了两分钟,给她发去了我的思路。然后学妹一眼就看懂了,立马秒过。 那么这道题到底是怎么做的呢? 题目要求将$n$个数切分成$k$块,求每块的序号乘上该
2021-08-26
  目录