[Paper] QLoRA 解读：LLM 4-bit 方案与双层量化

算法

2-Level (Double) Quantization

QLoRA 使用了两阶段量化的方案，我们先来说说量化是怎么个流程，需要保存哪些个变量。

First Level Quantization

对于输入的权重，假设其为大小 $R\times C$ 的矩阵，数据格式为 $\texttt{FP16/FP32/BF16}$，将其 flatten 为一维的数组 $w[0\dots N-1], N=R\cdot C$．

我们首先对这个数组按 blocksize 为一组进行切分，划分成多个 blocks．

以下，记 $\texttt{b}$ 对应 block 包含的下标集合为 $\text{block}(\texttt{b})$．

接着，对于每一个 block，求出 max absolute scale absmax[block]

$$\texttt{absmax[block]}=\max_{i\in\text{block}(\texttt{block})} |w_i|$$

求出 absolute scale 后，将权重进行 regularization

$$w_i\gets \frac{w_i}{\texttt{absmax[block]}}$$

于是，权重 $w_i$ 被归一化到 $[-1, 1]$ 了．最后，我们在 NF4 Quantization Look-Up Table $f:\mathbb{N_{[0,15]}}\mapsto\mathbb{R_{[-1,1]}}$ 里，将 $w_i$ 映射为最近的实数，然后记录其在 Lookup Table 中的索引．

所以，其实这里的 NF4 并没有利用其数值意义，而是利用其值域范围只有 $16$ 个数的特点，做 Look-Up Table 而不是数值计算。

Second Level Quantization

由于 $\texttt{absmax[]}$ 本身只是对原始数据分块求的 $\max$，所以其数据格式和输入矩阵的格式是相同的．为了进一步压缩显存，我们再对 $\texttt{absmax[]}$ 进行量化．这就是第二阶段量化

还是相同的套路，只不过这次用 mean 进行 normalization

$$\begin{aligned} \text{offset} &= \mu(\texttt{absmax[]}) \\ a_b &\gets \texttt{absmax[b]}-\text{offset} \end{aligned}$$

然后我们对 $a_b$ 进行 first level quantization，但这里是固定将 $256$ 个 $a_b$ 分为一个 group，对这个 group 求 absolute max．以及，不同于第一阶段只有 $16$ 个固定值，第二阶段量化有 $256$ 个固定值 (code2)