大模型显存计算指南 - 推理与训练显存计算详解

📚 在部署和训练大模型时,显存往往是最大的瓶颈。本文将帮助你理解显存计算的关键要素,让你能够更好地规划硬件资源。

一、推理显存计算 - 如何评估模型部署所需显存?

1. 影响推理显存的关键因素

(1) 模型参数

• 基础显存占用 = 参数量 * 精度字节数
• 例如:7B模型在不同精度下的参数占用
  • FP32(4字节): 7B * 4 = 28GB
  • FP16(2字节): 7B * 2 = 14GB
  • INT8(1字节): 7B * 1 = 7GB
  • INT4(0.5字节): 7B * 0.5 = 3.5GB

(2) 注意力缓存(Attention Cache)

• KV Cache大小 = batch_size * num_layers * 2 * seq_length * hidden_size * precision
• 对于长文本生成,注意力缓存可能占用大量显存

(3) 激活值(Activations)

• 模型推理过程中的中间计算结果
• 通常占用基础参数量的10%-20%显存

2. 实际显存估算示例

以Qwen-7B为例(hidden_size=4096):

def calculate_inference_memory(
    batch_size=1,
    seq_length=2048,
    model_size_b=7,
    precision="fp16"
):
    # 精度映射
    precision_map = {
   
        "fp32": 4,
        "fp16": 2