本文 Decoder Only 模型的参数量和计算量，主要包括 Transformer结构、参数量方面、显存占用、计算量方面、参数量和计算量的关系、训练时间预估等等。

https://zhuanlan.zhihu.com/p/830885346

Decoder Only 模型的参数量和计算量

Transformer

参数量方面

• 大模型 = Embedding + 若干层 Transformer Layer
• 每层 Transformer Layer 主要由 self attention 层 和 MLP 层构成

self attention

计算 self attention 层的参数：

• $W_Q$、偏置：$h\times h,1\times h$
• $W_K$、偏置：$h\times h,1\times h$
• $W_V$、偏置：$h\times h,1\times h$
• $W_0$、偏置：$h\times h,1\times h$

总计： self attention 层参数共，$4h^2+4h$

MLP

计算 MLP 层的参数，以 GPT-3 为例子

• h – 4h – h
• 上投影 $W_1$、偏置：$h\times 4h,1\times 4h$

• 下投影，$W_2$、偏置：$4h\times h,1\times h$

• 总计：MLP 层参数共，$8h^2+5h$

LayerNorm

• 在 self attention 层和 MLP 层后均有一个 LayerNorm 层，每个LN层均包含两个可训练的参数，因此LN层共有参数 $2h\times h=4h$

Embedding

• Embedding 层包含的参数： $Vh$
• V 词库token大小

综上：

• 一层 transformer decoder layer 共包含参数：$12h^2+13h$
• 一个包含l层 decoder layer 的生成模型，共包含参数：$l(12h^2+13h)+Vh$

eg: QWen-1.8B：l=24，h=2048，V=151851，按照上面的公式计算得到，参数量大约为1.92B

显存占用

推理

• 推理时，仅参数占用显存，每个fp16精度的参数占用2个byte，fp32精度的参数占用4个byte。
• 因此一个7B的fp16精度模型，大概需要占用 $7\times 2=14GB$ 的显存

训练

以fp16方式训练时，训练过程中需要保存的数据有：

• 模型参数：fp16
• 参数梯度：fp16
• 优化器状态：fp32一阶动量、fp32二阶动量、fp32模型参数、fp32参数梯度

一个fp16的数据占用2个byte，fp32占用4个byte。因此，对于参数量为 $\mathrm{\Phi }$ 的模型来说，共需要 $(2+2+4\times 4)\mathrm{\Phi }=20\mathrm{\Phi }$ 的显存空间。

eg:

一个7B的模型，大约需要 $20\times 7\times {10}^9/{1024}^3=130GB$ 显存空间

计算量方面

前置：对于矩阵 $A\in {\mathbb{R}}^{m\times p},B\in \mathbb{R}p\times n$，AB相乘的计算量：$m\ast n\ast p\ast 2$

• 其中，$m\ast n$ 表示结果矩阵包含 $m\ast n$ 个元素
• $p\ast 2$ 表示每个元素需要经过 p 次加法和 p 次乘法计算得到

self attention

self attention 的计算量：假设输入数据的形式为 $[b,L]$ (batch, sequence length)

• $x=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{h}}\right)\cdot V\cdot W_0+x$

• $Q=x{W}_Q,K=x{W}_K,V=x{W}_V$

  • eg: $Q=x{W}_Q$, $[b,L,h][h,h]=[b,L,h]$, 计算量 $2bLhh=2bL{h}^2$

具体计算量：

• 计算 $Q、K、V$: 三次 $x^{b\times L\times h}\cdot W_Q^{h\times h}$，计算量为 $3\ast 2bL{h}^2=6bL{h}^2$
• 计算 $Q{K}^T$: $Q^{b\times n_{head}\times L\times h_{head}}\cdot {K^T}^{b\times n_{head}\times h_{head}\times L}$，结果是 $b{n}_{head}LL$，计算量为 $2b{n}_{head}LL{h}_{head}=2b{L}^{2}h$ $(n_{head}{h}_{head}=h)$
• 计算 $S\cdot V$: $S^{b\times n_{head}\times L\times L}\cdot V^{b\times n_{head}\times L\times h_{head}}$，结果是 $b{n}_{head}L{h}_{head}$ = $bLh$，计算量为 $2b{n}_{head}L{h}_{head}L=2b{L}^{2}h$
• 结果线性映射 $O\cdot W_0$: $O^{b\times L\times h}\cdot W_0^{h\times h}$， 结果是 $bLh$，计算量为 $2bLhh=2bL{h}^2$

总计，self attention 层的计算量共为，$8bL{h}^2+4b{L}^{2}h$

MLP

MLP层的计算量可以表示为：$x=f\left(x_{out}{W}_1\right)W_2+x_{out}$

• 计算 $x_{mid}=x_{out}^{b\times L\times h}\cdot W_1^{h\times 4h}$，结果是 $bL4h$，计算量为 $2bL4hh=8bL{h}^2$
• 计算 $x_{mid}^{b\times L\times 4h}\cdot W_2^{4h\times h}$，结果是 $bLh$，计算量为 $2bLh4h=8bL{h}^2$

总计，MLP 层的计算量共为，$16bL{h}^2$

Embedding

隐藏向量映射到词表，运算量为 2bLhV

综上:

• 一层 transformer decoder 进行一次前向计算的计算量：$24bL{h}^2+4b{L}^{2}h$
• 一个包含l层decoder layer的生成模型，对于其中一个batch，进行一次前向计算的计算量是 $l($24bLh^2 + 4bL^2h$)+2bLhV$

参数量和计算量的关系

• 在前向传播中，输出 token 总数为 bL，模型总参数量为 $12h^2+13h$

• 可以计算出，每个 token，每个参数需要的浮点数计算次数为

• $$\frac{24bL{h}^2+4b{L}^{2}h}{bL\cdot (12h^2+13h)}=\frac{24h+4L}{12h+13}\approx 2$$

即，前向传播一次，每token、每个参数需要进行 2次浮点数运算 FLOPS，反向传播所需计算量是前向传播的2倍，因此前向+反响传播，每token、每个参数需要进行6次浮点数计算 FLOPS

以 GPT3-175B为例子，其参数 174600M，训练数据为300B，则训练所需要的总计算量为 $6\ast 174600\ast {10}^6\ast 300\ast {10}^9\approx 3.143\times {10}^{23}flops$

训练时间预估

在实际训练中，为了节省中间激活的显存占用，通常会在反向传播时进行一次重新计算，因此会引入一次额外的前向传播，此时，一次完整的前向反向传播，每token每参数共需要进行8次浮点数运算。训练时间估计可以参考如下公式：

以GPT3-175B为例子，需要的训练时间为