详解大模型推理训练占用的显存

最新推荐文章于 2025-03-11 11:02:34 发布

AI有温度

最新推荐文章于 2025-03-11 11:02:34 发布

阅读量1.9k

点赞数 13

文章标签：大模型

本文链接：https://blog.csdn.net/Antai_ZHU/article/details/144409222

版权

AI因你而升温，记得加个星标哦！

在当前开源大模型中，我们可以观察到一个普遍的现象：开源可下载模型的参数规模都集中在7B、13B、65B等特定的数值上，那这种现象背后有着怎样的原因呢？

开源模型是为了让大家易部署易传播，所以模型参数量的设计主要是为了能匹配大众使用显卡的显存，让大家在消费级显卡上就可以把模型跑起来。

我们大众消费级显卡的显存一般为4/8/16/24G，那么大模型参数量和占用显存怎么换算？可分为两部分，大模型推理部署和大模型训练。

目前模型的参数绝大多数都是float32类型，即1个参数占用4个字节。所以一个粗略的计算方法就是，每10亿（1B）个参数，约占用4G显存（其中1G=1024M，1M=1024K，1K=1024字节。除了参数外，在推理时比如输入输出数据、计算图等还需消耗一定的显存）：

$\frac{10^9*4}{1024^3}=3.725G\approx4G$

所以1个7B的大模型部署，大约需要28G，两张16G显卡为32G，刚好可以满足。

如果想用1张显卡部署，就可以用半精度（FP16）来加载模型，这样每个参数只占2个字节，所需显存就降为一半，只需要14G显存，那么一张16G的显卡就可以部署了。

如果有块3070显卡，一般是8G显存，还能部署吗？当然可以！资本会为你考虑的。可以采用int8精度，显存再降一半，仅需7G显存。但是相对应的，随着精度越低，模型效果也会更差一些。

同理，如果使用目前int4的最低精度，有3.5G显存就能跑起来7B模型。

根据上表，就可以推算出大模型部署所需要的显存。比如：

所以说，模型参数量的设计和我们显卡的显存紧密相关。

模型训练的实质就是反向传播，公式如下：

$\theta_{n+1}= \theta_{n} - \eta \nabla L(\theta_{n})$

所以模型训练时候的显存使用主要包括如下几部分：

模型权重。这是最基本的，训练模型肯定得先把模型加载进来吧，模型权重所需的显存和模型部署所需的权重是一样的。

优化器。我们训练最常用的优化是Adam，而Adam需要同时维护参数的一阶动量和二阶动量两个状态的变量，也就说优化器使用显存是模型权重的2倍。如果采用梯度下降法（SGD），则优化器占用显存就和模型权重一样，因SGD只需维护参数的一阶动量。

梯度值。梯度占用显存空间与模型权重相同。

所以说，如果模型想要训练，只看这3部分，需要的显存是至少推理的3或4倍。除此之外，还有以下部分：

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

目前很多工程师都在为降低显存做着优化工作，比如：

随着模型的变大，主要训练方法还是分布式计算：

在这里插入图片描述