[NLP]LLM 训练时GPU显存耗用量估计-Toy模板网

这篇具有很好参考价值的文章主要介绍了[NLP]LLM 训练时GPU显存耗用量估计。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

一只进行推理

全精度llama2 7B最低显存要求：28GB
全精度llama2 13B最低显存要求：52GB
全精度llama2 70B最低显存要求：280GB

16精度llama2 7B预测最低显存要求：14GB
16精度llama2 13B预测最低显存要求：26GB
16精度llama2 70B预测最低显存要求：140GB

8精度llama2 7B预测最低显存要求：7GB
8精度llama2 13B预测最低显存要求：13GB
8精度llama2 70B预测最低显存要求：70GB

4精度llama2 7B预测最低显存要求：3.5GB
4精度llama2 13B预测最低显存要求：6.5GB
4精度llama2 70B预测最低显存要求：35GB

目前模型的参数绝大多数都是float32类型, 占用4个字节。所以一个粗略的计算方法就是，每10亿个参数，占用4G显存(实际应该是10^9*4/1024/1024/1024=3.725G，为了方便可以记为4G)。

比如LLaMA的参数量为7000559616，7B(70亿), 那么全精度加载这个模型参数需要的显存为：

7000559616 * 4 /1024/1024/1024 = 26.08G | 4G * 7 = 28G

[NLP]LLM 训练时GPU显存耗用量估计,自然语言处理,人工智能

如果用用半精度的FP16/BF16来加载，这样每个参数只占2个字节，所需显存就降为一半，只需要13.04G。

目前int4就是最低精度了，再往下效果就很难保证了。比如百川给的量化结果对比如下：

[NLP]LLM 训练时GPU显存耗用量估计,自然语言处理,人工智能

注意上面只是加载模型到显存，模型运算时的一些临时变量也需要申请空间，比如你beam search的时候。所以真正做推理的时候记得留一些Buffer，不然就容易OOM。

如果显存还不够，就只能采用Memery Offload的技术，把部分显存的内容给挪到内存，但是这样会显著降低推理速度。

dtype	每10亿参数需要占用内存
float32	4G
fp16/bf16	2G
int8	1G
int4	0.5G

二进行模型训练

以LLM中最常见的Adam + fp16混合精度训练为例，分析其显存占用有以下四个部分：

[NLP]LLM 训练时GPU显存耗用量估计,自然语言处理,人工智能

GPT-2含有1.5B个参数，如果用fp16格式，只需要1.5G*2Byte=3GB显存, 但是模型状态实际上需要耗费1.5B*16=24GB.

llama2 7B最低显存要求：如果用fp16格式，只需要7G*2Byte=14GB显存, 但是模型状态实际上需要耗费7B*16=112GB.

比如说有一个模型参数量是1M，在一般的深度学习框架中(比如说PyTorch)，一般是32位存储。32位存储的意思就是1个参数用32个bit来存储。那么这个拥有1M参数量的模型所需要的存储空间的大小即为：1M * 32 bit = 32Mb = 1M * 4Byte = 4MB。因为1 Byte = 8 bit。现在的quantization技术就是减少参数量所占的位数：比如用16位存储，那么：所需要的存储空间的大小即为：1M * 16 bit = 16Mb = 2MB。

[NLP]LLM 训练时GPU显存耗用量估计,自然语言处理,人工智能

显存计算

全参数微调LLaMA-2-7B:

1. 开启zero3且不offload时，全参数微调最少需要显存可以估计为n_params(in Billion)16个GB。所以7*16=112GB，大约是1120/80=1.4张80G的A100显卡，大概是1台机器的2张卡。这里估计的只是把模型、梯度和优化器放下需要的显存，前向计算还需要额外的显存。

全参数微调LLaMA-2-70B:

1. 开启zero3且不offload时，全参数微调最少需要显存可以估计为n_params(in Billion)16个GB。所以70*16=1120GB，大约是1120/80=14张80G的显卡，大概是两台机器。这里估计的只是把模型、梯度和优化器放下需要的显存，前向计算还需要额外的显存。

	测试1	测试2
setting	huggingface trainer deepspeed zero3 no-offload flash attention 2 model_max_length 4096 数据长度吃满 A100 80G with IB RDMA	同
nnodes	4	4
ngpus_per_node	8	8
batch_size_per_device	1	3
gradient_accumulate_steps	16	3
global_batch_size (前面4个的乘积)	512	288
每步需要的时间，秒（稳定训练若干步后）	249.56	105.30
nvidia-smi看到的显存占用，MB	71000	75000
训练120k个样本需要的时间	14h50m	11h57m

可以看出，要想训练快，还是要把batch_size_per_device尽量开大一些。

checkpoint大小计算

保存checkpoint的时候只需要模型参数(fp16)和优化器状态(fp32)就行了。

对于70B的模型，使用AdamW训练时优化器的参数量是模型本身的两倍，所以最后算起来每个checkpoint需要70 * 2 + 70 * 2 * 4 = 700GB ，还是非常大的。建议设置一下hf trainer的--save_total_limit number ，把太早的checkpoint删掉，避免集群的磁盘满了。

结论如下：

不考虑Activation，3090的模型容量上限是 24/16=1.5B，A100的模型容量上限是 80/16=5B
- 假设训练的过程中batchsize恒定为1，也即尽最大可能减少Activation在显存中的占用比例，使得我们的理论计算值16Φ更接近真实的显存占用，那么24G的3090的模型容量上限是1.5B（差不多是GPT-2的水平），80G的A100的模型容量上限是5B
考虑Activation，3090的模型容量上限是 0.75B，A100的容量上限是 2.5B
- batchsize为1的训练效率非常低，batchsize大于1才能充分发挥GPU的效率，此时Activation变得不可忽略。经验之谈，一般需要给Activation预留一半的显存空间（比如3090预留12G，A100预留40G），此时3090的模型容量上限是0.75B，A100的容量上限是2.5B，我们实际测试结果接近这个值
- 激活在训练中会消耗大量的显存。一个具体的例子，模型为1.5B的GPT-2，序列长度为1K，batch size为32，则消耗显存为60GB。
[1B, 5B] 是目前市面上大多数GPU卡的分水岭区间
- [0, 1B) 市面上绝大多数卡都可以直接硬train一发
- [1B, 5B] 大多数卡在这个区间的某个值上触发模型容量上限，具体触发值和显存大小有关
- (5B, ~) 目前没有卡能裸训