1. 前言
本文分析采用decoder-only框架transformer模型的模型参数量、计算量、中间激活值、KV cache。 为了方便分析,先定义好一些数学符号。记transformer模型的层数为$l$, 隐藏层维度为$h$,注意力头数为$a$。词表大小为$V$,训练数据的批次大小为$b$,序列长度为$s$
2. 模型参数量
transformer模型由$l$个相同的层组成,每个层分为两部分:self-attention块和MLP块。 self-attention块的模型参数有$Q$、$K$、$V$的权重矩阵$W_Q$、$W_K$、$W_V$和偏置,输出权重矩阵$W_O$和偏置。4个权重矩阵的形状为$[h, h]$, 4个偏置的形状为$[h]$。self- attention块的参数量为$4h^2+4h$。
MLP块由2个线性层组成,一般地,第一个线性层是先将维度从$h$映射到 $4h$,第二个线性层再将维度从$4h$映射到$h$。第一个线性层的权重矩阵$W_1$的形状为 $[h,4h]$ ,偏置的形状为 $[4h]$ 。第二个线性层权重矩阵 $W_2$ 的形状为 $[4h,h]$ ,偏置形状为 $[h]$ 。MLP块的参数量为 $8h^2+5h$ 。
self-attention块和MLP块各有一个layer normalization,包含了2个可训练模型参数:缩放参数 $\gamma$ 和平移参数 $\beta$ ,形状都是 $[h]$ 。2个layer normalization的参数量为 $4h$ 。
Refence
分析transformer模型的参数量、计算量、中间激活、KV cache 理解 Linux 动态链接库依赖