大模型GPU显存占用计算_大模型显存计算

以参数量13B大模型为例，

其中B是Billion，代表十亿参数，13B就是130亿参数

其中每个参数全精度是fp32，也就是float32，占用32位bit，也就是4byte字节。

$1GB= 1024MB= 1024^{2}KB=1024^{3}Byte$

那么全精度13B模型占用$13\times 10^{9}\times 4Byte\div 1024^{3}\approx 48.4GB$

全精度全参数微调训练

训练的显存主要由四部分组成：模型参数、梯度参数、优化器参数、数据的中间计算结果

如果进行全量fp32精度训练，模型参数需占用48.4GB，每个参数对应一个梯度，那么梯度参数同样需占用48.4GB。

优化器参数，是根据梯度来计算的，以Adam优化器来说，其计算和更新如下：

$m_{t}=\beta _{1}m_{t-1}+\left ( 1-\beta _{1} \right )g_{t}$

$v_{t}=\beta _{2}v_{t-1}+\left ( 1-\beta _{2} \right )g{_{t}}^{2}$

$\hat{m}_{t}= \frac{m_{t}}{1-\beta _{1}^{t}}$

$\hat{v}_{t}=\frac{v_{t}}{1-\beta _{2}^{t}}$

$\theta _{t+1}=\theta _{t}-\frac{\eta }{\sqrt{\hat{v}_{t}}+\epsilon }\hat{m}_{t}$

其中，$g_{t}$是梯度，$\beta _{1}$和$\beta _{2}$是两个指数加权平均值的衰减系数，$\hat{m}_{t}$和$\hat{v}_{t}$是梯度的偏差纠正后的移动平均值，$\theta_{t}$和$\theta _{t+1}$是更新前和更新后的参数，$\eta$是学习率，$\epsilon$是一个很小的常数，用于避免除零。

可以看到，其中$m_{t}$和$v_{t}$要计算并且存储留作下一步计算使用，$\hat{m}_{t}$和$\hat{v}_{t}$两个也需要计算占用显存，所以总共至少占用4倍的参数显存量，$4\times 48.4=193.6GB$。

对于数据的中间计算结果，我们来大致计算一下需要多少显存空间。

输入embedding层：

输入参数维度是[batch_size,sequence_length]，经embedding层后维度是[batch_size,sequence_length,hidden_size]。该计算仅需一次，不需要像transformer结构有多层

transformer中的计算：

attention计算 ： $Attention(Q,K,V)=softmax(\frac{QK^{T}}{\sqrt{d_{k}}})V$

        经QKV矩阵embedding后，QKV三矩阵都是[hidden_size,hidden_size]，计算后有3个输出结果维度都是[batch_size,sequence_length,hidden_size]。

        $Q\cdot K^{^{T}}/\sqrt{d_{k}}$计算attention_score，维度[batch_size,num_attention_heads,sequence_length,sequence_length]。

        softmax及dropout后获得attention_probs，维度[batch_size,num_attention_heads,sequence_length,sequence_length]。

        乘V获取结果，维度[batch_size,num_attention_heads,sequence_length,$d_{k}$]。这里$d_{k}$=hidden_size/num_attention_heads。

        维度变换获取attention计算结果，维度[batch_size,sequence_length,hidden_size]。

        进行add&norm，具体要经过dense、dropout、残差连接、LayerNorm，维度[batch_size,sequence_length,hidden_size]。

FFN前向神经网络：

        经过前向神经网络计算后的维度[batch_size,sequence_length,hidden_size]，但它计算中有个状态参数量会更大，我们计算占用取这个维度[batch_size,sequence_length,intermediate_size]。

        与attention计算结果进行add&norm，维度[batch_size,sequence_length,hidden_size]。

以上1层transformer encoder的计算参数量有：

$7\times batch\_size\times sequence\_length\times hidden\_size+2\times batch\_size\times num\_attention\_heads\times sequence\_length^{2}+batch\_size\times sequence\_length\times intermediate\_size$

我们以LLaMA或Baichuan的13B模型为例，它们的参数是：隐藏层维度hidden_size: 5120，最大序列长度model_max_length: 4096，intermediate_size: 13696，num_attention_heads: 40，层数num_hidden_layers: 40。虽然模型有40层，但下一层的计算结果会覆盖上一层的结果，不然40层占用显存太大了，现在的显存都放不下。

以batch_size为1，长度最大计算参数量$1\times 4096\times 5120+7\times 1\times 4096\times 5120+2\times 1\times 40\times 4096^{2}+1\times 4096\times 13696=1566048256$

显存占用$1566048256\times 4Byte/1024^{3}\approx 5.83GB$

所以经目前粗略估算，全参数全精度微调，batch_size=1，至少需要显存$48.4\times 6+5.83\approx 296GB$，这里我们还忽略了一些中间变量的情况下。

全参数混合精度训练

混合精度训练时，计算时用半精度（fp16,bf16），参数更新用全精度，所以需要全精度时一半的显存占用，再加上一个模型全精度的参数量，显存占用$296\div 2+48.4\approx 196GB$。