[NLP]LLM 训练时GPU显存耗用量估计

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

 

 

 

 

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

比如说有一个模型参数量是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。

结论如下：

• 不考虑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, ~) 目前没有卡能裸训

LLM Training GPU显存耗用量估计 - 知乎 (zhihu.com) 

[深度学习]大模型训练之框架篇-DeepSpeed_奇思闻影的舒克与贝克的博客-CSDN博客