LLMs模型参数与显存占用的估算

最近找实习面试的时候会被问如何利用模型参数来估算显存占用，正好在这里总结一下：

1. 推理阶段

        推理阶段只需考虑将模型加载到显卡上就可以了，假如说现在有个nB的LLMs，如果以半精加载，那么显存的占用情况如下：

                        ​​​​​​​        

也就是说在以fp16加载模型的话,1B需要1.86GB的显存，其他精度也可以这也计算，参考下表：

LLMs加载精度	1B参数量所需显存(GB)
fp32	3.72
fp16 / bf16	1.86
int8	0.93

2. 训练阶段

        训练阶段如果采用Full finetuning的策略，依旧是考虑以半精训练，训练过程大致需要考虑三个因素，即模型参数、优化器状态参数、梯度缓存。

1. 模型参数 

        与推理阶段一样的参数占比，1.86nGB

2. 优化器状态参数

        其实主要是维护两个状态一阶矩估计和二阶矩估计，每个矩估计的参数量都跟原模型差不多，其中两个状态的参数如果都设置为半精，那么就是1.86nGB+1.86nGB=3.72nGB

3. 梯度缓存 

        与推理阶段一样的参数占比，1.86nGB‘

那么对于nB的LLMs进行Full finetuning就至少需要7.44nGB的显存。

如果全精的话，进行Full finetuning就至少需要14.88nGB的显存。

如果采用lora训练策略的话，可训练的参数其实与target_model有关，替换就可以，混合训练策略这部分可能会导致模型加载精度与优化器状态参数的精度不一致，但是仍然可以做估算，另外如果训练过程中使用了deepspeed的offload的话，会把部分优化器参数状态给copy到CPU上，可以减少显存的占用。同样的deepspeed库也给出了显存估计的方法，可以参考DeepSpeed使用指南(简略版)_Reza.的博客-CSDN博客讲的比较详细。

以上就是根据模型参数大概估计显存的方法，但是实际训练过程中影响显存占用的因素很多，所以只能粗略估计个数量级。