文章内容主要摘自：https://zhuanlan.zhihu.com/p/624740065

1. 前言

最近，OpenAI推出的ChatGPT展现出了卓越的性能，引发了大语言模型(Large Language Model, LLM)的研究热潮。大规模语言模型的“大”体现在两个方面：模型参数规模大，训练数据规模大。以GPT3为例，GPT3的参数量为1750亿，训练数据量达到了570GB。进而，训练大规模语言模型面临两个主要挑战：显存效率和计算效率。

现在业界的大语言模型都是基于transformer模型的，模型结构主要有三大类：

• encoder-decoder（代表模型是T5）、
• encoder-only（代表模型是BERT，也叫Masked-Language Model，MLM）以及
• decoder-only，具体的，decoder-only结构又可以分为Causal LM（因果语言模型-自回归的预测，代表模型是GPT系列）和Prefix LM（代表模型是GLM，UniLM）。归因于GPT系列取得的巨大成功，大多数的主流大语言模型都采用Causal LM结构。
• 

因此，针对decoder-only框架，为了更好地理解训练训练大语言模型的显存效率和计算效率，本文分析采用decoder-only框架来分析transformer模型的模型参数量、计算量、中间激活值、KV cache。

Transformer 原始架构和decoder-only架构的GPT对比：

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

2. 模型参数量

decoder-only架构的大模型有两部分组成，输入&输出的词嵌入层（大部分模型词嵌入层共享参数）以及堆叠$l$层的transformer块。

transformer块（上图2所示）由两部分组成：Attention和FFN（MLP），其中穿插了两层LN（Pre-Norm）。

• Attention模块模型参数包括将输入投影为$Q,K,V$的权重矩阵 $W_Q,W_K,W_V$ 和偏置，输出矩阵为 $W_O$ 和偏置，四个权重矩阵的size为 $h\times h$，四个偏置的size为 $h$。忽略激活函数，此时Attention模块的参数量为 $4h^2+4h$，参数量来源于线性层投影矩阵。
• FFN模块，由两个线性层组成，第一个线性层: $h \rightarrow 4h $ 权重矩阵和偏置分别为 $h\times 4h,4h$，第二个线性层：$4h\to h$ 的权重矩阵和参数量为 $4h\times h,h$。忽略激活函数，此时FFN模块的参数量为 $8h^2+5h$。
• 两个LN层，包含了2个可训练模型参数：缩放参数 $\gamma$和平移参数 $\beta$ ，形状都是 $h$ 。2个layer normalization的参数量为 $4h$ 。

此外，还包括一个输入输出的词嵌入层，词嵌入层的参数量取决于输入字典的大小 $V$, 参数量为 $Vh$

关于位置编码，如果采用可训练式的位置编码，会有一些可训练模型参数，数量比较少。如果采用相对位置编码，例如RoPE和ALiBi，则不包含可训练的模型参数。我们忽略这部分参数。

综上Transformer decoder模型的**参数量为$12l{h}^2+13lh+Vh$**。当隐藏维度 $h$ 较大时，可以忽略一次项，模型参数量近似为 $12lh2$ 。