如何估算transformer模型的显存大小

在微调GPT/BERT模型时，会经常遇到“ cuda out of memory”的情况。这是因为transformer是内存密集型的模型，并且内存要求也随序列长度而增加。所以如果能对模型的内存要求进行粗略的估计将有助于估计任务所需的资源。

如果你想直接看结果，可以跳到本文最后。不过在阅读本文前请记住所有神经网络都是通过反向传播的方法进行训练的， 这一点对于我们计算内存的占用十分重要。

1. total_memory = memory_modal + memory_activations + memory_gradients

这里的memory_modal是指存储模型所有参数所需的内存。memory_activations是计算并存储在正向传播中的中间变量，在计算梯度时需要使用这些变量。因为模型中梯度的数量通常等于中间变量的数量，所以memory_activations= memory_gradients。因此可以写成：

1. total_memory = memory_modal + 2 * memory_activations

所以我们计算总体内存的要求时只需要找到memory_modal和memory_activations就可以了。

 

完整文章：

https://avoid.overfit.cn/post/6724eec842b740d482f73386b1b8b012