02-06.模型参数计算.md

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

采用decoder-only框架transformer模型的模型参数量、计算量、中间激活值、KV cache。

为了方便分析，先定义好一些数学符号。记transformer模型的层数为 $l$ ，隐藏层维度为 $h$ ，注意力头数为 $a$ 。词表大小为 $V$ ，训练数据的批次大小为 $b$ ，序列长度为 $s$ 。

transformer模型由 $l$ 个相同的层组成，每个层分为两部分：self-attention块和MLP块。

参数量

self-attention块的参数量

包含$W_Q$,$W_K$,$W_V$三个权重矩阵和偏置，每个权重矩阵形状是$[h,h]$,偏置形状是$[h,1]$

Self-attention参数量 $4h^2+4h$

MLP块

包含两个线性层。

一般地，

第一个线性层是先将维度从 $h$ 映射到 $4h$ ，第一个线性层的权重矩阵 $W1$ 的形状为 $[h,4h]$ ，偏置的形状为 $[4h]$

第二个线性层再将维度从 $4h$映射到 $h$。 权重矩阵 $W2$ 的形状为 $[4h,h]$ ，偏置形状为 $[h] $。

MLP块的参数量为 $8h^2+5h $

layer normalization

self-attention块和MLP块各包含了2个可训练模型参数：缩放参数 γ 和平移参数 β ，形状都是 $[h]$ 。2个layer normalization的参数量为 $4h$ 。

总的，每个transformer层的参数量为 $12h^2+13h$

词嵌入

词向量通常等于隐藏维度$h$,词嵌入矩阵的参数量为$Vh$

总结

l层transformer的可训练模型参数量为$l(12h^2+13h)+Vh$

当h较大时，近似$12lh^2$

不同版本LLaMA模型

实际参数量	隐藏维度h	层数l	近似值 12lh^2
6.7B	4096	32	6,442,450,944
13.0B	5120	40	12,582,912,000
32.5B	6656	60	31,897,681,920
65.2B	8192	80	64,424,509,440

显存占用分析

训练显存占用

在训练神经网络的过程中，占用显存的大头主要分为四部分： 模型参数、前向计算过程中产生的中间激活、后向传递计算得到的梯度、优化器状态 。

在一次训练迭代中，每个可训练模型参数都会对应1个梯度，并对应2个优化器状态（Adam优化器梯度的一阶动量和二阶动量）

设模型参数量为 $\phi$，那么梯度的元素数量为 $\phi$ ，AdamW优化器的元素数量为 $2\phi$ 。

float16数据类型的元素占2个bytes，float32数据类型的元素占4个bytes。

在混合精度训练中，会使用float16的模型参数进行前向传递和后向传递，计算得到float16的梯度；

在优化器更新模型参数时，会使用float32的优化器状态、float32的梯度、float32的模型参数来更新模型参数。

$$ 2+4+2+4+4+4 = 20 bytes

$$

推理显存

推理阶段模型参数占用的显存大概是 $2\phi bytes$