Transformers量化模型降低显存占用

我在本地尝试加载语言模型时，最先遇到的通常不是推理代码，而是模型权重还没有放进显存，进程就已经报 CUDA out of memory。Transformers 的量化入口可以把这个问题前移：在 from_pretrained() 加载权重时，通过 BitsAndBytesConfig 选择 8-bit 或 4-bit 表示。

官方地址：https://huggingface.co/docs/transformers/

本文结论：如果目标是先把较大的模型放进现有 GPU，8-bit 是更稳妥的起点；如果显存仍然不足，或者要给 QLoRA 一类参数高效微调留空间，再考虑 4-bit。量化只改变主要权重的表示方式，激活、KV cache、临时张量和运行时分配仍然占用显存。

Transformers 量化不是把“模型大小”简单除以 4，而是在加载阶段降低权重的位宽，再由计算 dtype、硬件后端和任务质量要求共同决定最终是否适合。

1. 先估算精度变化带来的权重占用

模型有多少参数，乘以每个参数占用多少字节，可以得到一个很粗的权重空间估算。bf16 或 fp16 通常按每个参数 2 字节计算，8-bit 约为 1 字节，4-bit 约为 0.5 字节。因此，一个 8B 参数模型的权重空间大约是 16GB、8GB 和 4GB 这三个量级。

这个估算不能当成显存实测。加载过程还会出现临时张量，推理时还要保存激活和 KV cache；如果使用较长上下文，KV cache 甚至可能成为新的主要开销。量化的价值是先降低最容易预测的权重部分，而不是承诺任何输入长度下都能刚好塞进显卡。

2. 准备 Transformers 与 bitsandbytes

Transformers 的 bitsandbytes 集成需要加载器、Accelerate 和量化后端一起工作。依赖安装可以从一个干净环境开始：

# 安装模型加载器、设备分配组件和 bitsandbytes 后端
python -m pip install -U transformers accelerate bitsandbytes

# 查看当前环境中的关键包版本，便于定位加载失败
python -m pip show transformers accelerate bitsandbytes

安装成功并不等于所有硬件都适合。bitsandbytes 的不同量化路径对 GPU、CPU、XPU 或其他后端有不同支持范围；如果机器没有可用的兼容后端，错误往往出现在真正加载模型时。先确定目标设备，再决定是 8-bit 还是 4-bit，通常比只复制一段配置更省时间。

如果文章讨论的是某个具体模型，还要确认模型卡允许的加载方式、权重格式和访问权限。量化配置负责告诉 Transformers 如何加载，不会替模型补齐缺失的访问令牌，也不会自动修复模型架构与当前类不匹配的问题。

3. 用 BitsAndBytesConfig 加载量化模型

8-bit 版本适合先验证链路。下面以一个支持 CausalLM 接口的模型标识为例；实际使用时，把 MODEL_ID 换成目标模型的仓库名，并根据模型卡选择对应的 AutoModel 类。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

MODEL_ID = "your-org/your-causal-model"

# 让 tokenizer 与模型使用同一个仓库配置，避免词表和权重错配
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# 8-bit 主要降低权重占用，适合作为第一次量化尝试
quant_config = BitsAndBytesConfig(load_in_8bit=True)

# device_map=auto 让 Accelerate 根据可用设备安排量化后的模块
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=quant_config,
    device_map="auto",
)

# 这不是完整显存实测，但能帮助确认模型对象的内存估算
print(model.get_memory_footprint())

这里的关键点有三个：quantization_config 把量化策略交给模型加载器，bitsandbytes 提供对应的量化线性层实现，device_map="auto" 负责安排模块位置。它们是相互配套的入口，不能只把 load_in_8bit=True 放进一个普通字典里就期待模型自动量化。

如果 8-bit 仍然不够，可以把配置改成 4-bit。推理时常用 NF4 作为权重类型，并单独指定计算 dtype：

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# NF4 是 4-bit 权重表示；计算仍可使用 bfloat16，减少低精度计算的额外损失
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

# 量化发生在加载阶段，设备分配仍由 Accelerate 配合 device_map 处理
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=quant_config,
    device_map="auto",
)

# 读取模型自身的占用估算，不把它当成整段程序的显存峰值
print(model.get_memory_footprint())

bnb_4bit_compute_dtype 不是把 4-bit 权重重新变成 4-bit 计算。它描述的是计算时使用的 dtype；权重存储、计算类型和非量化模块的类型是三个不同概念。若设备不支持 bfloat16，应按目标硬件改用兼容的 dtype，而不是只修改模型名称。

4. 按任务选择 8-bit 或 4-bit

我会把选择分成三层。第一层是显存：8-bit 通常已经能把权重空间压到原来的一半左右，4-bit 还能继续压缩。第二层是质量和兼容性：位宽越低，越需要用实际任务的输入样本检查输出，而不能只看“能否加载”。第三层是任务类型：官方文档对 8-bit/4-bit 权重的训练支持主要针对额外参数，因此全量训练和参数高效微调不能混为一谈。

还有一个容易忽略的取舍：如果模型已经是预量化检查点，通常直接调用 from_pretrained()，让 Transformers 从检查点配置中读取量化方法；如果手里是全精度检查点，才是在加载时传入相应配置。不要把“加载预量化权重”和“加载时量化”当成同一条路径。

5. 用内存统计与常见错误收尾

get_memory_footprint() 可以帮助判断模型对象的估算占用，但它不是 nvidia-smi 的进程峰值。检查时最好同时记录模型配置、输入长度、生成长度和设备布局，这样才知道显存上涨来自权重、KV cache 还是临时张量。

常见错误可以按下面的边界处理：

• 导入或后端错误：先检查 transformers、accelerate、bitsandbytes 版本和当前硬件支持，别先改模型代码。
• 加载时仍然 OOM：先缩短输入和生成长度，再确认是否启用了量化；权重节省不代表 KV cache 没有上限。
• 输出质量下降：固定一组代表性问题，对比 bf16/fp16、8-bit 和 4-bit 的结果，不要用单次回答下结论。
• 想做全量训练：重新检查训练目标。8-bit/4-bit 通常用于推理或只训练额外参数的方案，不能直接等价为全量低精度训练。
• 自动分配位置不符合预期：把 device_map 当成部署决策的一部分记录下来，必要时再为关键模块提供明确映射。

总结

Transformers 降低显存占用的最小路径是：先按位宽估算权重空间，再安装兼容的 bitsandbytes 后端，用 BitsAndBytesConfig 传给 from_pretrained()，最后用模型占用统计和固定样本检查实际效果。8-bit 更适合稳妥地验证推理链路，4-bit 更适合显存紧张或需要为额外参数训练让出空间的场景。

官方资料入口：

• https://huggingface.co/docs/transformers/quantization/overview
• https://huggingface.co/docs/transformers/main/quantization/bitsandbytes
• https://huggingface.co/docs/transformers/main_classes/quantization

相关问答

4-bit 一定比 8-bit 更快吗？

不一定。4-bit 首要收益是降低权重占用，实际速度还取决于 GPU、量化内核、输入长度、生成长度和数据搬运。应该用目标模型和目标输入做基准，而不是从位宽直接推导吞吐。

量化后为什么显存仍会继续上涨？

生成阶段会保存 KV cache，并创建中间激活和临时张量；上下文或输出越长，额外占用越明显。量化主要压缩权重，不能消除这些运行时对象。

已经有 4-bit 模型，还要再传 BitsAndBytesConfig 吗？

不一定。预量化检查点通常把量化配置写在检查点的配置文件中，Transformers 可以据此选择对应加载器。只有需要覆盖加载选项时，才根据该量化方法的文档传入同类型配置。