本文 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 = 14 GB\) 的显存

训练

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

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

一个fp16的数据占用2个byte，fp32占用4个byte。因此，对于参数量为 \(\Phi\) 的模型来说，共需要 \((2+2+4\times4) \Phi = 20\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 * n * p * 2\)

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

self attention

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

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

• $Q = xW_Q, K = xW_K, V = xW_V$

  • eg: \(Q = xW_Q\), \([b, L, h][h,h] = [b, L, h]\), 计算量 \(2bLhh = 2bLh^2\)

具体计算量：

• 计算 \(Q、K、V\): 三次 \(x^{b \times L \times h} \cdot W_Q^{h \times h}\)，计算量为 \(3 * 2bLh^2 = 6bL h^2\)
• 计算 \(QK^T\): \(Q^{b \times n_{head} \times L \times h_{head}} \cdot {K^T}^{b \times n_{head} \times h_{head} \times L}\)，结果是 \(bn_{head}LL\)，计算量为 \(2bn_{head}LLh_{head} = 2bL^2h\) \((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}}\)，结果是 \(bn_{head}Lh_{head}\) = \(bLh\)，计算量为 \(2bn_{head}Lh_{head}L = 2bL^2h\)
• 结果线性映射 \(O \cdot W_0\): \(O^{b \times L \times h} \cdot W_0^{h \times h}\)， 结果是 \(bLh\)，计算量为 \(2bLhh = 2bLh^2\)

总计，self attention 层的计算量共为，\(8bLh^2 + 4bL^2h\)

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 = 8bLh^2\)
• 计算 \(x_{mid}^{b \times L \times 4h} \cdot W_2^{4h \times h}\)，结果是 \(bLh\)，计算量为 \(2bLh4h = 8bLh^2\)

总计，MLP 层的计算量共为，\(16bLh^2\)

Embedding

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

综上:

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

参数量和计算量的关系

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

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

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

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

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

训练时间预估

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

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