GPU显存计算

获取基本参数

  1. 获取参数量

  2. 获取 KV 的 kv_dim 和 num_KV_heads

  3. 获取支持的最大token数(max_position_embeddings)

算显存占用

  • 总公式

= <KV缓存大小> * <KV缓存数量>

  • 一个 <KV缓存大小> 公式

= <KV_dim> * <num_KV_heads> * 2

以 Qwen2.5-7B 为例

获取基本参数

config.json 文件内容:

{
  "architectures": [
    "Qwen2ForCausalLM"
  ],
  "attention_dropout": 0.0,
  "bos_token_id": 151643,
  "eos_token_id": 151645,
  "hidden_act": "silu",
  "hidden_size": 2048,
  "initializer_range": 0.02,
  "intermediate_size": 11008,
  "max_position_embeddings": 32768,
  "max_window_layers": 70,
  "model_type": "qwen2",
  "num_attention_heads": 16,
  "num_hidden_layers": 36,
  "num_key_value_heads": 2,
  "rms_norm_eps": 1e-06,
  "rope_theta": 1000000.0,
  "sliding_window": 32768,
  "tie_word_embeddings": true,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.43.1",
  "use_cache": true,
  "use_sliding_window": false,
  "vocab_size": 151936
}

  • 获取关键参数:

    • 模型维度相关

      • “hidden_size”: 2048

      • “intermediate_size”: 11008 # FFN 层的中间层维度

      • “num_hidden_layers”: 36 # Transformer 层数

      • “num_attention_heads”: 16 # 注意力头数

      • “num_key_value_heads”: 2 # 启动GQA机制，表示只用 2 个 key/value 头，节省 KV-cache 开销。

    • 注意力机制与位置编码

      • “max_position_embeddings”: 32768 # 支持的最大 token 数

    • Token 配置

      • “vocab_size”: 151936 # 词表大小

    • 激活函数

      • “hidden_act”: “silu” # 对应的 FFN 乘以 3

  • 计算用到的其他参数

    • group_size: num_attention_heads/num_key_value_heads

      • = 16/2 = 8

    • head_dim: hidden_size / config.num_attention_heads

      • = 2048 / 16 = 128

    • kv_dim: num_key_value_heads × head_dim

      • = 2 × 128 = 256

计算参数量

参数量：
    1. 嵌入层：151936 * 2048
    2. 注意力层：36 * 2048 * 2048 * (2 + 2/8)
    3. FFN层：36 * 2048 * 11008 * 3
    4. 输出层：2048 * 151936
总参数量：
    = 151936 * 2048 + 36 * 2048 * 2048 * (2 + 2/8) + 36 * 2048 * 11008 * 3 + 2048 * 151936
    = 151936 * 2048 * 2 + 36 * 2048 * (2048*2.25 + 3*11008)
    = 151936 * 4k + 72 * (2048*2.25 + 3*11008)k
    = 593.5M + 72 * (4.5 + 32.25)M
    = 3239.5M

显存占用

  • 以 FP16 为例

  • 模型参数占显存:

    • = <模型参数量> * 2

    • = 3239.5M * 2

    • = 6479M

  • KV缓存占显存:

    • 总公式 = <KV缓存大小> * <KV缓存数量> * 2

    • 一个 <KV缓存大小>

      • = <KV_dim>

      • = <KV_head> * <num_KV_heads>

      • = 128 * 2

      • = 256

    • 一个K, 一个V 有

      • = 2 * 256

      • =512

    • 使用FP16

      • = 512 * 2

      • = 1024

    • Transformers 层数为 36

      • = 36 * 1024

      • = 36 K

    • Sequence_size:

      • 如果序列长度为 1024

        • = 1024 * 36K

        • = 36M

      • 最大max_position_embeddings: 32768

        • = 32768 * 36K

        • = 1152M

        • 约 1.152G

  • 说明1：batch_size

    • 如果有 batch_size 还要乘以 batch_size

  • 说明2：不使用 GQA 的情况

    • 首先模型参数会增加

    • KV Cache方面

      • 一个 <KV缓存大小> = 2048

      • 比不使用 GQA 占用显存多 8 倍