0x05 计算资源和效率
概述
先前介绍了 Transformer 模型的架构，本节讨论对于 Transformer 模型的参数量等相关计算。
首先回顾一下计算机中基本的数据类型，及其所占的字节数
| 数据类型 | 字节数 |
|---|---|
| int4 | 0.5 bytes |
| int8 | 1 byte |
| int16 | 2 bytes |
| int32 | 4 bytes |
| float16 (fp16，半精度) | 2 bytes |
| bfloat16 (bf16，半精度) | 2 bytes |
| float32 (fp32，单精度) | 4 bytes |
| float64 (fp64，双精度） | 8 bytes |
在大模型领域，经常提及的模型参数量用 B (billion) 来表示，我们自然会问 一个 7B 的模型权重，占用磁盘空间权重是多少呢？
为了回答这个问题，首先需要搞清楚数字进制，
- 1K = 1,000 = 1e3，千
- 1M = 1,000,000 = 1e6，百万
- 1B = 1,000,000,000 = 1e9，十亿
一个 7B 参数量的模型，就包含 七十亿 的参数。另外，还需要看模型的数据类型，比如 7B 的 fp16 存储的模型，和 7B fp32 存储的模型，显然是不同的。在存储计量中，一般按照 byte 也就是字节作为最小的计算单元，有如下的转换：
- 1 KB = 1024 bytes = 2^10 bytes
- 1 MB = 1024 KB = 2^10 KB = 2^20 bytes
- 1 GB = 1024 MB = 2^10 MB = 2^30 bytes
而 2^30 = 1,073,741,824 ≈ 10^9。一般为了简便转换，可以认为 1B 字节 的参数，约等于 2^30 bytes 也就是 1 GB。因此，如果一个 7B 模型是 bf16 格式，那么他的总参数量权重大概就是 14 GB，如果是 fp32 保存的，那么需要占用的磁盘空间大概为 28 GB。
如果考虑到将模型权重加载到显卡进行推理呢，甚至需要进行训练呢？大概可以分别以下几个部分，用表格形式更直观总结如下：
| 使用方式 | 空间 |
|---|---|
| 存储 | 模型权重 |
| 推理 | 模型权重+激活值（+KV-Cache) |
| 训练 | 模型权重+激活值+梯度+优化器 |
到这里读者肯定有很多疑问，比如：
- 模型权重大小如何根据 Transformer 结构计算出来？
- 模型推理中的激活值如何计算？
- 梯度和优化器状态在训练中是如何考虑的？
这些将在后续的章节尝试解答。
Transformer 的参数量
此处以 Decoder Only 的 LLM 为例，其参数主要包括三个部分：Embedding, Decoder, LM head。其中占大头的是 Decoder，一般有多个 decoder-layer 组成，而如前所述，每个 decoder-layer 又分为 Multi-head Self-Attention 以及 FFN 两部分。
首先给出一些符号定义
| 符号 | 说明 |
|---|---|
| \(h\) | 词表维度 / hidden states 维度 |
| \(a\) | 多头注意力中注意力头个数 |
| \(l\) | Decoder Layer 的层数 |
| \(V\) | 词表大小 |
| \(b\) | Batch Size |
| \(s\) | 输入的序列长度 |
Embedding 层
词表是一个二维的矩阵，参数量为 \(V\times h\)。
当一个输入为 \(b\times s\) 输入到 Embedding 进行 Embedding look up 之后，维度会变为 \(b\times s \times h\)
Transformer 层
先计算单层的参数量，然后再乘上层数。每一层的 Decoder Layer 大概有如下几个部分：MHA 中的 \(Q,K,V,O\) 变换矩阵，LayerNorm 中的可学习参数，FFN 中的升维和降维矩阵。
-
MHA
MHA 包含四个权重矩阵，分别用于不同的变换。此处认为每个线性变换都包含偏置项，因此参数量为 \(4\times(h\times h +h) = 4h^2 + 4h\)
-
LayerNorm
每个 LN 包含缩放和平移两个参数，维度均为 \(d\)。一层包含两个 LN，总参数为 \(4h\)。
-
FFN
包含一个升维变换，和一个降维变换，参数量分别为 \(h\times 4h +4h\) 和 \(4h \times h + h\)，共计为\(8h^2 + 5h\)。
一层 Transformer 层的参数量为 \(12h^2 + 13h\)。
LM Head 层
LM Head 是自然语言处理模型中的一个组件，主要作用是将模型的输出（通常是经过 Transformer 编码器处理后的隐藏状态）转换成预测下一个词的概率分布。
Head 与 embedding的参数量相同。如果是 tied embedding（即，head 权重矩阵与词嵌入矩阵是参数共享的），则两者共用一个参数。
总参数量
\(l\) 层的 Transformer 的总参数量为 \(l(12h^2 + 13h) + 2Vh\)，当 \(h\) 很大的时候，可以忽略一次项，近似为 \(12lh^2\)。
LLama3 的参数量估计
再用LLaMA3来看看在工业界落地中的一些特殊之处。
-
SwiGLU
LLaMA 等模型在 FFN 中会使用 SwiGLU 激活，这也就导致其会额外多了一个权重矩阵。LLaMA论文中提到，使用 SwiGLU 后将 \(d_{ffn}\) 从 \(4d\) 降低到了 \(8h/3\)。这样 3 个权重矩阵的参数量还是 \(8h\)，总的参数量依然可以使用 \(12×nlayer×h×h\) 来 预估。
Transformer 的显存占用
根据之前的分析，在训练过程中，显存的占用分为 模型权重+激活值+梯度+优化器 这个四部分。此处以 AdamW 优化器为例，并用混合精度来加速训练，假设模型的参数量为 \(\theta\)，除激活值外，其他几个显存部分占用分别为：
- 模型权重：\(2\theta\)，权重需要加载进显存，同时会保存一份全精度副本，\(4\theta\)，共计 \(6\theta\)
- 梯度：\(2\theta\)，和权重保持一致
- 优化器：\(4\theta + 4\theta\) 全精度的一阶和二阶动量
总计为 \(16 \theta\) 。表格总结如下：
| 类别 | 占用 |
|---|---|
| 模型权重 | \(2\theta + 4\theta = 6\theta\) |
| 梯度 | \(2\theta\) |
| 优化器 | \(4\theta + 4\theta = 8\theta\) |
| 共计 | \(16\theta\) |
Transformer 的激活值
以 Megatron 架构为例，论文中的符号表如下：
| 符号 | 含义 |
|---|---|
| \(a\) | Attention 中，注意力头的个数 |
| \(b\) | 每个 GPU 上的 batch size |
| \(h\) | Transformer 的隐层维度 |
| \(d\) | 单个Attention Head 的隐层维度，即 \(d = h / a\) |
| \(L\) | Transformer Decoder 的层数 |
| \(s\) | 输入的 token 个数 |
| \(V\) | 词表的大小 |
此处计算按照半精度浮点计算，也就是 fp16 或者 bf16 的情况，每个参数占用 2 bytes。根据刚才介绍的 Transformer 架构，主要激活可以大致分为三部分，一部分是 Multi-head Self-Attention，另一部分是 FFN，再者就是两者中间的 Add & LayerNorm 层。下面分别针对这三部分分别计算激活值
Multi-head Self-Attention
| 张量内容 | 张量形状 | 对应激活值 | 保存内容 | 保存原因 |
|---|---|---|---|---|
| \(X\) | \([b,s,h]\) | \(2bsh\) | 输入到 Transformer Block 的张量 \(X\) | 用于后续 \(QKV\) 计算 |
| \(Q、K\) | \([b,s,h]\) | \(4bsh\) | \(X\) 经过转换后的 \(Q\) 和 \(K\) 矩阵 | 用于计算 Attention 的 \(QK^\top\) |
| \(QK^\top\) | \([b,a,s,s]\) | \(2bas^2\) | 每个注意力头的注意力分数 | 用于计算后续的 Softmax 操作 |
| Dropout Mask | \([b,a,s,s]\) | \(bas^2\) | Softmax 后接的 Dropout Mask 矩阵，注意只需要 1 个 byte 的 Mask 矩阵 | 用于经过 Softmax 得到注意力权重后，做 Attention Weights 的 Drop |
| \(V\) | \([b,s,h]\) | \(2bsh\) | \(X\) 经过转换后的 \(V\) 矩阵 | 用于计算 \(\text{Softmax}(\frac{QK^\top}{\sqrt{d}})V\) |
| Attention Weights | \([b,a,s,s]\) | \(2bas^2\) | 经过 \(\text{Softmax}()\) 和 Dropout 的 Attention Weights | 同样用于计算 \(\text{Softmax}(\frac{QK^\top}{\sqrt{d}})V\) |
| Attention Context | \([b,s,h]\) | \(2bsh\) | 经过多头自注意力计算后的张量 | 用于作为后续的输出线性层的输入 |
| Output Mask | \([b,s,h]\) | \(bsh\) | 经过线性层输出的 Dropout Mask 矩阵 | 用于后续 LayerNorm 等的输入 |
| 共计 | \(11bsh+5bas^2\) |
MLP
| 张量内容 | 张量大小 | 对应激活值 | 保存内容 | 保存原因 |
|---|---|---|---|---|
| \(X\) | \([b,s,h]\) | \(2bsh\) | 输入进 MLP 的张量 \(X\) | 用于第一个 Linear 的计算 |
| Linear1 Output | \([b,s,4h]\) | \(8bsh\) | 第一个 Linear 的升维输出 | 用于激活函数的计算 |
| Act Function Output | \([b,s,4h]\) | \(8bsh\) | 激活函数的输出 | 用于第二个 Linear 的计算 |
| Dropout Mask | \([b,s,h]\) | \(bsh\) | Dropout Mask 矩阵 | |
| 共计 | \(19bsh\) |
LayerNorm
Self-Attention 和 MLP 后续分别接了一个 Layer Normalization。每个 LayerNorm 需要保存其输入，大小为 \(2bsh\)。2 个 Layer Norm需要保存的中间激活为 \(4bsh\)。
综上，每个transformer层需要保存的中间激活占用显存大小为\(34bsh+5bas^2\)。对于 \(l\) 层的 Transformer 模型，中间激活占用的显存大小可以近似为 \(l\times (34bsh+5bas^2)\)。
Transformer 的浮点计算量
计算一般即为 FLOPs，在 LLM 的解码阶段，主要涉及的计算和操作是矩阵/向量的乘法。首先看一下，对于一个\((m,n) \times (n,k)\)的矩阵乘法，按照 FLOPs=乘法+加法执行的次数计算，需要执行\(2mkn\) 的计算。
值得注意的是，FLOPs 和 FLOPS 的含义不同，区别如下：
FLOPs：Floating Point Operations 指 浮点运算次数，用于描述计算量。-
FLOPS：Floating Point Operations per Second 指 每秒浮点运算次数，用于描述计算速度，常作为计算卡/显卡的性能指标。 -
QKV 转换
- 输入\(X\in \mathbb{R}^{b\times s\times h}\)，权重\(W_Q,W_K,W_v \in \mathbb{R}^{h \times h}\)
- Flops=\(3 \cdot 2bsh^2 = 6bsh^2\)
- QK的转置得到 Attention Score
- 输入\([b,a,s,d]\)
- 每个 Attention Head 的计算是\(2s^2d\)
- 整个的 Flops 是 \(2bas^2d = 2bs^2h\)
- Softmax乘 V得到 Attention Context
- 输入 \([b,a,s,s]\) \(V=[b,a,s,d]\)
- 每个 Attention Head 的 Flops 为 \(2s^2d\)
- 整个的 Flops 为 \(2bas^2d = 2bs^2h\)
- 输出投影
- 输入 \([b,s,h]\) 权重 \(W_O \in \mathbb{R}^{h\times h}\)
- 整个的 FLOPs 为 \(2bsh^2\)
- Self-Attention 的整体的 FLOPS为
- \(6bsh^2 + 2bs^2h + 2bs^2h + 2bsh^2 = 8bsh^2 + 4bs^2h\)
在考虑 FFN 部分， FFN 主要是由两个线性层组成
- 第一个线性层
- 输入 \([b,s,h]\) \(W_{l1}\in \mathbb{R}^{h\times 4h}\)
- Flops 为 \(2bsh\cdot 4h = 8bsh^2\)
- 第二个线性层
- 输入 \([b,s,4h]\) \(W_{l2} \in \mathbb{R^{4h\times h}}\)
- Flops 为 \(2bs4h \cdot h = 8bsh^2\)
- FFN 部分的整体 FLOPS为
- \(8bsh^2 \cdot 2 = 16bsh^2\)
所以，一层 Transformer Block 的 前向传播的 FLOPs 为 \(24bsh^2 + 4bs^2h\)。反向传播过程的 FLOPs是前向传播过程的 2 倍，也就是\(48bsh^2+8bs^2h\)，因为对一个矩阵乘操作，反向传播需要分别把梯度回传给两个矩阵，按照链式法则会出现两次矩阵乘。
综上，一层 Transformer Block 总体前向+反向的 FLOPs 为 \(72bsh^2 + 12bs^2h\)。
最后，再来考虑一下 Logits 的计算部分。在经过 \(l\) 层的 Transformer 之后，需要将隐层映射为词表，具体将 \([b,s,h]\) 的张量，经过一个 LM Head，转换为 \([b,s,V]\) 的形状，总体的FLOPs为 \(2bshV\)，如果再加上反向传播，那么是 \(6bshV\)。
总体的 FLOPs 为，\(l\times(72bsh^2 + 12bs^2h) + 6bshV\)
结合模型参数量对 FLOPs 训练和推理的估计
一个显示的场景是，给出模型的参数量，以及训练的语料 Token 数量，我们想估计一下，总计需要多少 FLOPS 能够将这些语料过一遍。首先看一下，在推理场景中 FLOPs 和模型参数量的关系，即：
对于常见大模型的参数而言，一般可以认为 \(l(12h+13) \gg 2V\)，例如 Qwen3-8B 中的参数：
- 层数 \(l\) ：36
- 隐层\(h\)：4096
- 词表\(V\): 151936
那么 \(l(12h+13) = 1,769,508 \approx 1.7M\)，\(2V= 303,872 = 0.3M\)，此时推理计算量和参数量的比值近似于 \(2bs\)，而 \(bs\) 实际上就是输入到模型中，一个批次下的 token 总数。因此，推理的 FLOPs 近似为 2 倍的参数量，再乘 token 数量，即：
考虑模型训练的 FLOPs，根据刚才的计算，每个参数反向传播的时，是前向传播的 2 倍，按照上述估计，训练的 FLOPs 为
在获得了总的 FLOPs 后，可以根据 GPU 的参数，估计整体的训练时间，近似按照下面的计算方式：
有不少工具和网站可以快速查阅不同 GPU 的相关参数，如 gpupoet。 关于 A100 PCIe 80GB 的相关信息如下 https://gpupoet.com/gpu/learn/card/nvidia-a100-pcie：
和 Nvidia 官方的 数据 一致：
细心的朋友会发现，同样的一款显卡有不同的型号，而且在性能上差异很大，导致差异的主要原因之一是卡间通讯方式的不同：一个通过主板 PCIe 通讯， 而 SXM 通过 NVLink 互联，带来更高的通信带宽。这也引出下面的内容，LLM 另外一个除了计算量以外的受制约的因素，带宽。
One More Thing: 带宽
推理、训练不仅取决于 GPU 的算力大小，GPU 内存带宽也可能成为瓶颈。模型在训练过程中，参数、梯度需要不断在 GPU内存中进出，也会在不同的 GPU 之间通信。根据 LLM 在计算时候的特点，常把 LLM 在不同时期的运算操作，分为 ”计算密集型“ 和 ”访存密集型“ 两种类型。如何定量化区分这两种类型，可以引入算术强度 (Arithmetic Intensity, AI) 的概念：
举例来说，对于 GEMM ，也就是尺寸为 \((m,n) \times (n,k)\) 矩阵乘法的算数强度为：
其中 \(b\) 指代数据类型占用的字节数。
从直觉直觉来理解，AI 的含义是每搬运 1 字节数据，能做多少计算。因此，AI 越高，说明属于 计算密集型，主要受制于 GPU 的 FLOPS；而 AI 越低，说明是 访存密集型，此时主要受制于 GPU 显存带宽等因素。更严谨则需要结合 roof-line 模型来及时，有兴趣的读者可以深度阅读相关资料 [知乎: Roofline Model与深度学习模型的性能分析]，此处就不再展开。
对 LLM 推理而言，prefill（处理 prompt）阶段主要由大矩阵乘组成，AI 相对更高，往往更接近计算上限；而 decode（逐 token 生成）阶段每步计算规模变小，但仍需频繁读写权重与 KV cache，因此更容易被显存带宽限制。