量化推理按显存预算选择权重精度的实现方法

量化推理选精度，先看显存预算，再谈 FP16、INT8 还是 INT4。实际可用显存不能全部给模型权重：KV cache、激活、框架分配和并发请求都要留位置。我的做法是先用参数量做一轮保守估算，再用目标上下文和固定请求集确认，避免模型虽然“加载成功”，一到长输入就频繁溢出。

先把显存预算拆成四个部分

假设一张卡有 24 GB 显存，真正给权重的预算应写成：可用权重预算 = 总显存 - KV cache - 激活 - 框架与并发余量。例如长上下文服务，KV cache 可能随着 token 数和并发增长；只按“模型参数量能不能塞下”做判断，结论会偏乐观。

图中的数字是规划示意，不是某个模型的实测。它表达的是顺序：先固定运行条件，再把剩余空间交给权重精度选择。

用参数量和位宽估算权重占用

第一轮可以只算权重：7B 参数模型按 FP16/BF16 约需 14 GB，按 INT8 约需 7 GB，按 INT4 约需 3.5 GB。它们是理想化下界，还没有计入缩放因子、分组元数据、embedding 或 lm_head 等可能保留更高精度的部分。

因此，24 GB 显存并不意味着 7B INT4 一定能稳定跑满长上下文。若还要承担 8K 上下文、批量请求和加载瞬间的峰值，就应把安全余量提前扣除。反过来，如果是短上下文、单请求离线任务，INT8 可能比 INT4 更容易保住质量和算子兼容性。

把预算转换成可复用的加载配置

Transformers 的量化入口可以通过 BitsAndBytesConfig 声明 8-bit 或 4-bit 加载。下面的代码只展示推理加载边界；模型名称、设备和版本应换成自己的实际目标，不能把示例 footprint 当作压测结果。

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

model_id = "bigscience/bloom-1b7"  # 中文注释：示例模型，实际项目替换成已评估的模型

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,                 # 中文注释：显存紧张时先用 4-bit 权重加载
    bnb_4bit_quant_type="nf4",         # 中文注释：采用 NF4，训练场景还要单独评测质量
    bnb_4bit_compute_dtype=torch.bfloat16,  # 中文注释：计算类型要和硬件能力匹配
    bnb_4bit_use_double_quant=True,    # 中文注释：进一步压缩量化元数据，仍需实测收益
)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",                # 中文注释：让 Accelerate 按设备可用空间分配模块
    dtype="auto",                     # 中文注释：其他模块沿用模型配置的 dtype
    quantization_config=quantization_config,
)

print(model.get_memory_footprint())    # 中文注释：只记录加载后的 footprint，不等于峰值显存

如果 4-bit 加载后仍然触顶，先检查上下文长度、批量和 offload，而不是立刻再降精度。官方说明也区分了推理和训练：8-bit、4-bit 训练主要用于训练额外参数，不能把量化推理配置直接当作完整微调方案。

用固定请求集确认吞吐和精度边界

最后不要只记录“能不能启动”。至少固定一组短输入、长输入和并发请求，分别记录首 token 延迟、生成吞吐、峰值显存、实际上下文长度与任务质量。若 INT4 只在短输入下省显存，长输入时却被 KV cache 限制，那么它解决的是权重问题，不是整体容量问题。

我的落地顺序通常是：预算足够先用 BF16/FP16 做基线；预算开始紧张时试 INT8；只有需要更大模型或更高并发时才进入 INT4。每次只改一个因素，并保留同一模型、同一 tokenizer、同一请求集和同一生成参数，结果才有可比性。

常见问题

按位宽计算出的显存为什么比实际小？

因为公式只覆盖理想化权重。量化元数据、未量化层、激活、KV cache、框架缓存和加载峰值都会增加实际占用。

INT4 一定比 INT8 吞吐更高吗？

不一定。吞吐还受 GPU 架构、量化算子、上下文长度、batch 和内存带宽影响，必须在目标环境中实测。

显存不够时先降精度还是先减少上下文？

如果权重本身已经接近预算，先降位宽；如果权重能容纳而长输入触顶，先减少上下文或并发。两者解决的是不同的内存来源。

参考资料：Hugging Face Transformers 量化总览与 bitsandbytes 文档。量化档位最终应由目标硬件、上下文、并发和质量评测共同决定。