Pytorch显存机制与显存占用(一) —— 理论分析(memory)

PyTorch在进行深度学习训练的时候，有4大部分的显存开销，分别是

      ①模型参数(parameters)

      ②模型参数的梯度(gradients)

      ③优化器状态(optimizer states)

      ④中间激活值(intermediate activations) 或者叫中间结果(intermediate results)。

分别在深度学习的4个步骤中产生

      1. 模型定义：定义了模型的网络结构，产生模型参数；

while(你想训练):

      2. 前向传播：执行模型的前向传播，产生中间激活值；

      3. 后向传播：执行模型的后向传播，产生梯度；

      4. 梯度更新：执行模型参数的更新，第一次执行的时候产生优化器状态。

在模型定义完之后，2~4循环执行。

Pytorch显存分析

      用nvidia-smi或者gpustat来看pytorch程序的显存占用不是很合适的

      因为PyTorch是有缓存区的设置的，意思就是一个Tensor就算被释放了，进程也不会把空闲出来的显存还给GPU，而是等待下一个Tensor来填入这一片被释放的空间。

      有什么好处？进程不需要重新向GPU申请显存了，运行速度会快很多，有什么坏处？他不能准确地给出某一个时间点具体的Tensor占用的显存，而是显示的是已经分配到的显存和显存缓冲区之和。

      这也是令很多人在使用PyTorch时对显存占用感到困惑的罪魁祸首。

torch.cuda is all you need

在分析PyTorch的显存时候，一定要使用torch.cuda里的显存分析函数

我用的最多的是

torch.cuda.memory_allocated()
和
torch.cuda.max_memory_allocated()

前者可以精准地反馈当前进程中Torch.Tensor所占用的GPU显存，后者则可以告诉我们到调用函数为止所达到的最大的显存占用字节数。

还有像

torch.cuda.memory_reserved()

这样的函数则是查看当前进程所分配的显存缓冲区是多少的。

memory_allocated+memory_reserved就等于nvidia-smi中的值啦。

Pytorch context开销

      很多人在做显存分析的时候是为了在训练的时候可以把卡的显存用满，其实PyTorch context是我们在使用torch的时候的一个大头开销。

      什么是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。

      想测试的话可以执行下面这两行代码，然后用nvidia-smi去看看自己的环境里context的大小。然后用总大小减去context的大小再做显存分析。

import torch
temp = torch.tensor([1.0]).cuda()

      怎么去减小这个开销...官方也给了一个办法，看看自己有哪些cuda依赖是不需要的，比如cuDNN，然后自己重新编译一遍PyTorch。编译的时候把对应的包的flag给设为false就好了。我是还没有试过，要搭编译的环境太难受了，而且还要经常和库做更新。

Pytorch显存分配机制

      在PyTorch中，显存是按页为单位进行分配的，这可能是CUDA设备的限制。就算我们只想申请4字节的显存，pytorch也会为我们分配512字节或者1024字节的空间。

     即就算我们只想申请4字节的显存，pytorch也会先向CUDA设备申请2MB的显存到自己的cache区中，然后pytorch再为我们分配512字节或者1024字节的空间。这个在使用torch.cuda.memory_allocated()的时候可以看出来512字节；用torch.cuda.memory_cached()可以看出向CUDA申请的2MB。直观点来说，如图所示，PyTorch的显存管理是一个层级结构。

Pytorch显存释放机制

      在PyTorch中，只要一个Tensor对象在后续不会再被使用，那么PyTorch就会自动回收该Tensor所占用的显存，并以缓冲区的形式继续占用显存。

      要是实在看缓冲区不爽的话，也可以用torch.cuda.empty_cache()把它归零，或者加一个环境变量PYTORCH_NO_CUDA_MEMORY_CACHING=1，但是程序速度会变慢哦(试过在一个实验里慢了三倍)

import torch

# 模型初始化
linear1 = torch.nn.Linear(1024,1024, bias=False).cuda() # + 4194304
print(torch.cuda.memory_allocated())
linear2 = torch.nn.Linear(1024, 1, bias=False).cuda() # + 4096
print(torch.cuda.memory_allocated())

# 输入定义
inputs = torch.tensor([[1.0]*1024]*1024).cuda() # shape = (1024,1024) # + 4194304
print(torch.cuda.memory_allocated())

# 前向传播
loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 + 512
print(torch.cuda.memory_allocated())

# 后向传播
loss.backward() # memory - 4194304 + 4194304 + 4096
print(torch.cuda.memory_allocated())

# 再来一次~
loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304  (512没了，因为loss的ref还在)
print(torch.cuda.memory_allocated())
loss.backward() # memory - 4194304
print(torch.cuda.memory_allocated())

在模型的定义部分，显存占用量约为参数占用量*4 (tensor的类型默认是float32，一个float32占4字节)

所以linear1是1024*1024*4 = 4194304

linear2是1024*4 = 4096

所以第一个输出是4194304，第2个输出是4194304+4096=4198400

输入数据是1024*1024，显存占用是1024*1024*4

所以第三个输出是4198400+1024*1024*4=8392704

在前向传播过程，显存增加等于每一层模型产生的结果的显存之和，且跟batch_size成正比

inputs的shape是1*1024*1024, 经过linear1的结果还是1*1024*1024

所以第四个输出是8392704+1024*1024*4=12587008

经过linear2的结果是1*1024

所以第五个输出的12587008+1024*4=12591104

然后loss存储最后的结果求和，按理说loss只用个4字节就好了，但是在PyTorch中，显存是按页为单位进行分配的。就算我们只想申请4字节的显存，pytorch也会为我们分配512字节

所以第六个输出12591104+512=12591616

反向传播中，后向传播会将模型的中间激活值给消耗并释放掉，并为每一个模型中的参数计算其对应的梯度。在第一次执行的时候，会为模型参数分配对应的用来存储梯度的空间。

反向传播部分没太看懂。。。。

以后再补吧

PyTorch显存机制分析 - 知乎