LLM推理与训练显存的计算方法
值得一看的链接：
分析transformer模型的参数量、计算量、中间激活、KV cache
需要多少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 于苏州