Wnma's Blogs
0%

LLM的推理相关计算公式

发表于 更新于 分类于

记录 LLM 在推理上的理论计算公式

参数量

embedding + lm head + transformer block \(\times\) N

  • Embedding / LM Head 的参数量：vocab_size \(\times\) hidden_size
  • transformer block 的参数量
    • self-attention：qkvo 四个线性层，需要分别考虑两种 Attention
      • multi-head attention（MHA）：4 \(\times\) hidden_size \(^2\)
      • grouped query attention（GQA）：

        • q 和 o，2 $$ hidden_size \(^2\)

        • k 和 v，2 $$ hidden_size \(^2\) / num_attention_heads $$ num_key_value_heads

    • FFN：三个线性层，3 \(\times\) hidden_size \(\times\) intermediate_size
    • layer norm：2 \(\times\) 2 \(\times\) hidden_size

以 llama2-13B 为例，关键参数

1
2
3
4
5
6
7
"hidden_size": 5120,
"intermediate_size": 13824,
"mlp_bias": false,
"num_attention_heads": 40,
"num_hidden_layers": 40,
"num_key_value_heads": 40,
"vocab_size": 32000,

由于 num_attention_heads == num_key_value_heads，所以是 MHA

  • Embedding/LM Head： \(32000\times 5120\)
  • transformer block
    • self-attention：\(4\times 5120\times 5120\)
    • FFN: \(3\times 5120\times 13824\)
    • layer norm: \(4\times 5120\)

综上，列出计算

\[2\times 32000\times 5120 + 40\times (4\times 5120\times 5120 + 3\times 5120\times 13824 + 4\times 5120) = 13,016,268,800 \approx 13B\]

如果是 (b)float16，则占用空间/显存为 13 $$ 2 = 26G

如果是 GQA，self-attn 参数量公式为

KV cache 显存计算公式

\(4\times b\times l\times num\_heads\times embed\_size\_per\_head \times (s+n)\)

参数说明

  • b: 句子条数
  • l：层数
  • num_heads：隐层大小 (num_key_value_heads)
  • embed_size_per_head：每个头的大小 (hidden_size / num_attention_heads)
  • s：输入长度
  • n：输出长度
  • 4：k cache+v cache，均为 float16，所以是(1+1)*2

如果只有1条句子，输入+输出 token 长度由 512 -> 1024，则会增加 \(4\times 512 \times l\times h\)：

Llama3-8B：$4 / 1024 / 1024 = 64 $ M

Llama3-70B：\(4\times 512\times 1024\times 80 / 1024 / 1024 = 160\) M

简单来说，每增加一个 token

8B 就会增加 0.125 M 的显存

70B 就会增加 0.3125 M 的显存

计算量

Embedding 可以视作一个哈希表，没有计算量

  • LM Head：\(2 \times b \times s\times hidden\_size\times V\)
  • Self-attention：
    • q_proj: \(2 \times b \times s \times hidden\_size^2\)
    • k_proj和v_proj:
      • MHA: \(2\times(2 \times b \times s \times hidden\_size^2)\)
      • GQA: \(2\times(2 \times b \times s \times hidden\_size^2 / (num\_attention\_heads \times num\_key\_value\_heads)^2 )\)
    • attn_weights: \(2 \times b \times s^2 \times hidden\_size\)
    • attn_output: \(2 \times b \times s^2 \times hidden\_size\)
    • o_proj: \(2 \times b \times s \times hidden\_size^2\)
  • FFN:
    • gate_proj 和 up_proj: \(2\times 2\times b \times s\times (hidden\_size\times intermediate\_size)\)
    • down_proj: \(2\times b \times s\times (intermediate\_size\times hidden\_size)\)

参数说明 - b: 句子条数 - s: 输入长度 - 2: 一次乘法操作和一次加法操作

还是以 llama2-13B 为例，关键参数

1
2
3
4
5
6
7
"hidden_size": 5120,
"intermediate_size": 13824,
"mlp_bias": false,
"num_attention_heads": 40,
"num_hidden_layers": 40,
"num_key_value_heads": 40,
"vocab_size": 32000,

假设 b=1，s=1，则整体计算量有

  • LM Head: \(2\times 5120\times 32000\)
  • self-attention:
    • \(2\times 5120\times 5120\)
    • \(4\times 5120\times 5120\)
    • \(2\times 5120\)
    • \(2\times 5120\)
    • \(2\times 5120\times 5120\)
  • MLP:
    • \(4\times 5120\times 13824\)
    • \(2\times 13824\times 5120\)

模型有 40 层，所以

\(2\times 5120\times 32000+ 40\times(8\times 5120\times 5120+4\times 5120+ 2\times 5120\times 13824+2\times 13824\times 5120) = 20,041,728,000\)

参考资料：

https://blog.csdn.net/wxc971231/article/details/135434478