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

godweiyang

发布于 2021-12-02 11:33:35

9800

发布于 2021-12-02 11:33:35

文章被收录于专栏：算法码上来

作者 | godweiyang

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

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

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

那么我们考虑浮点数矩阵乘法

C = AB

，如何将它转化为整数矩阵乘法，并且得到几乎相同的乘法结果呢？

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

首先我们需要将一个浮点数矩阵

用整数矩阵

A_I

来表示。我们假设

的数值范围在

[-c_1, c_1]

之间，其实这个假设是合理的，例如一般深度学习模型参数初始化都是正态分布，那么数值范围就在

[-1, 1]

之间。然后整数矩阵

A_I

的数值范围其实就是有符号整数的表示范围

[-128， 127]

，为了实现的简单，我们只量化到

[-127, 127]

，这样就和

一样关于零点左右对称了。我们令

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

是取过整的，所以还原回去的

并不完全等于原始的

，是有误差的。举个通俗的例子，两个浮点数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的。

总结一下流程

输入两个浮点数矩阵

和

，先分别转化为各自的整数矩阵

A_I = int(\frac{s}{c_1}A)

和

B_I = int(\frac{s}{c_2}B)

。

然后计算整数矩阵乘法结果

A_I B_I

。

最后乘上系数还原为浮点数的乘法结果

C = \frac{c_1c_2}{s^2}A_I B_I

。

进阶（relu激活函数）

熟悉Transformer的同学应该知道，FFN第二层输入分别是relu的结果

和参数

。那么这里就存在一个问题，relu结果的数值范围是

[0, c_1]

，而不可能是

[-c_1, c_1]

。

如果我们强行还按照

[-c_1, c_1]

的范围来量化relu结果

的话会怎么样呢？这样会导致整数区间

[-127, 0)

永远不会有数字，因为根本没有负数浮点数的存在。这样就白白浪费了127个整数，就会导致量化的精度大大受损。

那按照

[0, c_1]

来量化的话，怎么计算整数矩阵乘法的结果呢？

稍稍推导一下就可以得出，

可以表示为

A = \frac{c_1}{2s}(A_I + s \textbf{1})

，其中

\textbf{1}

表示和

相同形状的全1矩阵。而

的话依然表示为

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

的维度和

相同，并且它的矩阵元素等于

中同一列的元素之和。那么问题就很简单了，我们只需要提前计算出矩阵

每一列的元素和，再乘上系数

\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模型量化过程中做的一些尝试。

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2021-11-26，如有侵权请联系 cloudcommunity@tencent.com 删除

神经网络

存储

本文分享自算法码上来微信公众号，前往查看

如有侵权，请联系 cloudcommunity@tencent.com 删除。

本文参与腾讯云自媒体同步曝光计划，欢迎热爱写作的你一起参与！

神经网络

存储

登录后参与评论

0 条评论

热度