相信很多人，包括我在内，都对GPU的显存抱有不小的怨念，CUDA out of memory之类的问题一直困扰着咱等，今天这篇文章就是来浅析一下，说不定会有所帮助
首先就一笔带过说一下GPU的重要性吧，以Pytorch为例，就是使用CUDA，cuDNN对深度学习的模型推理时执行的各种计算转换为矩阵乘法进行加速，来达到从猴年马月的运行，到现在几十成百倍的提速。
至于我们爱之深恨之切的显存呢，以数据读取时GPU运行内存的变化作为参考，具体实现机制一般是通过worker进程+队列的方式，让多个worker异步地读取、预处理数据入列，然后主管训练主进程从队列的另一头取数据。如果队列满了、训练进程来不及取数据，worker进程就会阻塞住，运行内存也就不会无限制地增长了。
torch.utils.data.DataLoader(datasets[x], batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)
注意：打印日志、计算eta等无关数据处理和计算的操作也会占用不少CPU处理时间。
先粗略讲一下 batch_size管显存（不完全成正比，模型自身的参数及其延伸出来的数据也要占据显存），num_workers管GPU利用率(使得CPU加载数据和GPU处理速度能够匹配，联想操作系统就能理解)，后面两者都会提到
显而易见, 将数据的加载，预处理，后处理等放在CPU上进行，也就是其他读写的IO任务。并通过调节num_workers数量来调节GPU利用率是个不错的办法。数量设置在一个比较大的范围内比较好（可以考虑4-8），但不是越大越好。因为越大，虽然线程多了，但是切分到每一个线程消耗也大了，所以会增加CPU的负荷，从而降低对GPU的利用。num_workers数一般和batch_size数配合使用。
除此之外，当你的服务器或者电脑的内存较大，性能较好的时候，建议打开pin_memory打开，就省掉了将数据从CPU传入到缓存RAM里面，再给传输到GPU上；为True时是直接映射到GPU的相关内存块上，省掉了一点数据传输时间。
这里注意一下，数据loading的时候建议在dataset的__init__里或者pipeline中通过调用数据导入方法全部load进来，而不是在getitem方法中写。
因此，针对于此，一个粗糙的方法就是在多GPU的时候采取ddp模式(线程不要太多)，用共享内存的方法，让rank0的进程把这些标注数据和扔到共享内存里，然后其他所有进程都只读、映射这段内存，实现了零拷贝。
若只是GPU利用率低的话，第一是增加batch size，增加GPU的内存占用率，尽量用完内存，而不要剩一半，空的内存给另外的程序用，两个任务的效率都会非常低。第二，在数据加载时候，将num_workers线程数设置稍微大一点，推荐是8,16等，且开启pin_memory=True。不要将整个任务放在主进程里面做，这样消耗CPU，且速度和性能极为低下。
其次，考虑模型本身的影响，分别是模型参数和各层的响应，均与模型复杂度有关，以Resnet50为例，有26 million参数，若使用32-bit浮点精度，那么占内存：26M * 32 bit = 99MB。同时他有16 million响应，占内存：16M*32bit = 64MB。
响应和模型参数的数量并没有直接关系。卷积层可以有很大尺寸的响应，但只有很少的参数；激活层甚至可以没有参数。
可能到这里你会觉得才100M多一些，影响也不大啊，非也非也，这时候就是batchsize登场的时候了。
为了有效利用GPU的SIMD机制，要把数据以mini-batch的形式输入网络。
如果要用32 bit的浮点数填满常见的1024 bit通路，需要32个样本同时计算。
在使用mini-batch时，模型参数依然只保存一份，但各层响应需要按mini-batch大小翻倍。
50层的ResNet，mini-batch=32，各层响应占内存：64MB*32 = 2GB
不过，2G还不是终点，在深度学习库中，一般会采用**lowering的方式，把卷积计算转换成矩阵乘法**。
在计算此类卷积时，前层响应X XX需要扩大K^2 倍。
50层的ResNet，考虑lowering效应时，各层响应占内存7.5GB
还要注意的一个问题就是，你的网络训练结束的标准是什么。是以epoch还是iteration为标准。如果是以iteration为标准，那么你加大batch_size大小，网络训练时间自然就会增加，此时需要在batch_size翻倍时iteration打对折。如果是以epoch为标准，那么更改batch_size时不需要做其他更改。
显存占用 ≈ 模型显存占用 + batch_size × 每个样本的显存占用，可以这么简单来看待
之前都是以ResNet50举例，下面换成GPT-2可能更接近现在的网络结构。
引言
对于一个1.5B参数量的GPT-2模型，在16比特精度下需要3GB的内存来存储它的权重（或参数），但它却没法在具有32GB显存的一块GPU上进行训练。模型训练过程中，大部分的内存开销用于model states，例如优化器的状态、梯度和参数。除了模型状态之外，剩下的内存开销则来自于residual states，例如激活值、临时区缓存和内存碎片。
Model States
以Adam为例，它需要存储两部分的优化器状态：time averaged momentum和variance of the gradients。因此，在使用Adam进行模型训练的时候需要有足够的的内存空间来存储动量估计和梯度方差的复制值，除此之外还需要有足够的空间存储模型本身的梯度和权重。优化器的状态通常会占到很大一部分的内存，这在混合精度训练中尤为明显。
在混合精度（fp16/32）训练中，模型参数和激活值是以fp16的格式进行保存，前向和反向传播中也使用fp16的权重和激活值进行计算。然而为了更高效的计算和保证梯度更新的正确（混合精度训练会出现舍入误差），通常会同时拷贝一份fp32的权重和优化器状态。
举例：对于一个参数量大小为 γ 的模型使用Adam和混合精度进行训练。首先它分别需要 2γ 的内存来存储fp16拷贝的参数和权重，分别需要 4γ 的内存来存储fp32拷贝的参数、动量估计和梯度方差。这里用 K 来表示存储优化器状态的所需的额外内存开销倍数，混合精度Adam对应的 K = 12 。综上，训练这样一个模型所需要的内存开销为 2γ + 2γ + Kγ = 16γ。对于一个参数量为1.5B的GPT-2模型，所需的训练内存消耗至少为24GB，远大于fp16格式下模型本身3GB大小的参数量。
Residual States
Activation
在训练过程中激活值会占到很大一部分的内存。Transformer结构的模型的激活值所占内存与 transformers_layers * hidden_dimensions * sequence_length * batch_size 的大小成正比，对于GPT-2模型来说，激活值内存为 12 * hiddem_dim * bsz * seq_length * transformer_layers 。因此，对于一个1.5B参数量、sequence length为1K、batch size为32的GPT-2模型，需要的内存为60GB。
虽然activation checkpointing技术能够以33%的重新计算激活值的运算代价实现激活值所占显存近似根号倍的降低（这就能够将上述模型所需的激活值内存降低到8GB左右），对于更大的模型来说效果还是有限的（100B参数量的GPT模型在使用activation checkpointing的基础上仍需要60G的内存）。
Temporary buffers
保存中间计算结果的临时区缓存，例如梯度的all-reduce，或者为了提升throughput而在all-reduce前将所有的梯度结果融合成为single flattened buffer的梯度标准化计算（gradient norm computation）。对于1.5B参数量的模型，fp32下的flattened buffer内存占用就会达到6GB。
Memory Fragmentation
即使内存大小大于理论上模型训练需要的内存开销，由于内存碎片的存在（空闲内存以不连续的方式出现在不同的位置，导致没有连续的内存来应对模型训练中的需求），还是有可能出现内存不足的错误。训练超大模型时，在某些极端情况下，内存剩余还有30%的时候还会出现OOM（out of memory）的问题。
显存用到刚刃上的简单方法推荐
1.为了有效利用SIMD，如果精度降低一倍，batch大小要扩大一倍。不能降低内存消耗。但可以使用英伟达的apex混合精度加速，pytorch显存直接砍一半。
2. in-place运算：不开辟新内存，直接重写原有响应。很多激活函数都可以这样操作。如nn.ReLU(inplace=True)，也就是尽量使用inplace操作flag。
复杂一些，通过分析整个网络图，可以找出只需要用一次的响应，它可以和后续响应共享内存。例如MxNet的memory sharing机制。
3. 计算换存储：找出那些容易计算的响应结果（例如激活函数层的输出）不与存储，在需要使用的时候临时计算。使用这种方法，MxNet的这个例子能够把50层的ResNet网络占用的内存减小四倍。
还可以看看这里的10多种方法
Trick
最后，如果你能优化算法运算，在网络图结构上优化单位batchsize的显存占用量，自然就可以装下更多的batchsize，来更充分的利用GPU core。
参考
- https://www.zhihu.com/question/476086761
- https://blog.csdn.net/shenxiaolu1984/article/details/71522141
- https://zhuanlan.zhihu.com/p/362556069
- https://blog.csdn.net/qq_45756171/article/details/122910838
- https://zhuanlan.zhihu.com/p/348122818
- https://blog.csdn.net/qq_34405401/article/details/108519823
- https://zhuanlan.zhihu.com/p/520898642
- https://zhuanlan.zhihu.com/p/31558973
- https://www.cnblogs.com/jiangkejie/p/10098995.html
近期评论