值得一看的链接：

分析transformer模型的参数量、计算量、中间激活、KV cache

大模型GPU显存占用计算_大模型显存计算-CSDN博客

需要多少GPU显存才能运行预训练大语言模型？大语言模型参数规模与显存大小的关系估算方法

众所周知，自从Transformer出现后，大语言模型基本上抛弃了原来的RNN/LSTM的思路。然而Transformer类模型也有自己的不足之处，其中最被人诟病的可能就是大量的计算量。今天我就来学习一下，训练/运行一个Transformer类的模型，究竟要占用多少显存？

整数与浮点数

在开始之前，需要了解一下不同的数据格式占用的存储大小：

• FP64：双精度浮点数（64位），8字节/单位
• FP32：单精度浮点数（32位），4字节/单位
• FP16：半精度浮点数（16位），2字节/单位
• BF16：Google推出的半精度浮点数（16位）
• INT16：整数（16位），2字节/单位
• INT8：整数（8位），1字节/单位
• INT4：整数（4位），0.5字节/单位

浮点数是如何变成整数的？下图是一个简单的演示。在这过程中，会出现精度和存储空间的取舍。

模型占用多大空间？

什么是B？B代表Billion，十亿参数。因此13B就是130亿参数，7B则是70亿参数。

我们以参数量13B大模型为例，13B是130亿参数。如果我们以全精度fp32运行，也就是float32，将会占用32位bit，也就是4byte字节。我们给出计算的公式：

$1\text{GB} = 1024 \text{MB} = 1024^{2} \text{KB} = 1024^{3} \text{Byte}$

那么全精度13B模型占用显存为 数量 * 类型大小：$13 \times 10^{9} \times 4 \text{Byte} \div 1024^{3} \approx 48.4 \text{GB}$

也就是说，我们只是单纯以全精度加载运行一个13B的模型，就需要大约48.4G的显存空间。

实际计算中，通常会把$1000^{3}$个字节估计为1G，而非$1024^{3}$个字节。以此类推，常见的13B模型，需要的显存约为：

• FP32 13*4 = 52 GB
• BF16 13*2 = 26 GB
• INT8 13*1 = 13 GB
• INT4 13*0.5 = 6.5 GB

模型本身的参数量是怎么计算的？

首先有一些符号需要给出定义：

• transformer模型的层数为 𝑙
• 隐藏层维度为 ℎ
• 注意力头数为 𝑎
• 词表大小为 𝑉
• 训练数据的批次大小为 𝑏
• 序列长度为 𝑠 。

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

self-attention块的模型参数有 𝑄、𝐾、𝑉 的权重矩阵：$W_{Q}$、$W_{K}$、$W_{V}$ 和偏置，输出权重矩阵 $W_{O}$ 和偏置，4个权重矩阵的形状为 [ℎ,ℎ] ，4个偏置的形状为 [ℎ] 。self- attention块的参数量为 $4h^2 + 4h$ 。

MLP块由2个线性层组成，第一个线性层一般将维度从 ℎ 映射到 4ℎ ，第二个线性层再将维度从4ℎ映射到ℎ。第一个线性层的权重矩阵 $W_{1}$ 的形状为 [ℎ,4ℎ] ，偏置的形状为 [4ℎ] 。第二个线性层权重矩阵 $W_{2}$ 的形状为 [4ℎ,ℎ] ，偏置形状为 [ℎ] 。MLP块的参数量为 $8h^2 + 5h$ 。

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

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

除此之外，词嵌入矩阵的参数量也较多，词向量维度通常等于隐藏层维度 ℎ ，词嵌入矩阵的参数量为 𝑉ℎ 。最后的输出层的权重矩阵通常与词嵌入矩阵是参数共享的。

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

综上， 𝑙 层transformer模型的可训练模型参数量为$l(12h^2 + 13h) + Vh$。当隐藏维度 ℎ 较大时，可以忽略一次项，模型参数量近似为 $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

模型训练所需的存储空间

当模型训练时，需要的存储空间就不止参数量本身了，因为反向传播当中处处需要计算和存储。我们只要记住：训练参数是模型参数的4倍就可以了。

在训练中，存储空间主要分为了三个部分：

• 梯度：1$\times$参数量
• 参数：1$\times$参数量
• 优化器部分：Adam-2$\times$参数量 /SGD-1$\times$参数量

对于每个参数，都会有一个对应的梯度，因此梯度占据的内存和参数一样。

在训练过程中，优化器则会为存储额外的信息，例如一阶和二阶动量。以Adam为例，它会存储2倍参数量（即每个参数都有两个动量）的信息。而对于SGD则只需要1倍的参数量。

LoRA微调参数量

此前了解过LoRA的微调方法，是利用低秩矩阵来训练一个Adaptor，对原有模型的参数进行调整。LoRA训练的参数量取决于秩Rank的大小。

假如全量微调时需要$1024\times 512$的模型参数量，使用LoRA时，就只需要$1024\times8\+512\times8+$原始模型参数量的总合参数量，这大大降低了的训练所需的显存占用。

2024/3/31 于苏州