跳到主要内容

模型尺寸与显存估算

总体公式

推理时的显存占用可以粗略拆成四部分:

总显存 ≈ 权重显存 + KV Cache 显存 + 运行时开销 + 冗余空间

其中最主要的两项通常是:

  • 权重显存:模型参数本身占用的显存。
  • KV Cache 显存:长上下文、高并发场景下会快速增长。

1. 权重显存

权重显存 ≈ 参数量 × 每个参数占用字节数

以 9B 模型为例:

精度字节/参数9B 模型权重显存估算备注
FP3249B × 4 bytes ≈ 36 GB训练或少量特殊场景
FP16 / BF1629B × 2 bytes ≈ 18 GB常见推理精度
INT8 / FP819B × 1 byte ≈ 9 GB需要量化或对应硬件支持
INT4 / 4bit0.59B × 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_sizeKV Cache 精度
324096FP16,2 bytes

请求配置为:

并发tokens/请求总 token 数
440964 × 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
层数32transformer layers
hidden_size4096标准 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 一起算进去。

快速判断

估算显存时可以按这个顺序看:

  1. 先用参数量和精度估算权重显存。
  2. 再根据并发数和上下文长度估算 KV Cache。
  3. 最后根据推理框架、batch 策略和多卡方式预留运行时开销。

如果是长上下文或高并发场景,KV Cache 往往比直觉中更重要。