transformer模型的参数量和计算量

参考https://zhuanlan.zhihu.com/p/624740065，本文为个人学习记录。

前言

大规模语言模型（Large Language Model, LLM）的大体现在两个方面：模型参数规模大，训练数据规模大。以GPT3为例，参数量为1750亿，训练数据达到了570GB。进而，训练大语言模型面临两个主要挑战：显存效率和计算效率。

现在业界的大语言模型都是基于transformer模型的，模型结构主要是encoder-decoder（代表模型是T5）和decoder-only（分为Causal LM——代表模型GPT，Prefix LM——代表模型GLM）。本文分析的是decoder-only框架transformer模型。

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

模型参数量

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$。

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

此外，词嵌入矩阵的参数量也较多，词向量维度通常等于隐藏层维度$h$，词嵌入矩阵的参数量为 V h Vh