谁动了我的显存？——深度学习训练过程显存占用分析及优化

在大语言模型时代，不仅语言模型变得越来越大，而且几乎所有的模型都想变得越来越大，试图在模型变大之后观察到一些涌现出来的能力。 模型变大之后，最突出的问题就是显存不够用了。本文对深度学习训练过程中的显存占用问题进行一些具体分析，加深我对训练过程的理解，能够进行一些简单的显存优化操作。如果读者们有更多的相关资料、优化技巧，欢迎在评论区补充。 显存占用概述 深度学习训练过程中的显存占用，大致可以分为三部分： 框架占用，例如pytorch框架的cuda context会占用大约几百MB显存 模型参数相关的占用，比如7B的模型以FP16格式要占用14GB显存。此处还包括优化器、梯度相关的参数占用，全量微调的情况下，梯度与参数一样大，优化器状态是梯度的1~2倍（SGD为1倍，Adam为2倍）。如果使用DDP进行多卡训练，则还需要乘以显卡数量；如果使用FSDP进行多卡训练，显存占用与显卡数无关，但是会增加通信开销。 特征相关的占用，这部分显存占用是最复杂的，因为它与模型的具体计算流程有关。很多地方只会笼统地说这类占用与batchsize成正比，但是具体的比例系数很难分析。 本文希望详细解析特征相关的显存占用到底是多少。 统计方法 我们用一个样例程序，来使用不同的方法、在不同的情况下计算 这样一个简单的函数。具体程序为： import torch# Create two tensors