序列模型（3）—— LLM的参数量和计算量

• 本文说明以下重要结论

  设模型参数量为 $N$，训练数据量（Token）为 $D$，LLM训练中计算量（FLOPs）$C\approx 6ND$

• 参考：
  1. 模型训练计算量到底怎么算
  2. 分析transformer模型的参数量、计算量、中间激活、KV cache

---

0. 符号约定

• 本文遵循以下符号约定

  参数	符号	说明
  Decoder 层数	$l$
  Token 嵌入维度	$d$
  Attention 层嵌入维度	$d$
  MLP 隐藏层维度	$4d$	通常设置为嵌入维度4倍
  Attention head 数量	$n$	要求其整除 $d$
  词表尺寸	$V$
  batch_size	$b$
  模型输入长度	$s$
  训练 batch data	$x$	张量尺寸 ${\mathbb{R}}^{b\times s\times d}$
  模型参数量	$N$
  训练数据量（Token）	$D$
  训练计算量（FLOPs）	$C$

• 本文主要考察 GPT2/GPT3 等标准 Decoder-only 模型
  

1. LLM 的参数量分析

1.1 Embedding & LM Head

• Token Embedding 嵌入层本质是一个不带偏置的线性层，总参数量为 $Vd$
• 用于自回归预测的多分类头线性层参数通常和嵌入层共享，不占用额外参数

1.2 Transformer Decoder Block

• 一个 Transformer Decoder Block 的总参数量最多 $12d^2+13d$

  • Self-attention 层总参数量 $4d^2(+4d)$：

    1. 投影矩阵 $W_Q,W_K,W_V,W_O$，尺寸均为 ${\mathbb{R}}^{d\times d}$，参数量共 $4d^2$

       参数量和注意力头数量无关，$n>1$ 的情况下每个注意力头有自己的投影矩阵 $W_Q^{(i)},W_K^{(i)},W_V^{(i)}$，尺寸均为 ${\mathbb{R}}^{d\times \frac{d}{n}}$，参数量共 $3\times n\times d\times \frac{d}{n}=3d^2$，再加上汇总各个注意力头信息的 $W_O\in {\mathbb{R}}^{d\times d}$，一共$4d^2$

    2. 如果每个矩阵带偏置，则参数量再加 $4d$

  • MLP 块总参数量 $8d^2+5d$：

    1. 从输入层到隐藏层的线性层尺寸 ${\mathbb{R}}^{d\times 4d}$，带偏置，参数量 $4d^2+4d$
    2. 从隐藏层到输出层的线性层尺寸 ${\mathbb{R}}^{4d\times d}$，带偏置，参数量 $4d^2+d$

  • LayerNorm 层共两个，总参数量 $2\times 2d=4d$

    1. 缩放参数 $\gamma$，尺寸 ${\mathbb{R}}^{1\times d}$，参数量 $d$
    2. 平移参数 $\beta$，尺寸 ${\mathbb{R}}^{1\times d}$，参数量 $d$

1.3 总结

• Decoder-Only 模型各模块参数量如下

  模块	数量	单个参数量	总参数量
  Token Embedding	1	$Vd$	$Vd$
  LM Head	1	-	-
  Self-attention	$l$	$4d^2(+4d)$	$4l{d}^2(+4ld)$
  MLP FFN	$l$	$8d^2+5d$	$8l{d}^2+5ld$
  LayerNorm	$2l$	$2d$	$4ld$

  总参量 $12l{d}^2+13ld+Vd$，通常不考虑嵌入层参数 $Vd$，再忽略一次项后模型总参数量近似为 $12l{d}^2$

• 考察一些常见模型的参数量估算结果

  模型	实际参数量	Transofmer Block 层数 $l$	隐层维度 $d$	$12l{d}^2$
  DB1 (TransformerXL)	1.2B	24	2048	$1,207,959,552\approx 1.20\text{B}$
  Gato (TransformerXL)	1.18B	24	2048	$1,207,959,552\approx 1.20\text{B}$
  Gato (TransformerXL)	364M	12	1536	$339,738,624\approx 340\text{M}$
  Gato (TransformerXL)	79M	8	768	$56,623,104\approx 57\text{M}$
  LLaMa	6.7B	32	4096	$6,442,450,944\approx 6.4\text{B}$
  LLaMa	13.0B	40	5120	$12,582,912,000\approx 12.6\text{B}$
  LLaMa	32.5B	60	6656	$31,897,681,920\approx 31.9\text{B}$

2 计算量估计（FLOPs）

• LLM 中的主要运算是矩阵乘法，故考察 LLM 计算量时，通常只关注矩阵乘法运算对应的浮点计算量

  • 设有矩阵 $A\in {\mathbb{R}}^{m\times n},B\in {\mathbb{R}}^{n\times p}$，矩阵乘法 $AB$ 可以看作 $m\times p$ 次 尺寸为 $n$ 的向量内积运算
    A m × n B n × p = [ a 1 ∗ a 2 ∗ ⋮ a m ∗ ] [ b ∗ 1 b ∗ 2 ⋯ b ∗ p ] \pmb{A}_{m\times n}\pmb{B}_{n\times p}= \begin{bmatrix} a_{1*} \\ a_{2*} \\ \vdots \\ a_{m*} \end{bmatrix} \begin{bmatrix} b_{*1}& b_{*2}& \cdots & b_{*p} \\ \end{bmatrix} Am×n​Bn×p​= ​a1∗​a2∗​⋮am∗​​ ​[b∗1​​b∗2​​⋯​b∗p​​] 向量内积运算中要进行 $n$ 次乘法和 $n-1$ 次加法运算，因此矩阵乘法的浮点运算量为 $2mnp-mp$
  • 但是 GPU 计算矩阵乘法时一般使用 FMA (fused multiply–add) 进行计算，一次 FMA 运算操作是 a <-- a + (b * c)，包括了一次乘法操作和一次加法操作，相当于一次 MACs(Multiply–Accumulate Operations)，而 1MACs = 2FLOPs。故通常把以上矩阵乘法运算的浮点计算量表述为 $2mnp$，以下分析中均基于此进行分析

• 本节分析一次训练迭代过程的计算量，设输入 batch 数据尺寸 $x\in {\mathbb{R}}^{b\times s\times d}$

2.1 Embedding & LM Head

• Token Embedding 可以直接查 lookup table 不进行矩阵运算；
• 预测多分类头将尺寸为 $d$ 的隐藏向量映射为词表大小，矩阵运算尺寸为 $[b\times s\times d]\times [d\times V]\to [b\times s\times V]$，总计算量 $2bsdV$。

2.2 Transformer Decoder Block

• 一个 Transformer Decoder Block 的总计算量最多 $24bs{d}^2+4b{s}^{2}d$

  • Self-attention 层计算过程如下
    $$\begin{array}{l}Q=x{W}_Q,K=x{W}_K,V=x{W}_V\\ \\ A=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{h}}\right)\\ \\ x_{out}=A\cdot V\cdot W_O+x\end{array}$$

    1. 计算 $Q,K,V$：矩阵运算 $xW$ 尺寸为 $[b\times s\times d]\times [d\times d]$，总计算量 $3\times 2bs{d}^2=6bs{d}^2$
    2. 计算注意力得分：矩阵运算 $Q{K}^T$ 尺寸为 $[b\times n\times s\times \frac{d}{n}]\times [b\times n\times \frac{d}{n}\times s]\to [b\times n\times s\times s]$，总计算量 $bn\times 2s^2\frac{d}{n}=2b{s}^{2}d$
    3. 汇聚价值信息：矩阵运算 $AV$ 尺寸为 $[b\times n\times s\times s]\times [b\times n\times s\times \frac{d}{n}]\to [b\times n\times s\times \frac{d}{n}]$，总计算量 $bn\times 2s^2\frac{d}{n}=2b{s}^{2}d$
    4. 输出投影：矩阵运算 $V_s{W}_O$ 尺寸为 $[b,s,d]\times [d,d]\to [b,s,d]$，总计算量 $2bs{d}^2$

  • MLP 块计算过程如下
    $$x=f_{\text{gelu }}\left(x_{\text{out }}{W}_1\right)W_2+x_{\text{out }}$$

    1. 第一个线性层：矩阵运算 $x_{\text{out }}{W}_1$ 尺寸为 $[b\times s\times d]\times [d\times 4d]\to [d\times s\times 4d]$，总计算量 $8bs{d}^2$
    2. 第二个线性层：矩阵运算 $(x_{\text{out }}{W}_1)W_2$ 尺寸为 $[b\times s\times 4d]\times [4d\times d]\to [d\times s\times d]$，总计算量 $8bs{d}^2$

2.3 反向传播的计算量

• 以上分析了正向传播过程中的计算量，本节进一步说明反向传播计算量约为正向传播的两倍。考虑简单的 MLP 模型，不考虑偏置，设激活函数为 $\sigma$，输入为 $x^1$，前向传播和对应的反向传播计算过程如下
  $$\begin{array}{rlrlrl}& 1& z^1=W^1{x}^1& & & \frac{\partial L}{\partial W^1}=\frac{\partial L}{\partial z^1}\frac{\partial z^1}{\partial W^1}={\left[(W^2{)}^T\left(\frac{\partial L}{\partial a^2}\odot {\sigma }^{\prime }(z^2)\right)\odot {\sigma }^{\prime }(z^1)\right]}^T{x}^1\\ & 2& a^1=\sigma (z^1)& & & \frac{\partial L}{\partial z^1}=\frac{\partial L}{\partial a^1}\frac{\partial a^1}{\partial z^1}=(W^2{)}^T\left(\frac{\partial L}{\partial a^2}\odot {\sigma }^{\prime }(z^2)\right)\odot {\sigma }^{\prime }(z^1)\\ & & & & & \\ & & & & & \frac{\partial L}{\partial a^1}=\frac{\partial L}{\partial z^2}\frac{\partial z^2}{\partial a^1}=(W^2{)}^T\left(\frac{\partial L}{\partial a^2}\odot {\sigma }^{\prime }(z^2)\right)\\ & 3& z^2=W^2{a}^1& & & \frac{\partial L}{\partial W^2}=\frac{\partial L}{\partial z^2}\frac{\partial z^2}{\partial W^2}={\left(\frac{\partial L}{\partial a^2}\odot {\sigma }^{\prime }(z^2)\right)}^T{a}^1\\ & 4& a^2=\sigma (z^2)& & & \frac{\partial L}{\partial z^2}=\frac{\partial L}{\partial a^2}\frac{\partial a^2}{\partial z^2}=\frac{\partial L}{\partial a^2}\odot {\sigma }^{\prime }(z^2)\\ & & & & & \\ & 5& L=loss(a^2,y)& & & \frac{\partial L}{\partial a^2}\end{array}$$ 其中 $\odot$ 是矩阵对应位置元素相乘
• 关注反向传播过程中的矩阵乘法操作，注意到反向传播过程中每个非第一层都有两次矩阵乘法操作（计算 $\frac{\partial L}{\partial W^i},\frac{\partial L}{\partial a^i}$），而相应的前向过程中只有一次（计算$z^i=W^i{a}^i$）。随着网络层数增加，反向传播计算量会越来越接近正向传播计算量的两倍。进一步的详细说明请参考：浅谈后向传递的计算量大约是前向传递的两倍

2.4 总结

• Decoder-Only 模型各模块前向传播的矩阵乘法计算量如下

  模块	数量	单个计算量	总计算量
  Token Embedding	1	-	-
  LM Head	1	$2bsdV$	$2bsdV$
  Self-attention	$l$	$8bs{d}^2+4b{s}^{2}d$	$8bsl{d}^2+4b{s}^{2}ld$
  MLP FFN	$l$	$16bs{d}^2$	$16bsl{d}^2$

  一次前向传播的总计算量为 $24bsl{d}^2+4b{s}^{2}ld+2bsdV$。通常不考虑嵌入层计算量 $2bsdV$，当 $d>>s$ 时计算量近似为 $24bsl{d}^2$

• 进一步考虑反向传播的计算量，一次训练迭代的总计算量（FLOPs）近似为 $C\approx 72bsl{d}^2$，注意到输入 Token 数据量 $D=bs$，模型参数量由第一节分析为 $N\approx 12l{d}^2$，三者间有以下重要关系
  $$C\approx 6ND$$ 这可以理解为：一次训练迭代中，对于每个token，每个模型参数，需要进行约 6 次浮点数运算

  可以用此结论正确估计 GPT3-175B 的计算量
  
  如图所示，GPT3-175B 参数量 174600M，训练数据量为 300B Token，总计算量为 3.14E+23，利用以上关系估计为
  $$6\times 174600\times 10^6\times 300\times 10^9=3.143\times 10^{23}\text{ FLOPs}$$

2.5 补充

• 有了计算量，就可以结合硬件平台评估自己的训练速度是否正常。A100显卡不做任何优化的情况下，tflop/s 应该达到90的底线，正常约为120，用 flash attention 2 优化的可达 200+。另外也可以预估训练时间
  $$\text{ 训练时间 }\approx \frac{6ND}{\text{GPU数 }\times \text{GPU峰值FLPOS }\times \text{GPU利用率}}$$ 数据加载、优化器更新、多卡通信和记录日志得都会影响 GPU 利用率，一般在 $[0.3,0.55]$ 之间。
• 注意到前向过程计算的所有激活值 $a$ 通常都会保存下来用于反向传播的梯度计算（如2.3节中的第3步），这会导致大量显存占用，对于 LLM 训练不利。一种改进方案是使用激活重计算技术来减少中间激活显存，它在反向传播过程中进行一次额外的前向传递来临时计算激活值，这时有
  $$C\approx 8ND$$
• 进一步考虑 Encoder-Decoder 模型的计算量。由于这类 LLM 中输入只会过encoder，输出只会过decoder，假设二者分别占一半的参数量和数据量，则参数量相同的情况下，其整体计算量约为
  $$6\cdot \frac{N}{2}\cdot \frac{D}{2}+6\cdot \frac{N}{2}\cdot \frac{D}{2}=3ND$$ 也就是说，参数量相同的情况下，Encoder-Decoder 架构的模型计算量约为 Decoder-only 模型的一半。大部分在两者间进行对比的文章会使 Encoder-Decoder 类模型具有两倍的参数，使其推理成本接近