Hugging Face 将 LLM.int8() 8 位矩阵乘法集成到 Transformers 和 Accelerate 中

TL;DR

Hugging Face 宣布，8 位 LLM.int8() 量化方法现已完整集成到 transformers 和 accelerate 库中，使得可以在大约一半的内存占用下推理诸如 BLOOM‑176B 等超大模型，并且在精度上没有可测量的损失。

为什么 8 位量化对大语言模型很重要

大语言模型（LLMs）现在已经超过数千亿参数（例如 PaLM 540B、OPT 176B、BLOOM 176B）。以全精度 FP32 存储模型需要每个权重 4 字节，导致内存需求达到数百 GB，远超大多数 GPU 的容量。将精度降低到半精度（FP16/BF16）可以将内存减半，但仍然使得像 BLOOM 176B 这样的模型占用约 350 GB。使用 8 位整数（INT8）量化可以再实现 2× 的压缩，但传统的粗糙量化往往会降低精度，尤其是对超过约 60 亿参数的模型。

LLM.int8() 的核心思想：零退化矩阵乘法

  1. 异常值提取 – 对隐藏状态矩阵的每列，识别出幅度超过阈值（≈6）的值。
  2. 混合精度矩阵乘法 – 对异常值使用 FP16 进行乘法，其余大部分矩阵量化为 INT8 并使用向量级（激活为行级，权重为列级）量化进行乘法。
  3. 反量化与聚合 – 将 INT8 结果反量化回 FP16 并与异常值的 FP16 结果相加，得到最终的 FP16 输出。

这一三步流程在保持原始 FP16/BF16 模型精确推理质量的同时，将内存使用量降低至四分之一。

量化机制：零点 vs. 绝对最大值

  • 零点量化 将浮点范围（例如 [-1, 1]）缩放到 INT8 范围 [-127, 127] 并对每个值进行四舍五入。逆向缩放可恢复原始值的近似。
  • 绝对最大值量化 将每个张量除以其绝对最大值，再乘以 127 并四舍五入。对于向量 [1.2, ‑0.5, ‑4.3, …, 5.4]，缩放因子为 127/5.4 ≈ 23.5，产生 [-127, 127] 范围的整数值。

这两种方案都可以按行或按列应用，这对于大规模的精确矩阵乘法至关重要。

零退化的实证证据

使用 lm‑eval‑harness 对 OPT‑175B 和 BLOOM‑176B 进行基准测试表明，INT8 与 FP16/BF16 分数之间的绝对差异低于所有任务的标准误差（例如 HellaSwag 准确率 0.7849 对 0.7849，Lambada 困惑度 3.0142 对 3.0152）。在一个案例（BLOOM‑176B 在 Lambada 上）中，INT8 模型甚至略有提升。论文 LLM.int8(): 8‑bit Matrix Multiplication for Transformers at Scale 提供了完整评估。

速度权衡

内存节省会带来对最大模型的轻微慢速：BLOOM‑176B 在 INT8 下比 FP16 慢 15 %–23 %。较小模型（如 T5‑3B、T5‑11B）最初出现更大慢速，但最近的优化将每 token 延迟从 312 ms 降至 173 ms（T5‑3B）以及从 45 ms 降至 25 ms（T5‑11B）。未来版本将继续缩小差距。

模型 精度 GPU 数量 Tokens / ms（批次 1）
BLOOM‑176B BF16 8 × A100 80GB 239
BLOOM‑176B INT8 4 × A100 80GB 282
T5‑11B FP16 2 × T4 15GB 11.7
T5‑11B INT8 1 × T4 15GB 43.5

与 transformers 的集成

关键组件是 bitsandbytes.nn.Linear8bitLt，它是 torch.nn.Linear 的直接替代。最小化的转换工作流如下：

import torch, bitsandbytes as bnb
from bnb.nn import Linear8bitLt

# Define a FP16 model and save its weights
fp16 = torch.nn.Sequential(torch.nn.Linear(64, 64), torch.nn.Linear(64, 64))
torch.save(fp16.state_dict(), "model.pt")

# Build an INT8 version
int8 = torch.nn.Sequential(
    Linear8bitLt(64, 64, has_fp16_weights=False),
    Linear8bitLt(64, 64, has_fp16_weights=False),
)
int8.load_state_dict(torch.load("model.pt"))
int8 = int8.to(0)   # quantization occurs on GPU

在调用 .to 之后，权重会以 int8 张量的形式存储在范围 [-127, 127] 中。原始的 FP16 值可以通过 (weight.CB * weight.SCB) / 127 恢复。

利用 accelerate 实现零内存模型构建

accelerate.init_empty_weights() 在 meta 设备上创建模型，不分配任何 RAM。集成对 accelerate 进行补丁，使参数在脱离 meta 设备时仍保持自定义类 (Int8Params)。递归辅助函数会将每个 nn.Linear 替换为 Linear8bitLt，同时保留诸如 lm_head 之类需要保持全精度的模块：

from accelerate import init_empty_weights
import torch.nn as nn, bitsandbytes as bnb

def replace_8bit_linear(model, threshold=6.0, exclude="lm_head"):
    for name, module in model.named_children():
        if list(module.children()):
            replace_8bit_linear(module, threshold, exclude)
        if isinstance(module, nn.Linear) and name != exclude:
            with init_empty_weights():
                model._modules[name] = bnb.nn.Linear8bitLt(
                    module.in_features,
                    module.out_features,
                    module.bias is not None,
                    has_fp16_weights=False,
                    threshold=threshold,
                )
    return model

两个针对 accelerate 的 PR 确保对每个 INT8 张量仅调用一次 set_module_tensor_to_device，从而避免双重量化错误。

硬件和安装要求

  • GPU 支持 – 需要 INT8 张量核心（NVIDIA Turing、Ampere、RTX 20/30、A40‑A100、T4）。CPU 和较旧的 Kepler GPU 不具备原生支持。
  • 安装 – 使用 Python ≥ 3.8：
pip install accelerate bitsandbytes
pip install git+https://github.com/huggingface/transformers.git

演示

Google Colab 笔记本展示了在仅 11 GB 内存下运行原本需要 42 GB FP32 的 T5‑11B，以及一个可以在单个 T4 上轻松运行的 BLOOM‑3B 示例。

未来工作与局限性

  • 小模型的速度 – 正在进行的工作旨在使 ≤6 B 模型的 INT8 延迟与 FP16 持平。
  • Kepler GPU 支持 – 计划为缺乏原生 INT8 张量核心的 GPU（如 GTX 1080）添加独立的软件栈。
  • 状态字典持久化 – 当前的 INT8 检查点缺少量化统计信息（CB、SCB），导致无法直接从 Hub 加载；添加这些元数据是优先事项。
  • CPU 执行 – CPU 上没有 8 位张量核心；未来的软件路径可能扩大可访问性。
  • 超越文本 – 将该技术扩展到大型视觉、音频和多模态模型是一个开放的研究方向。

致谢：Younes B.、Tim Dettmers 以及原博客文章中列出的贡献者。

Sources