Dominic-20240602

[DOMITIGA]

大模型训练显存分析

当前，大语言模型（Large Language Models，LLMs）普遍采用transformer模型结构，主要有两类：encoder-decoder结构和decoder-only结构。Decoder-only又分为Causal LM（e.g. GPT系列）和Prefix LM（e.g. GLM）。GPT系列取得的巨大成功使得目前主流LLMs都采用Causal LM结构。 本章以decoder-only结构为例，介绍大模型的参数量、计算量和显存占用。

Decoder-only结构和参数

本章以一个层数为l的decoder-only transformer结构为例，模型结构如图所示。首先，将input tokens输入嵌入层（embedding layer），包括词嵌入（word embedding ）和位置嵌入（position embedding）。词嵌入表（word embedding table）是一个大小为v×h的矩阵，其中v是指词表大小（vocabulary size），h被称为隐藏层大小（hidden dimension size）。位置嵌入是一个大小为s×h的矩阵，其中s是指句长（sequence length）。嵌入层的输出是一个大小为b×s×h的三阶张量，其中b是指microbatch size，这个三阶张量下一步会进入到transformer block中。每个transformer层包括一个头数为 a 的多头自注意力层（multi-head self-attention block）；和一个层数为2的multi-layer preceptron （MLP）层，MLP将隐藏层大小从h增加到4h，然后再将其减少回h。每个transformer层的输入输出都是大小为s×b×h的三阶张量。最后一个transformer层的输出会被映射回此表维度以计算交叉熵损失。假设词嵌入层和输出层使用同一份权重。

大模型参数量

大模型的参数量由其transformer结构的各个层的参数加和得到，具体包括：嵌入层、多头自注意力层、全连接层等。下面分别介绍每个层的参数量。

（1）嵌入层：

嵌入层包括词嵌入和位置嵌入。在transformer结构中，位置嵌入采用正弦函数、余弦函数公式计算位置编码的方法，这一过程不引入参数。

在词嵌入过程中，用于大模型训练的语料经过tokenizer后，输出的tokens常以b×s×v one-hot矩阵的形式输入词嵌入层。One-hot特征编码形式中， 每一个token都是独立的，有多少个token就会有多少个特征例，维度非常庞大，而且表达过于稀疏，这么大的维度在后续运算过程中会造成很大的一个计算量。所以为了方便后续计算，需要对输入矩阵做一个嵌入降维。

词嵌入层旨在将单词one-hot形式的特征表达降维转换成固定维度（h）的向量的形式，即b×s×h。词嵌入过程本质上是一个全连接神经网络层，输入为b×s×v的one hot矩阵、中间层节点数为字向量维数h、输出矩阵为b×s×h。在大模型训练过程中，训练的目的是通过训练数据迭代学习隐藏层的特征参数，即隐藏层的权重矩阵，大小为v×h。因此嵌入层的参数量大小是v×h。

值得注意的是，从运算上来看，one hot型的矩阵相乘，就像是相当于查表，于是把one hot型的矩阵运算简化为了查表操作，而不写成矩阵再运算，这大大降低了运算量。另一方面，词嵌入过程得到了这个全连接层的参数之后，直接用这个全连接层的参数作为特征，用这个全连接层的参数作为字、词的表示，从而得到了字、词向量，词向量的夹角余弦能够在某种程度上表示字、词的相似度。从特征角度来看，词嵌入过程有两个价值，第一个是固定维度，让原来 one-hot 编码产生特征爆炸的，或者维度不固定的把它固定下来。固定维度以后为了后续的可计算可以做各种各样的一些计算，例如比较最接近的单词有哪些，做一些加减计算等。

（2）多头自注意力层：

多头自注意力层参数包含Q，K，V 权重矩阵W_Q，W_K，W_V，以及输出W_O，4个权重矩阵形状为[h, h]，因此，因此自注意力层自身的参数量为4h²。需要注意的是参数量和注意力头数无关。如果每个矩阵带偏置，则参数量再加上4h。

（3）全连接层：

全连接层由2个线性层组成，第一个线性层将维度从h映射到4h，权重矩阵形状为[h，4h]；第二个线性层将维度从4h映射回h，权重矩阵形状为[4h，h]。因此，全连接层参数量为8h²。

（4）归一化层

Layer Normalization（层归一化）是一种用于训练深度神经网络的归一化技术，它是由 Jimmy Lei Ba, Jamie Ryan Kiros 和 Geoffrey Hinton 在 2016 年提出的。层归一化的主要目的是在训练过程中稳定和加速神经网络的收敛。在传统的Batch Normalization（批量归一化）中，归一化是在小批量（mini-batch）的样本上进行的，这意味着每个特征的均值和方差是在每个小批量中计算的。然而，批量归一化在处理小批量大小或者在循环神经网络（RNN）中可能不太有效。层归一化与批量归一化的主要区别在于，它不是在小批量上进行归一化，而是在单个样本的单个层的所有激活上进行归一化。具体来说，对于给定的层，层归一化会计算该层所有激活的均值和方差，然后使用这些统计量来归一化每个激活值。这样，层归一化不受小批量大小的影响，并且在RNN中表现更好。

层归一化的公式如下：

其中：

x是输入的激活值。

​是激活值的均值。

​是激活值的标准差。

​和​是可学习的参数，用于在归一化后重新缩放和平移激活值。

层归一化有助于减少内部协变量偏移（Internal Covariate Shift），这是深度学习中常见的问题，即网络每层的输入分布随着训练过程而变化。通过层归一化，可以使得每层的输入分布更加稳定，从而有助于网络更快地收敛，并且提高模型的泛化能力。

自注意力层和全连接层各有归一化层，包含2个可训练参数：缩放参数​和平移参数​，2层归一化层参数量为4h。

（5）大模型参数量和模型大小的关系

大模型总参数量是嵌入层、多头自注意力层、全连接层、归一化层等各层参数量之和。l层transformer结构可训练参数量为l(12h²)+Vh。如果模型使用FP16/BF16精度保存和加载到显存，那么占用2个Byte；如果模型使用FP32精度保存和加载到现存，那么占用4个Byte。

模块 | 参数量 -- | -- 嵌入层 | Vh 多头自注意力层 | 4h2 全连接层 | 8h2 归一化层 | 4h

因此在大模型训练过程中最少需要 8.X 倍于模型权重的大小内存。例如：LLAMA-65B模型权重参数 130 GB，训练时候总消耗 NPU 内存最小为 130GB x 8 ≈ 1TB。1TG/64 GB≈17，除去激活值需要的内存，仅仅放下一个大模型就需要 32 张 NPU 。

在大模型训练过程中，模型权重参数、模型梯度使用半精度保存，一个参数占2个字节；优化器状态包括：模型权重参数、Adam优化器一阶动量m、二阶动量v使用单精度保存，一个参数占4个字节；在前向传播过程中保存激活值，用于反向传播。

大模型训练显存分析

当前，大语言模型（Large Language Models，LLMs）普遍采用transformer模型结构，主要有两类：encoder-decoder结构和decoder-only结构。Decoder-only又分为Causal LM（e.g. GPT系列）和Prefix LM（e.g. GLM）。GPT系列取得的巨大成功使得目前主流LLMs都采用Causal LM结构。 本章以decoder-only结构为例，介绍大模型的参数量、计算量和显存占用。

Decoder-only结构和参数

本章以一个层数为l的decoder-only transformer结构为例，模型结构如图所示。首先，将input tokens输入嵌入层（embedding layer），包括词嵌入（word embedding ）和位置嵌入（position embedding）。词嵌入表（word embedding table）是一个大小为v×h的矩阵，其中v是指词表大小（vocabulary size），h被称为隐藏层大小（hidden dimension size）。位置嵌入是一个大小为s×h的矩阵，其中s是指句长（sequence length）。嵌入层的输出是一个大小为b×s×h的三阶张量，其中b是指microbatch size，这个三阶张量下一步会进入到transformer block中。每个transformer层包括一个头数为 a 的多头自注意力层（multi-head self-attention block）；和一个层数为2的multi-layer preceptron （MLP）层，MLP将隐藏层大小从h增加到4h，然后再将其减少回h。每个transformer层的输入输出都是大小为s×b×h的三阶张量。最后一个transformer层的输出会被映射回此表维度以计算交叉熵损失。假设词嵌入层和输出层使用同一份权重。

大模型参数量

大模型的参数量由其transformer结构的各个层的参数加和得到，具体包括：嵌入层、多头自注意力层、MLP层等。下面分别介绍每个层的参数量。

（1）嵌入层：

嵌入层包括词嵌入和位置嵌入。在transformer结构中，位置嵌入采用正弦函数、余弦函数公式计算位置编码的方法，这一过程不引入参数。

在词嵌入过程中，用于大模型训练的语料经过tokenizer后，输出的tokens常以b×s×v one-hot矩阵的形式输入词嵌入层。One-hot特征编码形式中， 每一个token都是独立的，有多少个token就会有多少个特征例，维度非常庞大，而且表达过于稀疏，这么大的维度在后续运算过程中会造成很大的一个计算量。所以为了方便后续计算，需要对输入矩阵做一个嵌入降维。

词嵌入层旨在将单词one-hot形式的特征表达降维转换成固定维度（h）的向量的形式，即b×s×h。词嵌入过程本质上是一个全连接神经网络层，输入为b×s×v的one hot矩阵、中间层节点数为字向量维数h、输出矩阵为b×s×h。在大模型训练过程中，训练的目的是通过训练数据迭代学习隐藏层的特征参数，即隐藏层的权重矩阵，大小为v×h。因此嵌入层的参数量大小是v×h。

值得注意的是，从运算上来看，one hot型的矩阵相乘，就像是相当于查表，于是把one hot型的矩阵运算简化为了查表操作，而不写成矩阵再运算，这大大降低了运算量。另一方面，词嵌入过程得到了这个全连接层的参数之后，直接用这个全连接层的参数作为特征，用这个全连接层的参数作为字、词的表示，从而得到了字、词向量，词向量的夹角余弦能够在某种程度上表示字、词的相似度。从特征角度来看，词嵌入过程有两个价值，第一个是固定维度，让原来 one-hot 编码产生特征爆炸的，或者维度不固定的把它固定下来。固定维度以后为了后续的可计算可以做各种各样的一些计算，例如比较最接近的单词有哪些，做一些加减计算等。

（2）多头自注意力层：