量化模型本地推理的显存估算方法

量化模型能不能放进本地显卡，不能只看“模型多少 B”或权重文件大小。一个更稳妥的估算是把显存拆成三块：量化权重、随上下文增长的 KV cache，以及激活和框架分配的运行时余量。先用公式算出下界，再给结果留出空间，才知道 4 bit 是否真的适合当前显卡。

官方资料：https://huggingface.co/docs/transformers/quantization/overview

先记住：权重显存大致等于“参数量 × 每参数字节数”，KV cache 则随 batch、上下文长度、层数和 KV 头数线性增长。量化权重不会自动把 KV cache 和运行时开销一起变小。

先把量化权重和运行时余量分开

如果模型有 P 个参数，权重使用 b bit，理想权重下界可以写成：

weight_bytes ≈ P × b ÷ 8

例如 7B 模型用 4 bit 计算，纯权重约为 3.5 GB（十进制估算）。实际加载还要考虑分组量化的 scales、zero points、未量化层、对齐和框架对象，因此这个数只能当下界。8 bit 通常约为 4 bit 的两倍，16 bit 则约为 4 bit 的四倍。

估算时不要把整个显卡容量都分给权重。可以先扣出 10%～20% 的工程余量，再看剩下的空间能否容纳 cache 和临时张量。这个比例是保守的规划起点，不是任何模型、后端都固定适用的常数。

再按上下文长度计算 KV cache

自回归生成会为每层保存过去 token 的 K、V。对于使用 GQA 或 MQA 的模型，应该使用 KV 头数，不要误把总 attention 头数代入。一个常见的近似公式是：

kv_bytes ≈ 2 × batch × context × layers × kv_heads × head_dim × bytes_per_value

前面的 2 代表 K 和 V；bytes_per_value 取决于缓存 dtype，例如 FP16/BF16 通常按 2 字节估算。上下文从 4K 增加到 8K，其他参数不变时，KV cache 也大致翻倍；batch 增加同样会线性放大缓存。

下面的 Python 片段只用于做纸面估算，不会替代实际后端的显存统计。它保留了 GQA 的 kv_heads 参数，也把 K/V 两份缓存写进公式。

def estimate_vram(params_billion, weight_bits, layers, kv_heads,
                  head_dim, context, batch=1, cache_bytes=2,
                  runtime_ratio=0.18):
    # 先按十亿参数估算量化权重的理论下界
    weight_bytes = params_billion * 1e9 * weight_bits / 8
    # K 和 V 各占一份；KV 头数适配 GQA/MQA 模型
    kv_bytes = (2 * batch * context * layers * kv_heads
                * head_dim * cache_bytes)
    # 运行时余量覆盖激活、临时 buffer 和框架分配
    total_bytes = (weight_bytes + kv_bytes) * (1 + runtime_ratio)
    return {
        "weight_gib": weight_bytes / 2**30,
        "kv_gib": kv_bytes / 2**30,
        "estimated_gib": total_bytes / 2**30,
    }

这个结果适合比较方案：例如固定 4 bit 后，分别把 context 设为 4096、8192，再观察 cache 项的变化。不要把返回值当作某个推理框架的最终峰值，因为后端可能使用不同的临时 buffer、分页缓存或 offload 策略。

用显存余量反推量化档位和上下文

可以按下面的顺序做配置：

1. 先用参数量和位宽估算权重，确认 4 bit、8 bit 或半精度的数量级。
2. 再按目标 batch 和 context 计算 KV cache，优先限制最容易失控的上下文长度。
3. 从显卡容量中扣除运行时余量，剩下的空间才用于权重和 cache。
4. 如果权重刚好塞满，优先换更低位宽或更小模型，不要只把上下文硬改到极限。

Transformers 文档还列出了 Dynamic Cache、Static Cache、Quantized Cache 和 offloading 等策略：动态缓存会随生成增长，静态缓存预留最大长度，量化缓存能进一步降低缓存占用但可能牺牲延迟，offload 则用 CPU 内存换 GPU 空间。它们改变的是缓存管理方式，不会改变“权重 + cache + 运行时开销”的预算思路。

三个容易误判的显存问题

• 把文件大小当峰值显存：文件可能使用打包格式，加载时还会产生元数据、临时张量和未量化模块。
• 只算权重不算 KV：短提示词看不出问题，长上下文或并发 batch 会让 cache 成为主要增量。
• 把总头数当 KV 头数：GQA/MQA 模型的 KV 头数更小，代入错误会把缓存估算放大。

如果已经遇到 OOM，先记录实际目标：模型位宽、batch、最大上下文、缓存 dtype 和是否启用 offload。通常先降低 context 或 batch，再考虑量化 cache；只有在权重本身就超预算时，才需要更换量化档位或模型规模。

相关问题

4 bit 一定比 8 bit 省一半显存吗？
理论权重部分接近，但量化元数据、未量化层、KV cache 和运行时分配不会按同一比例缩小，所以整机峰值不一定刚好减半。

为什么模型能加载却一生成就 OOM？
加载阶段主要占权重；生成阶段会继续写入 KV cache，并分配临时张量，长上下文或 batch 往往在这里触发峰值。