PyTorch显存机制是什么样的？今天来彻底分析一下。
我从研一开始学习NLP自然语言处理，经常使用PyTorch框架。一开始用的时候对于PyTorch的显存机制也是一知半解，连蒙带猜的，经常来知乎上来找答案。经过两年的研究，现在回过头来看，能从大家的答案中找出不足的地方。但是两年过去了，也没有一篇很好的文章来总结PyTorch的显存机制的方方面面，那么我就吸收大家的看法，为PyTorch的显存机制做个小的总结吧。
实验环境：
OS: Window 11
python: 3.7.4
PyTorch: 1.9.1
GPU: RTX 3060
1 理论知识
1. 深度学习训练过程
开门见山的说，PyTorch在进行深度学习训练的时候，有4大部分的显存开销，分别是模型参数(parameters)，模型参数的梯度(gradients)，优化器状态(optimizer states)以及中间激活值(intermediate activations) 或者叫中间结果(intermediate results)。
为了后面显存分析阐述的方便，我将深度学习的训练定义4个步骤：
1. 模型定义： 定义了模型的网络结构，产生模型参数；while(你想训练):
2. 前向传播： 执行模型的前向传播，产生中间激活值；
3. 后向传播： 执行模型的后向传播，产生梯度；
4. 梯度更新： 执行模型参数的更新，第一次执行的时候产生优化器状态。
在模型定义完之后，2~4循环执行。
2. 前向传播
拿Linear层（或者叫Dense层，前馈神经网络，全连接层等等…）举例：假设他的权重矩阵为W，偏置向量为b，那么他的前向计算过程就是：
这里的X为该层的输入向量，Y为输出向量（中间激活值）
3. 后向传播（反向传播）
参考了这篇文章
https://zhuanlan.zhihu.com/p/22473137
后向传播回来了一个第l+1层的输出误差矩阵
，用以计算该层的梯度和输入误差
4. 梯度更新
接下来就是利用W_diff和b_diff进行更新了：
当然使用Adam优化器的时候，实际的更新过程并没有上面的这么简单。目前用的最多的是AdamW，可以看看这篇文章。
https://zhuanlan.zhihu.com/p/38945390
https://www.ilfishs.com/archives/80
但是使用这一类优化器，也会带来额外的显存开销。对于每一个参数，Adam都会为它准备对应的2个优化器状态，分别是动量(momentum)和方差(variance)，用以加速模型的训练。
2 显存分析方法与Torch机制
1. 分析方法
(1) No Nvidia-smi
我看很多人现在还在用nvidia-smi来看pytorch的显存占用，盯着跳来跳去的torch缓存区分析真的不累吗。（贴一个Torch为什么不用Nvidia-smi看的图）。而且PyTorch是有缓存区的设置的，意思就是一个Tensor就算被释放了，进程也不会把空闲出来的显存还给GPU，而是等待下一个Tensor来填入这一片被释放的空间。有什么好处？进程不需要重新向GPU申请显存了，运行速度会快很多，有什么坏处？他不能准确地给出某一个时间点具体的Tensor占用的显存，而是显示的已经分配到的显存缓冲区和torch在创建cuda进程时所需开销的和，也就是reserved_memory和torch context显存之和。这也是令很多人在使用PyTorch时对显存占用感到困惑的罪魁祸首。
(3) torch.cuda is all you need.
在分析PyTorch的显存时候，一定要使用torch.cuda里的显存分析函数，我用的最多的是torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()，前者可以精准地反馈当前进程中torch.Tensor所占用的GPU显存(注意是只包括torch.Tensor)，后者则可以告诉我们到调用函数为止所达到的最大的显存占用字节数。还有像torch.cuda.memory_reserved()这样的函数则是查看当前进程所分配的显存缓冲区是多少的。
Torch 官方文档
2. PyTorch context开销
-----之前没有提到PyTorch context的开销，做个补充…
我注意到有很多同学在做显存分析的时候是为了在训练的时候可以把卡的显存用满，这个之前没有考虑到呢。其实PyTorch context是我们在使用torch的时候的一个大头开销。
主要参考的是论坛里的这篇讨论：
https://discuss.pytorch.org/t/how-do-i-create-torch-tensor-without-any-wasted-storage-space-baggage/131134
什么是PyTorch context? 其实官方给他的称呼是CUDA context，就是在第一次执行CUDA操作，也就是使用GPU的时候所需要创建的维护设备间工作的一些相关信息。如下图所示
这个值跟CUDA的版本，pytorch的版本以及所使用的设备都是有关系的。目前我在ubuntu的torch1.9上测过RTX 3090和V100的context 开销。其中3090用的CUDA 11.4，开销为1639MB；V100用的CUDA 10.2，开销为1351MB。
感兴趣的同学可以在shell中执行下面这两行代码，然后用nvidia-smi去看看自己的环境里context的大小。然后用总大小减去context的大小再做显存分析
import torch
temp = torch.tensor([1.0]).cuda()
我估计会有人问怎么去减小这个开销…官方也给了一个办法，看看自己有哪些cuda依赖是不需要的，比如cuDNN，然后自己重新编译一遍PyTorch。编译的时候把对应的包的flag给设为false就好了。我是还没有试过，要搭编译的环境太难受了，而且还要经常和库做更新。
3. PyTorch显存分配机制
在PyTorch中，显存是按页为单位进行分配的，这可能是CUDA设备的限制。就算我们只想申请4字节的显存，pytorch也会先向CUDA设备申请2MB的显存到自己的cache区中，然后pytorch再为我们分配512字节或者1024字节的空间。这个在使用torch.cuda.memory_allocated()的时候可以看出来512字节；用torch.cuda.memory_cached()可以看出向CUDA申请的2MB。直观点来说，看图吧，PyTorch的显存管理是一个层级结构。
然后他们又是包含于被包含的关系，即PyTorch Allocated memory使用的是PyTorch Cached Memory里的显存，PyTorch Cached Memory则用的是GPU的显存。
关于Reserved_Memory的具体的分配逻辑，在这里写的很清楚了↓
https://github.com/pytorch/pytorch/blob/master/c10/cuda/CUDACachingAllocator.cpp
总结一下：
- pytorch中的reserved_memory以block的形式存在。
- 一个allocation的显存被释放后，他所在的block可以被重新被allocate.
- 分配器尝试寻找能满足requested size的最小cached block，如果这个block 的大小大于requested size，那么这个block可以被split. 如果没有block了，那么分配器就会调用cudaMalloc向CUDA设备申请显存。
- 如果cudaMalloc失败了，分配器会先尝试释放掉一个足够大的，且没有被split的cached block，并重新尝试allocate。
- 大于1MB的allocation和小于等于1MB的allocation会被存储在不同的pool中。小的请求会放进2MB的buffer里，大的请求会先尝试使用最小的可用free block，或者用cudaMalloc申请一个新的block。
- 为了减小碎片化，在所有可用block都没有充足的大小的时候，1MB到10MB的allocation会使allocator申请一个20MB的block，并在上面进行split；为了进一步减小碎片化，大于200MB的块则不能够被split。大于200MB的超大cached blocks仍可满足小于20MB的请求。
4. PyTorch显存释放机制
感兴趣的可以去看官方文档原文：
https://pytorch.org/docs/1.10.0/notes/cuda.html
CUDA semantics - PyTorch 1.10.0 documentation
https://pytorch.org/docs/1.10.0/notes/cuda.html
简单总结一下，就是在PyTorch中，只要一个Tensor对象在后续不会再被使用，那么PyTorch就会自动回收该Tensor所占用的显存，并以缓冲区的形式继续占用显存。
要是实在看缓冲区不爽的话，也可以用torch.cuda.empty_cache()把它减少，或者加一个环境变量PYTORCH_NO_CUDA_MEMORY_CACHING=1，但是程序速度会变慢哦，(曾经试过，在我的一个实验里慢了3倍)。不得不说层级存储结构yyds。
3 训练过程显存分析
为了让大家方便理解，我这里用torch.nn.Linear(1024, 1024, bias=False) 来做例子。为了省事，loss函数则直接对输出的样本进行求和得到。没办法，想直接执行loss.backward()的话，loss得是标量才行呢。
示例代码：
import torch
model = torch.nn.Linear(1024,1024, bias=False).cuda()
optimizer = torch.optim.AdamW(model.parameters())
inputs = torch.tensor([1.0]*1024).cuda() # shape = (1024)
outputs = model(inputs) # shape = (1024)
loss = sum(outputs) # shape = (1)
loss.backward()
optimizer.step()
1. 模型的定义
结论：显存占用量约为参数量乘以4
import torch
model = torch.nn.Linear(1024,1024, bias=False).cuda()
print(torch.cuda.memory_allocated())
打印出来的数值为4194304，刚好等于1024×1024×4。
2. 前向传播过程
结论：显存增加等于每一层模型产生的结果的显存之和，且跟batch_size成正比。
inputs = torch.tensor([1.0]*1024).cuda() # shape = (1024) memory + 4096
outputs = model(inputs) # memory + 4096
代码中，outputs为产生的中间激活值，同时它也恰好是该模型的输出结果。在执行完这一步之后，显存增加了4096字节。(不算inputs的显存的话)。
3. 后向传播过程
后向传播会将模型的中间激活值给消耗并释放掉掉，并为每一个模型中的参数计算其对应的梯度。在第一次执行的时候，会为模型参数分配对应的用来存储梯度的空间。
loss = sum(outputs) # memory + 512(torch allocate分配最小单位)
temp = torch.cuda.memory_allocated()
loss.backward()
print(torch.cuda.memory_allocated() - temp) # 第一次增加4194304
第一次执行时显存增加: 4194304字节 - 激活值大小
第二次以后执行显存减少: 激活值大小
Note: 由于这个中间激活值被赋给了outputs，所以后面在后向传播的时候会发现，这个outputs的显存没有被释放掉。但是当层数变深的时候，就能明显看到变化了。
为了让大家看到变化，再写一段代码~
import torch
# 模型初始化linear1 = torch.nn.Linear(1024,1024, bias=False).cuda() # + 4194304print(torch.cuda.memory_allocated())linear2 = torch.nn.Linear(1024, 1, bias=False).cuda() # + 4096print(torch.cuda.memory_allocated())
# 输入定义inputs = torch.tensor([[1.0]*1024]*1024).cuda() # shape = (1024,1024) # + 4194304print(torch.cuda.memory_allocated())
# 前向传播loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 + 512print(torch.cuda.memory_allocated())
# 后向传播loss.backward() # memory - 4194304 + 4194304 + 4096print(torch.cuda.memory_allocated())
# 再来一次~loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 (512没了，因为loss的ref还在)print(torch.cuda.memory_allocated())loss.backward() # memory - 4194304print(torch.cuda.memory_allocated())
4. 参数更新
optimizer.step() # 第一次增加8388608，第二次就不增不减了哦
print(torch.cuda.max_memory_allocated()) # = torch.memory_allocated + 8388608
第一次执行时，会为每一个参数初始化其优化器状态，对于这里的AdamW而言，每一个参数需要4*2=8个字节。
第二次开始，不会再额外分配显存。
显存开销:
第一次: 增加8388608字节
第二次及以后: 无增减
5. Note
由于计算机计算的特性，有一些计算操作在计算过程中是会带来额外的显存开销的。但是这种开销在torch.memory_allocated中是不能被察觉的。
比如在AdamW在进行某一层的更新的时候，会带来2倍该层参数量大小的临时额外开销。这个在max_memory_allocated中可以看到。在本例中就是8388608字节。
参考
[1] https://www.ilfishs.com/archives/134-PyTorch显存机制是什么样的？今天来彻底分析一下-小余同学的博客