0x00 Materials

本文是对LM量化的学习笔记，其中有不少内容是摘自业内的前辈们的文章，在此一并感谢。所参考的资料、摘录的文章来源在下面列出：

经典论文：

解析：

课程：

TinyML and Efficient Deep Learning Computing from MIT-Han-Lab，看量化和剪枝两章。

0x01 前言

数据类型

FP32，FP16，FP8 and BF16, …

这里复习下最基础的计算机组成里面的知识。浮点数可以被符号(Sign)、指数（exponent）、尾数（mantissa）三部分来表示。忘记了可以去看IEEE754 wikipedia去复习下计算方法。

Fig 1. FP32, FP16, BF16,etc...
https://www.servethehome.com/nvidia-h100-hopper-details-at-hc34-as-it-waits-for-next-gen-cpus/nvidia-h100-hopper-fp8-tensor-core/

这里主要是强调表示方法，比如FP8(E5M2、E4M3)，E表示指数、M表示尾数。我们接下来看一下4bit的精度，因为4bit的精度可以更容易说明subnormal numbers。

下面的公式的公式表示了FP的计算方式，$p$是指数，$b$是bias。

$$ f = (-1)^s \times 2^{p-b} \times (1 + \frac{d_1}{2^1} + \frac{d_2}{2^2} + \dots + \frac{d_m}{s^m}) $$

我们首先看E2M2的FP4，FP4不保留符号位，使用完整的4bit来表示数据：

指数	尾数	指数实值$e=2^{p-b}$这里使用$b=0$	尾数实值$m = (1 + \frac{d_1}{2^1} + \frac{d_2}{2^2} + \dots + \frac{d_m}{s^m})$
00	00	$1$	$1$
01	01	$2$	$\frac{5}{4}$
10	10	$4$	$\frac{6}{4}$
11	11	$8$	$\frac{7}{4}$

将4个指数和4个尾数相乘，我们可以得到16个数：

$$ \frac{4}{4}, \frac{5}{4}, \frac{6}{4},\frac{7}{4},\frac{4}{2},\frac{5}{2},\frac{6}{2},\frac{7}{2},4,5,6,7,8,10,12,14 $$

我们发现，上述的16个数字是没办法表示0的，这肯定不是我们想要的表示思路。之所以不表示0，是因为如果我们使用下面的式子：

$$ f = 2^{p-b} \times (0 + \frac{d_1}{2^1} + \frac{d_2}{2^2} + \dots + \frac{d_m}{s^m}) $$

那么如果尾数是0(即二进制的00)，无论指数是什么情况，最终解析出来的fp4数字都是0，也就是说我们为了引入0浪费了3个数据表示，这种情况在只能表示16个数字的fp4中显然是不允许存在的。因此，fp4规则引入了不同的处理方法，即subnormal numbers(Fig 2 中的 x x x x 四个点)。

当指数是0的时候，强行令指数是1，subnormal numbers使用下述公式来计算：

$$f = 2^{1-b} \times (0 + \frac{d_1}{2^1} + \frac{d_2}{2^2} + \dots + \frac{d_m}{s^m})$$

Fig 2. FP4数据分布 from LLM-FP4
2023-10, LLM-FP4: 4-Bit Floating-Point Quantized Transformers

如果不使用subnormal numbers，$b=-2$，那么16个数的分布是：

$$\frac{4}{16},\frac{5}{16},\frac{6}{16},\frac{7}{16},\frac{4}{8},\frac{5}{8},\frac{6}{8},\frac{7}{8},\frac{4}{4},\frac{5}{4},\frac{6}{4},\frac{7}{4},\frac{4}{2},\frac{5}{2},\frac{6}{2},\frac{7}{2}$$

使用subnormal numbers，$b=-2$后是：

$$\frac{0}{8},\frac{1}{8},\frac{2}{8},\frac{3}{8},\frac{4}{8},\frac{5}{8},\frac{6}{8},\frac{7}{8},\frac{4}{4},\frac{5}{4},\frac{6}{4},\frac{7}{4},\frac{4}{2},\frac{5}{2},\frac{6}{2},\frac{7}{2}$$

Int32、Int16、Int8 and Int4 …

整型的编码就十分简单了，没有什么可以讲的。与FP不同的是，整型的每个数是均匀分布的。

Fig 3. 整型 vs 浮点
https://arxiv.org/pdf/2305.12356

0x02 常见量化计算方法

Linear Quantization

Storage: Int

Computation: Int，实际上Int/Float应该是都可以的？

线性量化就是将一组整型映射到实际数值的仿射变换，可以用如下公式来表示：

$$r=S(q-Z)$$

其中$r$表示量化前的数值，$S$表示缩放因子，$q$表示量化后的数值，$Z$表示零点偏移量。如果不对零点进行偏移就是对称量化，但是可能会损失很多的精度；如果对零点进行偏移就是非对称量化，会有更好的量化效果。量化的方式如下图所示：

Fig 4. Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference [Jacob et al., CVPR 2018]
TinyML and Efficient Deep Learning Computing

接下来的问题就是我们怎么求解出$S,Z$:

$$r_{max} = S(q_{max} - Z)$$

$$r_{min} = S(q_{min}-Z)$$

两个式子，两个变量，可以解，解得：

$$S=\frac{r_{max}-r_{min}}{q_{max}-q_{min}}$$

$$Z = \text{Round}(\frac{q_{min}}{S})$$

此时我们来考虑一下量化后的矩阵怎么做矩阵乘法，这在Transformer中是最频繁的操作。对于矩阵乘法：

$$\bf{Y} = \bf{W}\bf{X}$$

$$S_{\mathbf{Y}}\left(\mathbf{q_{Y}}-Z_{\mathbf{Y}}\right)=S_{\mathbf{W}}\left(\mathbf{q_{W}}-Z_{\mathbf{W}}\right)\cdot S_{\mathbf{X}}\left(\mathbf{q_{X}}-Z_{\mathbf{X}}\right)$$

$$\mathbf{q_{Y}}=\frac{S_{\mathbf{W}}S_{\mathbf{X}}}{S_{\mathbf{Y}}}\left(\mathbf{q_{W}}-Z_{\mathbf{W}}\right)\left(\mathbf{q_{X}}-Z_{\mathbf{X}}\right)+Z_{\mathbf{Y}}$$

$$\mathbf{q_{Y}}=\frac{S_{\mathbf{W}}S_{\mathbf{X}}}{S_{\mathbf{Y}}}\left(\mathbf{q_{W}}\mathbf{q_{X}}-Z_{\mathbf{W}}\mathbf{q_{X}}-Z_{\mathbf{X}}\mathbf{q_{W}}+Z_{\mathbf{W}}Z_{\mathbf{X}}\right)+Z_{\mathbf{Y}}$$

这里的$S_{\bf{Y}}, Z_{\bf{Y}}$是在PTQ中使用一个验证集来计算的全局值。比如W4A4，也就是$\bf{W},\bf{X}$都使用4bits量化，然后计算的时候使用int32，然后此时就可以运用预先计算好的$S_{\bf{Y}}, Z_{\bf{Y}}$了，算完后再计算量化后的$\bf{Y}$。如下图：

优化 1：Fix-Point Multiplication

在实验中发现，$\frac{S_{\mathbf{W}}S_{\mathbf{X}}}{S_{\mathbf{Y}}} \in (0,1)$，所以在实现的时候我们可以实用定点数而非浮点数来实现。

$$\frac{S_{\mathbf{W}}S_{\mathbf{X}}}{S_{\mathbf{Y}}} = 2^{-n} M_0, M_0 \in [0.5,1)$$

其中，$M_0$是定点数。而定点数可以实用Int来模拟，Int的计算会比浮点数快得多。

优化 2：$Z_{\bf{W}} = 0$，对称量化

不难发现，若$Z_{\bf{w}} = 0$，我们可以省下很大一部分的计算。如果参数矩阵$\bf{W}$的值近乎于正态分布，也就是说均值为0，那么我们就可以不使用$Z_{\bf{w}}$。这时候的$S$就非常的直观了:

$$S = \frac{\vert r \vert _{max}}{2^{N-1}}$$

其中，$N$是量化的位数。非常好理解，就是要把最大的动态范围的$r$尽数映射到$N$bits量化能够表示的全部空间内部。Pytorch 和 ONNX 的native量化实现就是用的这个方法。优化后，我们得到简化的矩阵乘法公式：

$$\mathbf{q_{Y}}=\frac{S_{\mathbf{W}}S_{\mathbf{X}}}{S_{\mathbf{Y}}}\left(\mathbf{q_{W}}\mathbf{q_{X}}-Z_{\mathbf{X}}\mathbf{q_{W}}\right)+Z_{\mathbf{Y}}$$

含有Bias的全连接层

$\bf{Y} = \bf{W}\bf{X} + \bf{b}$

$S_{\mathbf{Y}}\left(\mathbf{q_{Y}}-Z_{\mathbf{Y}}\right)=S_{\mathbf{W}}\left(\mathbf{q_{W}}-Z_{\mathbf{W}}\right)\cdot S_{\mathbf{X}}\left(\mathbf{q_{X}}-Z_{\mathbf{X}}\right) + S_{\mathbf{b}}\left(\mathbf{q_{b}}-Z_{\mathbf{b}}\right)$

我们可以继续使用上述的两种优化方法，令$Z_{\bf{b}}=0,Z_{\bf{w}} =0$:

$$S_{\mathbf{Y}}\left(\mathbf{q_{Y}}-Z_{\mathbf{Y}}\right)=S_{\mathbf{W}}\left(\mathbf{q_{W}}\right)\cdot S_{\mathbf{X}}\left(\mathbf{q_{X}}-Z_{\mathbf{X}}\right) + S_{\mathbf{b}}\left(\mathbf{q_{b}}\right)$$

$$S_\mathbf{Y}\left(\mathbf{q_Y-Z_Y}\right)=S_\mathbf{W}S_\mathbf{X}\left(\mathbf{q_Wq_X-Z_Xq_W}\right)+S_\mathbf{b}\left(\mathbf{q_b}\right)$$

可以继续损失一点精度，令$S_{\mathbf{b}}= S_\mathbf{W}S_\mathbf{X}$:

$$S_\mathbf{Y}\left(\mathbf{q_Y-Z_Y}\right)=S_\mathbf{W}S_\mathbf{X}\left(\mathbf{q_Wq_X-Z_Xq_W}+\mathbf{q_b}\right)$$

$$\mathbf{q_{Y}}=\frac{S_{\mathbf{W}}S_{\mathbf{X}}}{S_{\mathbf{Y}}}\left(\mathbf{q_{W}}\mathbf{q_{X}}+\bf{q}_b-Z_{\mathbf{X}}\mathbf{q_{W}}\right)+Z_{\mathbf{Y}}$$

其中，我们可以预先计算$\bf{q}_{bias}$:

$$\bf{q}_{bias} = \bf{q}_b - Z_{\bf{X}}\bf{q}_{\bf{W}}$$

$$\mathbf{q_{Y}}=\frac{S_{\mathbf{W}}S_{\mathbf{X}}}{S_{\mathbf{Y}}}\left(\mathbf{q_{W}}\mathbf{q_{X}}+\bf{q}_{bias}\right)+Z_{\mathbf{Y}}$$

如何计算Activation的$Z,S$

Type 1. 训练时计算

使用Exponential Moving Averages(EMA)。在训练的时候计算每组Activations的$S, Z$，然后使用EMA来累计得到我们最终需要的$S,Z$。

$$\hat{r}_{\max,\min}^{(t)}=\alpha\cdot r_{\max,\min}^{(t)}+(1-\alpha)\cdot\hat{r}_{\max,\min}^{(t-1)}$$

Type 2. 在训练后，使用一些样本来推理

每一次输入一个Batch给网络，然后计算出Activations的平均$r_{min}, r_{max}$。
使用KL散度等方法来找到Clipping的准确的点。

Fig 7. Activations Values的统计
8-bit Inference with TensorRT [Szymon Migacz, 2017]

上图是每一个Activations Values的统计，横坐标是Activations的值，纵坐标是对应一个Value区间的Activations在整个网络中出现的次数。我们希望找到一个合理的截断点，使得量化损失最小，那么使用KL散度，将量化后的输出和FP32模型的输出做Loss就知道Clip的位置在哪里合适了。

Clip掉的越多，低精度的数据类型映射到原始的数据精度的区域就更细致，如下图。但是卡掉的Outliers越多也会导致模型的精度下降。

Fig 9. Clipping
Optimal Clipping and Magnitude-aware Differentiation for Improved Quantization-aware Training [Sakr et al., ICML 2022]

如何Round？四舍五入是好方法吗？

四舍五入的方法肯定不是最优的。因为Tensor中的每一个值是跟周围的值有关系的，我们要考虑他周围的值。

Fig 10. Rounding
Up or Down? Adaptive Rounding for Post-Training Quantization [Nagel et al., PMLR 2020]

我们希望知道每一个值到底是下取整好还是上取整好，我们可以用一个可以学习的Bias来学习得到如何Rounding。如：

$$\tilde{w}=\lfloor\lfloor w\rfloor+\delta\rceil,\delta\in[0,1]$$

在优化的时候，我们使用下述式子：

$$\argmin_{\mathbf{V}}\Vert \mathbf{W}\mathbf{x}- \lfloor\lfloor\mathbf{W} + \mathbf{h}(\mathbf{V})\mathbf{x}\Vert_\mathbf{F}^2+\lambda f_{\text{reg}}(\mathbf{V})$$

$\mathbf{x}$是每个Layer的输入，$\bf{V}$是需要被优化的参数。$\mathbf{h}()$是一个将值映射到$(0, 1)$的函数，比如rectified sigmoid。

$f_{\text{reg}}(\mathbf{V})$是一个正则化的函数，鼓励$\bf{h}(\bf{V})$的结果是二值的。

$$f_{\text{reg}}(\mathbf{V}) = \sum_{i,j}1- \vert 2\bf{h}(\bf{V}_{i,j}) - 1 \vert ^ {\beta}$$

K-Means Based

Fig 11. K-Means Compression
Deep Compression [Han et al., ICLR 2016], from han lab.

K Means方法比较直白，就是对weight做K-Means聚类，把weight转换成index和lookup table。

0x03 量化粒度

下面介绍三种PTQ（Post-Training Quantization）中使用的量化粒度。

Per-Tensor Quantization

$$\vert r \vert _{max} = \vert \bf{W} \vert _{max}$$

缩放系数$S$是在整个Tensor层面上进行的。但是对于Outliers的兼容性不好，我们需要更加细的量化粒度。

Fig 13. MobileNetV2 Tensor Per-Channel Range
Data-Free Quantization Through Weight Equalization and Bias Correction [Markus et al., ICCV 2019]

从图中MobileNetV2的一个Tensor的Per Channel Range中可以看出每一个Channel的动态范围非常不一样，对于一个Tensor使用同一个$S,Z$的粒度太粗了。我们可以考虑对每一个Channel做量化操作，即量化的粒度下放到Channel的程度。

Per-Channel Quantization

以2bit为例，将一个Tensor的每一个Channel进行缩放。

$$S = \frac{\vert r \vert _{max}}{q_{max}}$$

Group Quantization

Group Quantization是一种细粒度更小的量化方法，如下图中仅对一个框定的Vector进行量化。Group Quantization通常和层次化的量化方法一起使用，读者可以进一步看这篇文章来了解层次化的量化方法：2021-02, VS-Quant: Per-Vector Scaled Quantization for Accurate Low-Precision Neural Network Inference [Steve Dai, et al.]。一个简单的描述是：

$$r = S(q-Z) \to r = \gamma \cdot S_q(q-Z)$$

其中$\gamma$是一个浮点数的Per-Tensor缩放系数，$S_q$是一个Integer的Per-Vector的缩放系数

这样的量化方法可以被多层次使用，如下图所示：

Fig 17. VSQ
VS-Quant: Per-Vector Scaled Quantization for Accurate Low-Precision Neural Network Inference [Steve Dai, et al.]

上图中，VSQ的处理逻辑是：

先对四个元素一组的元素做Group Quantization，使用的$S$为UINT4类型的。
然后对整个Channel做Per-Channel的量化，使用的$S$为FP16类型的。

在计算Effective的时候，上图中是忽略了Per-Channel的16个bits的，顾仅仅是4（Per-Group的UINT4） / 16。

Fig 18. MX Type
With Shared Microexponents, A Little Shifting Goes a Long Way [Bita Rouhani et al.]

上图中描述的MX量化方法用了共享指数部分的操作，以MX4为例：

每个元素为S1M2，我们发现Exponent不见了，其实Export是L0 Group 量化的$S$，每2个元素共用一个Exponent。
在第二层次的Group量化也是这样的，只不过是每4个元素共用一个8bits的Exponent。

0x04 Quantization-Aware Training

在训练的时候需要保留FP32的副本，因为我们希望梯度被累计。比如梯度传回来+0.1，但是如果不保留FP32的副本，Int值的量化参数会Round掉+0.1导致网络无法收敛。使用FP32后，当梯度传播回来的值累积到一定的程度后，Round后的量化参数会发生变化。所以前向走量化参数，反向走FP32的全量副本。

0x05 Advance

Advance的内容可以看博客，这里会不定期更新链接：

0x00 Materials#

0x01 前言#

数据类型#

FP32，FP16，FP8 and BF16, …#

Int32、Int16、Int8 and Int4 …#

0x02 常见量化计算方法#

Linear Quantization#

优化 1：Fix-Point Multiplication#

优化 2：$Z_{\bf{W}} = 0$，对称量化#

含有Bias的全连接层#

如何计算Activation的$Z,S$#

Type 1. 训练时计算#

Type 2. 在训练后，使用一些样本来推理#

如何Round？四舍五入是好方法吗？#

K-Means Based#

0x03 量化粒度#

Per-Tensor Quantization#

Per-Channel Quantization#

Group Quantization#

0x04 Quantization-Aware Training#

0x05 Advance#