模型尺寸与显存估算
总体公式
推理时的显存占用可以粗略拆成四部分:
总显存 ≈ 权重显存 + KV Cache 显存 + 运行时开销 + 冗余空间
其中最主要的两项通常是:
- 权重显存:模型参数本身占用的显存。
- KV Cache 显存:长上下文、高并发场景下会快速增长。
1. 权重显存
权重显存 ≈ 参数量 × 每个参数占用字节数
以 9B 模型为例:
| 精度 | 字节/参数 | 9B 模型权重显存估算 | 备注 |
|---|---|---|---|
| FP32 | 4 | 9B × 4 bytes ≈ 36 GB | 训练或少量特殊场景 |
| FP16 / BF16 | 2 | 9B × 2 bytes ≈ 18 GB | 常见推理精度 |
| INT8 / FP8 | 1 | 9B × 1 byte ≈ 9 GB | 需要量化或对应硬件支持 |
| INT4 / 4bit | 0.5 | 9B × 0.5 byte ≈ 4.5 GB | 量化后显存更低,但可能影响效果或速度 |
注意:这是只看模型权重的估算,不包含 KV Cache、框架开销、CUDA workspace、张量并行通信缓冲等额外占用。
2. KV Cache 显存
2.1 标准MHA
标准 Multi-Head Attention 下,可以用下面的简化公式估算:
KV Cache 显存 ≈ 2 × 层数 × hidden_size × 每元素字节数 × 总 token 数
其中:
2:分别代表 Key 和 Value。总 token 数:约等于并发请求数 × 每个请求的上下文 token 数。每元素字节数:FP16 / BF16 通常按 2 bytes 估算。
示例
假设模型配置为:
| 层数 | hidden_size | KV Cache 精度 |
|---|---|---|
| 32 | 4096 | FP16,2 bytes |
请求配置为:
| 并发 | tokens/请求 | 总 token 数 |
|---|---|---|
| 4 | 4096 | 4 × 4096 = 16384 |
那么:
KV Cache 显存
≈ 2 × 32 × 4096 × 2 bytes × 16384
≈ 8.0 GiB
≈ 8.6 GB
这个例子说明:即使权重显存不变,只要并发数或上下文长度增加,KV Cache 也会线性增长。
2.2 GQA / MQA 的影响
上面的公式适用于标准 MHA。如果模型使用 GQA 或 MQA,KV Cache 会更小,应该使用更精确的公式:
KV Cache 显存 ≈ 2 × 层数 × num_key_value_heads × head_dim × 每元素字节数 × 总 token 数
标准 MHA 中:
num_key_value_heads × head_dim = hidden_size
而 GQA / MQA 中 num_key_value_heads 通常小于注意力头数,所以 KV Cache 占用会低于标准 MHA 估算值。
3. 运行时开销和冗余空间
除了权重和 KV Cache,还需要预留运行时开销:
- 小模型 / 低并发:额外预留 10% 到 20%。
- 长上下文 / 高并发 / 多卡并行:额外预留 20% 到 30%,甚至更多。
- 使用 vLLM、TensorRT-LLM、SGLang 等推理框架时,还要考虑框架自己的内存管理策略。
4. 示例
假设我们要部署一个 7B 级别的模型,使用 FP16 推理,并希望支持 8 个并发请求,每个请求最多 4096 tokens。
4.1 已知条件
模型配置假设如下:
| 参数 | 值 | 说明 |
|---|---|---|
| 参数量 | 7B | 约 70 亿参数 |
| 推理精度 | FP16 | 每个参数 2 bytes |
| 层数 | 32 | transformer layers |
| hidden_size | 4096 | 标准 MHA 简化估算 |
| KV Cache 精度 | FP16 | 每个元素 2 bytes |
请求配置如下:
| 参数 | 值 | 说明 |
|---|---|---|
| 并发请求数 | 8 | 同时有 8 个请求 |
| tokens/请求 | 4096 | 输入和已生成 token 合计估算 |
| 总 token 数 | 8 × 4096 = 32768 | 用于估算 KV Cache |
4.2 估算权重显存
权重显存 ≈ 参数量 × 每个参数占用字节数
≈ 7B × 2 bytes
≈ 14 GB
所以,7B FP16 模型仅权重就大约需要 14 GB 显存。
4.3 估算 KV Cache 显存
这里先按标准 MHA 的简化公式估算:
KV Cache 显存
≈ 2 × 层数 × hidden_size × 每元素字节数 × 总 token 数
≈ 2 × 32 × 4096 × 2 bytes × 32768
≈ 16.0 GiB
≈ 17.2 GB
这里可以看到,8 并发、4096 tokens/请求时,KV Cache 已经和模型权重显存处在同一个量级,甚至更大。
4.4 加上运行时开销
先把权重和 KV Cache 加起来:
权重显存 + KV Cache 显存
≈ 14 GB + 17.2 GB
≈ 31.2 GB
再预留 20% 运行时开销和冗余空间:
总显存估算
≈ 31.2 GB × 1.2
≈ 37.4 GB
4.5 结论
这个场景下,单张 24 GB 显卡大概率不够,因为:
- FP16 权重约 14 GB。
- KV Cache 约 17.2 GB。
- 两者相加已经超过 31 GB。
- 再考虑框架运行时开销,估算需要约 37 GB 以上显存。
如果要部署在 24 GB 显卡上,可以考虑:
- 降低并发数,例如从 8 降到 2 或 4。
- 降低最大上下文长度,例如从 4096 降到 2048。
- 使用 INT8 / INT4 量化降低权重显存。
- 使用支持 GQA / MQA 的模型,降低 KV Cache 占用。
- 使用 KV Cache 量化或更激进的推理框架内存优化。
- 多卡部署或换用 40 GB / 48 GB / 80 GB 级别显卡。
这个例子也说明:判断模型能不能跑,不应该只问“7B FP16 是不是 14 GB”,还要把并发、上下文长度和 KV Cache 一起算进去。
快速判断
估算显存时可以按这个顺序看:
- 先用参数量和精度估算权重显存。
- 再根据并发数和上下文长度估算 KV Cache。
- 最后根据推理框架、batch 策略和多卡方式预留运行时开销。
如果是长上下文或高并发场景,KV Cache 往往比直觉中更重要。