LLM.int8() 量化详解：混合精度分解的 8-Bit 矩阵乘法

讲解 LLM.int8() 的混合精度分解量化：先按列提取激活中的离群特征，对离群维度保留 FP16 高精度计算、对其余部分做 INT8 量化矩阵乘，最后反量化相加，在不掉点的前提下压低显存占用。

技术原理

LLM.int8()（论文：LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale）是一种采用混合精度分解的量化方法。该方案先做了一个矩阵分解，对绝大部分权重和激活用8bit量化（vector-wise）。对离群特征的几个维度保留16bit，对其做高精度的矩阵乘法。

LLM.int8() 通过三个步骤完成矩阵乘法计算:

1. 从输入的隐含状态中，按列提取异常值 (离群特征，即大于某个阈值的值)。
2. 对离群特征进行 FP16 矩阵运算，对非离群特征进行量化，做 INT8 矩阵运算；
3. 反量化非离群值的矩阵乘结果，并与离群值矩阵乘结果相加，获得最终的 FP16 结果。
   实验结果表明该方法效果良好。可以通过使用 LLM.int8() 的量化过程来恢复全部性能。您可以清楚地看到随着模型参数量逐渐变多 8 比特基线（即 vector-wise quantization）的性能大幅下降。 而 LLM.int8() 方法使用vector-wise quantization和混合精度分解来恢复全部性能。

虽然 LLM.in8() 带来的性能下降微乎其微，但是这种分离计算的方式拖慢了推理速度。对于 BLOOM-176B，相比于 FP16，LLM.int8() 慢了大约 15% 到 23%；对于更小的模型（3B 和 11B），速度差距更为明显，LLM.int8() 慢了三倍以上。

在 Transformers 中使用 LLM.int8() 只需提前安装 bitsandbytes 即可，使用 LLM.int8() 方法量化transformer模型具体示例如下：

8bit量化：

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
  'decapoda-research/llama-7b-hf',
  device_map='auto',
  load_in_8bit=True,
  max_memory={
    i: f'{int(torch.cuda.mem_get_info(i)[0]/1024**3)-2}GB'
    for i in range(torch.cuda.device_count())
  }
)

4bit量化：

from transformers import BitsAndBytesConfig

nf4_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_quant_type="nf4",
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype=torch.bfloat16
)

model_nf4 = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=nf4_config)