组成

PyTorch中显存的占用通常分为5个部分：

1. PyTorch context：只要把任何东西放在GPU中，就一定会占用一定数量的显存（视版本的不同有变化）。

   

   https://stackoverflow.com/questions/43244645/what-is-a-cuda-context

2. model parameters

3. gradients

4. optimizer states

5. intermediate activations

其中intermediate activations占绝对大头，optimizer states的显存使用量大约为model parameters的整数倍（Adam:2, SGD:1）

机制

在分析torch的显存占用时不能使用nvidia-smi，PyTorch使用缓存分配器来管理缓存分配，在缓存分配器的机制下, 一个Tensor就算被释放了，进程也不会把空闲出来的显存还给GPU，而是等待下一个Tensor来填入这一片被释放的空间(即只要一个Tensor对象在后续不会再被使用，那么PyTorch就会自动回收该Tensor所占用的显存，并以缓冲区的形式继续占用显存，所以在nvidia-smi/gpustat中看到的显存并没有减少)。

由上可知，这一指令不能准确地给出某一个时间点具体的Tensor占用的显存，而是显示的已经分配到的显存缓冲区和torch在创建cuda进程时所需开销的和。

API

有3个API用于分析显存：

• torch.cuda.memory_allocated()：反馈当前进程中torch.Tensor所占用的GPU显存(注意是只包括torch.Tensor)
• torch.cuda.max_memory_allocated()：到调用函数为止所达到的最大的显存占用字节数
• torch.cuda.memory_reserved()：当前进程所分配的显存缓冲区

大模型的优化

如果仔细观察传统推理时模型的代码，会发现其加载代码时的步骤如下：

1. 用随机初始化的权重创建模型
2. 从磁盘上加载模型的权重
3. 在模型中加载这些权重

这一过程其实有很多不合理之处：创建模型时的随机权重是无用的，此后在载入checkpoint的state_dict之后，程序中实际存在双份权重。为了解决这一问题，详见instantiating-an-empty-model。

https://huggingface.co/docs/accelerate/usage_guides/big_modeling

而传统推理时的模型分发机制如下：

1. 从模型中加载整个权重到单GPU上
2. 将模型分发给不同GPU

我们能否在载入模型时就将其分发给不同的GPU？详见Loading weights。

Reference

PyTorch显存机制分析