以 GPT 系列 decoder-only 架构为例，系统推导大语言模型的参数量、训练与推理计算量（FLOPs）以及显存占用，覆盖 Self-attention 与 MLP 的计算量估算及中间激活带来的显存开销。
背景知识
简单介绍下 decoder-only 架构的 llm 结构，只采用 Transformer 模型中的解码器（Decoder）部分，同时 decoder 结构去掉了 Encoder-Decoder attention（Decoder 中的第二个 attention），只保留了 Masked Self-Attention。这里以 gpt1 模型为例，其模型结构如下所示:
gpt 模型结构，llama 在细节上会有所区别，但是主要网络层不会变。
1 | [(masked)multi_headed_attention --> layer_normalization --> MLP -->layer_normalization]* N -> Linear -> softmax -> output probs |
与正常的 Attention 允许一个位置关注/看见到它两边的 tokens 不同，Masked Attention 只让模型看到左边的 tokens：
图： self attention vs mask self attention
在计算模型参数量/计算量之前，我们先定义好一些表示符号：
- $b$: 批量大小
batch_size。 - $s$: 输入序列长度
seq_len，即输入prompt字符串的长度。 - $o$: 输出
tokens数量，用于计算 kv cache 的形状。 - $h$: 隐藏层的维度，也叫$d_\text{model}$，即序列中每个
token的embedding向量的维度。它定义了输入和输出的特征向量的大小，也是模型内部各个组件（特别是注意力机制和前馈网络）操作的主要向量维度。 - $V$：词表大小
vocab_size。也是每个 token 在做 embedding 前的 one-hot 向量维度。 - $l$：模型中 decoder layers 层数，对应 hf 模型配置文件中的 num_hidden_layers。
- $a$：表示多头注意力中有多个头；
- $h_a$：表示hidden_dim_per_head，为多头注意力中每个头的隐藏层维度；
另外，在实际使用时一般都有 $h_a\times a=h$ 成立。
这些变量值都可以在模型配置文件中找到，以 llama-13b 模型配置文件为例，主要字段解释如下：
vocab_size：词汇表中标记的数量，也是嵌入矩阵的第一个维度。hidden_size：模型的隐藏层大小，其实就是 $d_\text{model}$。num_attention_heads：模型的多头注意力层中使用的注意力头数量。num_hidden_layers：模型中的块数（层数）, number of layers。max_sequence_length: 2048, 即代表预训练的 LLaMA 模型的最大 Context Window 只有 2048，也是模型支持的最大输入上下文长度。
后续的 llama2-3/qwen2 模型都用 max_position_embeddings 参数表示模型支持的最大输入上下文长度，比如qwen2.5-3b 模型支持的最大上下文长度为 32768(32k)。
注意，很多 decoder-only 架构的自回归模型的全连接层的偏置 bias 都设置为 False，故这里的计算公式中没有考虑偏置参数。
模型参数量分析
Transformer 模型由 $l$ 个相同的 decoder block 串联而成，每个 decoder block 又由 1 个带掩码（mask）多头注意力（MHA）层、1 个前馈神经网络（FFN）层和 2 个层归一化层组成。
Self-attention参数量
MHA 块有 4 个线性层（全连接层/映射层），对应的是 $Q$、$K$、$V$ 和输出映射层的权重矩阵 $W_Q,W_K,W_V,W_o \in \mathbb{R}^{h\times h}$ 及其偏置，4 个线性层权重参数形状都为 $[h,h]$，偏置形状为 $[h]$，MHA块的参数量为：
$$
\text{Self-attention的参数量} = 4h^2 + 4h
$$
MLP参数量
MLP/FFN 块由 2 个线性层组成，一般第一个线性层完成 $h$ 到 $4h$ 的升维，第二个将 $4h$ 降维到 $h$，对应权重矩阵为 $W_1\in \mathbb{R}^{h\times 4h}$, $W_2 \in \mathbb{R}^{4h\times h}$，偏置形状分别为 $4h$ 和 $h$，MLP块的参数量为：
$$
\text{MLP的参数量} = 8h^2 + 5h
$$
Layer Norm层有两个，分别连接在 MHA 和 MLP 块的后面，layer norm 层有两个可训练参数: $\mu_{\beta}$ 和 $\sigma_{\beta}$（scale factor and offset，缩放和平移参数），参数大小都是 $[h]$，Layer Norm层参数量为：
$$
\text{2 个 Layer Norm 层的参数量} = 4h
$$
Embedding层的参数量
Embedding 层包括两部分: Token Embedding (TE) 和 Positional Embedding (PE)。TE 层的输入张量形状是 $[b, s, V]$，输出维度是 $[b, s, h]$，对应的 TE 层权重矩阵形状为 $[V, h]$，即 TE 层参数量 = **$Vh$**。另外，最后的输出层通常是和 TE 层共享权重矩阵的。PE 层一般使用纯数学计算，无需经过训练，故忽略不计。
参数量总结
综上可知，参数量和输入序列长度无关。对于有 **$n$** 层 decode block 块的 llm 参数量为：
$$
\text{llm 总参数量} = l(12h^2 + 13h) + Vh \simeq 12lh^2
$$
当 **$h$** 较大时，可忽略一次项，llm 参数量近似为 **$12lh^2$**。
不同版本 LLaMA 模型的参数量估算如下：
| 实际参数量 | 隐藏维度$h$ | 层数$l$ | heads 数目 | 预估参数量 $12lh^2$ |
|---|---|---|---|---|
| 6.7B | 4096 | 32 | 32 | 6,442,450,944 |
| 13.0B | 5120 | 40 | 40 | 12,582,912,000 |
| 32.5B | 6656 | 60 | 52 | 31,897,681,920 |
| 65.2B | 8192 | 80 | 64 | 64,424,509,440 |
二 计算量分析
FLOPs：floating point operations 指的是浮点运算次数，一般特指乘加运算次数，理解为计算量，可以用来衡量算法/模型时间的复杂度。对于矩阵 $A \in\mathbb{R}^{1\times n}$ 和 $B \in \mathbb{R}^{n\times 1}$ 的矩阵乘法的 FLOPs 为$ 2n$；对于矩阵 **$A \in \mathbb{R}^{m\times n}$** 和 **$B\in\mathbb{R}^{n\times p}$** 的矩阵乘法的 FLOPs 为 **$2mnp$**。
2.1 MHA 层计算量
prefill 阶段：MHA 块的计算量 + Self-Attention 层 + 多头拼接和线性映射
$$
Q = x W_Q, K = x W_K, V = x W_V
$$
$$
x_{\text{out}} = \text{softmax}\left(\frac{Q K^T}{\sqrt{h}}\right) \cdot V \cdot W_o + x
$$
先分析 MHA 块的计算量：
计算 Q、K、V：对输入矩阵做线性变换，输入 tokens 序列的 embedding 向量的形状为 $[b, s, h]$，做线性变换的权重矩阵 $W_Q$、$W_K$、$W_V$ $\in \mathbb{R}^{h\times h}$，矩阵乘法的输入输出形状为: $[b,s,h] \times [h,h]\to [b,s,h]$，计算量为：
$$
3* 2bsh^2 = 6bsh^2
$$Self-Attention 层，MHA 包含 heads 数目的 Self-Attention 层。
- $QK^T$** 打分计算**：每个头需要计算 Query 和 Key 的点积，所有头的 $QK^T$ 矩阵乘法的输入和输出形状为: $[b,h,s,h_a] \times [b,h,h_a,s]\to [b,h,s,s]$，
- 注意力权重：计算在 $V$ 上的加权 $score\cdot V$，矩阵乘法的输入输出形状: $[b,h,s,s] \times [b,h,s,h_a]\to [b,s,h,h_a]$，
- 多头拼接和线性映射：所有注意力头输出拼接后通过线性映射，concat 不涉及数学运算，只涉及内存操作。矩阵乘法的输入和输出形状为: $[b,s,h] \times [h,h]\to [b,s,h]$，
$$
\text{FLOPs} = 2bsh^2
$$
*综上，prefill 阶段 MHA 块的 FLOPs: *
$$
\text{Prefill_FLOPS}= 6sh^2 + 4s^2h + 2sh^2 = 8sh^2 + 4s^2h
$$
2.2 MLP 层计算量
decode 阶段：MHA 块的计算量 + Self-Attention 层 + 线性映射
$$
x = f_{\text{gelu}}(x_{\text{out}} W_1) W_2 + x_{\text{out}}
$$
先分析 prefill 阶段 Feed-forward（MLP/FFN）层的计算量。包含两个线性层，以及一个 relu 激活层。MLP 两个线性层的权重参数矩阵: $W_1 \in \mathbb{R}^{h\times 4h}$、$W_2 \in \mathbb{R}^{4h\times h}$，MLP 的输入矩阵: $\in \mathbb{R}^{s\times h}$。
- 第一个线性层，矩阵乘法的输入和输出形状为: $[b,s,h] \times [h,4h]\to [b,s,4h]$，
- 第二个线性层，矩阵乘法的输入和输出形状为: $[b,s,4h] \times [4h,h]\to [b,s,h]$，
$$
\text{FLOPs} = 8bsh^2
$$
将上述计算量相加，得到每个transformer层的计算量大约为
$$
\text{Layer_FLOPS}= 6sh^2 + 4s^2h + 2sh^2 + 8sh^2 + 8sh^2 = 24sh^2 + 4s^2h
$$
此外，另一个计算量的大头是logits的计算，将隐藏向量映射为词表大小。矩阵乘法的输入和输出形状为$[b,s,h]×[h,V]→[b,s,V]$，
$$
\text{FLOPs} = 2bshV
$$
因此，对于一个$l$层的transformer模型，输入数据形状为$[b,s]$的情况下，一次训练迭代的计算量为
$$
\text{All_FLOPS}= l\times (24sh^2 + 4s^2h) + 2bshV
$$
三 显存占用量分析
3.1 训练显存占用量
GPU显存包括四个部分：模型参数 + 优化器状态 + 梯度+ 中间激活值
- 模型参数：模型的权重和偏置。
- 优化器状态：如动量和二阶矩估计等，取决于优化器的类型，单纯的 SGD 不占显存。
- 梯度：每个参数对应一个梯度，数据类型与模型参数相同。
- 中间激活值：前向传播过程中产生的激活值，需要在反向传播中使用，其显存占用与 Batch Size、序列长度以及模型架构相关。
- 批量大小（Batch Size）：一次处理的数据样本数量。
- 其他开销：CUDA 上下文、显存碎片等。
在模型训练过程中，存储前向传播的所有中间变量（激活）结果，称为 memory_activations，用以在反向传播过程中计算梯度时使用。而模型中梯度的数量通常等于中间变量的数量，所以 memory_activations = memory_gradients。
假设 memory_model 是指存储模型所有参数所需的内存、memory_optimizer 是优化器状态变量所需内存。综上，模型训练过程中，显存占用量的理论计算公式为：
$$
\text{total_mem} = \text{memory_model} + 2 \times \text{memory_activations} + \text{memory_optimizer}
$$
激活重计算 (activation recomputation，或称为激活检查点 (activation checkpointing) ) 主要的做法是重新计算某些层的激活而不是把它们存在 GPU 内存中，从而减少内存的使用量，内存的减少量取决于我们选择清除哪些层的激活。
训练阶段
以 1B 参数的模型为例，假设训练时使用 Adam，精度使用 FP32，仅考虑与模型参数挂钩的显存计算：
- 模型参数：4 GB
- 梯度：4 GB
- 优化器状态（Adam：一阶和二阶动量）：8 GB（2 × 4 GB）
总显存=4GB（参数）+4GB（梯度）+8GB（优化器状态）=16GB
3.2 推理显存占用量
GPU 显存占用包括三个部分：模型权重 + 中间激活结果 + kv cache：
推理时的显存占用主要包括：
- 模型参数：同训练阶段一致。
- 中间激活值：仅保留当前步的激活值，相较于训练阶段，小非常多。
- 批量大小（Batch Size）：一次处理的数据样本数量。
模型权重显存占用
模型权重参数所需的显存计算公式（params 是模型参数量，参数类型为 fp16）：
$$
\text{memory_model} = \text{params} * 2 = [l(12h^2 + 13h) + Vh] * 2
$$
MHA/MLP中间激活显存占用（额外开销）
每个transformer层包含了一个self-attention块和MLP块，并分别对应了一个layer normalization连接。
MHA 层需要保存的激活值，以及每个激活值的大小：
FFN 层需要保存的激活值，以及每个激活值的大小：
将 MHA 和 FFN 层全部加起来得到：
如果有 $l$ 层 transformer，那么这 $l$ 层 transformer 总的中间激活值占用的显存为：
$$
l * (34bsh+5bas^2)
$$
kv cache 显存占用
假设输入序列的长度为 $s$ ，输出序列的长度为 $o$，decoder layers 数目为 $l$，以 float16 来保存 KV cache，那么 KV cache 的峰值显存占用计算公式（不使用 GQA 优化）为:
$$
\text{memory_kv-cache} = b(s+o)h\times l \times 2 \times 2 = 4blh(s+o)
$$
上式，第一个 2 表示 K/V cache，第二个 2表示 float16 占 2 个 bytes。每个 token 的 kv 缓冲大小 **$ = 4lh$**，单位为字节 byte。
综上分析可知，llm 推理时，gpu 显存占用主要是：模型权重和 kv cache，总显存消耗计算如下:
$$
\begin{aligned}\text{inference_memory} &\simeq [l(12h^2 + 13h) + Vh]*2 + 8bsh + 4lhb(s+o) \
&\simeq 24lh^2 + 4lhb(s+o)\end{aligned}
$$
中间激活和 kv cache 显存和批次大小 $b$ 以及序列长度 $s$ 成正比，在 bs > 某个阈值时，占推理显存大头的是 kv cache。以 llama13b 为例分析，权重参数占用 26GB，当 b = 64, s = 512 时，输出序列长度 o = 512, kv cache 显存占用 = $4lhb(s+o) = 42,949,672,960\ bytes \simeq 42GB$，是模型参数显存的 1.6 倍。
b 的增加能带来近乎线性的 throughput 增加，llm 服务模块的调度策略就是动态调整批次大小，并尽可能让它最大。
3.3 定性分析和定量结论
模型推理阶段，当输入输出上下文长度之和比较小的时候，占用显存的大头主要是模型参数，但是当输入输出上下文长度之和很大的时候，占用显存的大头主要是 kv cache。每个 GPU kv cache 显存所消耗的量和输入 + 输出序列长度成正比，和 batch_size 也成正比。
有文档指出，13B 的 LLM 推理时，每个 token 大约消耗 1MB 的显存。以 A100-40G GPU 为例，llama-13b 模型参数占用了 26GB，那么剩下的 14GB 显存中大约可以容纳 14,000 个 token。在部署项目中，如果将输入序列长度限制为 512，那么该硬件下最多只能同时处理大约 28 个序列。
3.4 LLM 并发支持估算
以集群上的单节点 8 卡 V100 机器运行 llama-13b 模型为例，估算极端情况下聊天系统同时服务 10000 人并发所需要的节点数量。这里的极端情况是指每个请求的输入长度为 512、输出长度为 1536（即上下文长度为 2048）且没有 latency 要求。
LLaMA 系列模型配置文件中 “max_sequence_length”: 2048, 即代表预训练的 LLaMA 模型的最大 Context Window 只有 2048。
结合前面的显存分析章节可知，k、v cache 优化中对于每个 token 需要存储的字节数为 4nh2
对于 llama-13b 模型而言， 其推理时，每个 token 大约消耗 1MB 的显存
对于输入输出上下文长度（512+1536）和为 2048 的请求，其每个请求需要的显存是 2GB。这里对每个请求所需要显存的估算是没有计算推理中间结果所消耗显存（其比较小，可忽略），另外不同框架支持张量并行所需要的额外显存也各不相同，这里暂时也忽略不计。
- 在模型权重为 float16 的情况下，支持的理论 batch 上限为 （32*8-24.6）/ 2 = 115.7。
- 在模型权重为 int8 的情况下，支持的理论 batch 上限为 （32*8-24.6/2）/ 2 = 121.85。（deepspeed 框架不支持 llama 模型的 int8 量化）
以上是理论值即上限值，float16 权重的实际 batch 数量会小于 115.7，目前的 deepspeed 框架运行模型推理时实测 batch 数量只可以达到 50 左右。
10000/50 = 200 (台 8 卡 V100 服务器)。
实际场景中的并发请求具有稀疏性，不可能每个请求都是 2048 这么长的上下文长度，因此实际上 200 台 8 卡 V100 服务器能服务的并发请求数目应该远多于 10000，可能是几倍。
对于 llama-65b 模型而言，其推理时，每个 token 大约消耗 2.5MB（估算的 4***nh*** = 4 * 80 * 8192/(1024 * 1024) = 2.5 MB）的显存，因此，极限情况下每个请求需要的显存是 5GB。
- 在模型权重为 float16 的情况下，支持的理论 batch 上限为 （32 * 8 - 121.6）/ 5 = 26.88。
- 在模型权重为 int8 的情况下，支持的理论 batch 上限为 （32 * 8 - 121.6/2）/ 5 = 39.04。（deepspeed 框架不支持 llama 模型的 int8 量化）
另外，如果输入能量化为 int8 数据类型，理论上支持的 batch 数量会翻倍。
四 结论
对于典型自回归 llm，假设 decoder layers 层数为 $l$，隐藏层大小（Embedding 向量维度）为 *$h$，输入输入数据形状为$ [b, s]$。当隐藏维度 *$h$ 比较大，且远大于序列长度 $s$ 时，则参数量和计算量的估算都可以忽略一次项，则有以下关于参数量、计算量和显存占用计算分析结论。
一些定性结论：
参数量和输入序列长度无关。
每个 token 的 kv cache 占用显存大小是 $4lh$，kv cache 显存占用量随（输入 + 输出序列长度）以及批量大小 batch_size 呈线性增长。
self-attention 的内存和计算复杂度随序列长度$ s$ 呈二次方增长。注意力输出矩阵 $O = \text{softmax}(QK^T)V$ 要求 $O(N^2d)$的 FLOPs，并且除了输入和输出内存之外，需要额外的 $O(N^2)$ 内存。
参考资料
Formula to compute approximate memory requirements of Transformer models
How continuous batching enables 23x throughput in LLM inference while reducing p50 latency
