显存优化综述

动态图和静态图

动态图是边计算边搭建计算图，中间结果一目了然。

静态图是先搭建计算图，最后喂入输入，直接出结果。

因为静态图在编译的时候进行了一些优化，比如改变了代码的计算过程，以利于GPU更好的并行计算，那么想debug看中间结果就很困难了，有时候看到的和期望的可能会有些出入。另外就是每次当我们搭建完一个动态计算图，然后在反向传播结束之后，整个计算图就在内存中被释放了。如果想再次使用的话，必须从头再搭一遍。

静态图显存优化

对于大网络结构的训练场景，在静态图上的显存优化主要可以分为三个方向：

静态显存分配。由于获得了整张计算图，所以可以去分析每一个 tensor 和每个算子的生命周期。对于生命周期没有重叠的算子，它们是可以共享显存的。
梯度检查点（用计算换显存）。设置一些梯度检查点，剩下的中间结果就先释放掉，如果将来在反向传播的过程中发现前向结果不在显存中，就找到最近的梯度检查点，恢复出被释放的 tensor。
内存交换（用带宽换显存）。把暂时不用的数据从 GPU 上交换到 CPU 上，到了需要的时候，再把它交换回来。

动态图显存优化

动态图无法提前获得全局的计算图信息。因为无法得到每个 tensor 的生命周期，所以静态显存分配不再可用

用计算换显存，也就是动态图版的 Sublinear 显存优化；
用带宽换显存，在 GPU 和 CPU 之间交换内容。

如上图所示，交换耗时比计算耗时高出很多，因此用带宽换显存不合理。

如上图所示，在前向传播中（第一行从左到右），蓝色圆圈表示模型的中间计算结果开始占用显存。一直到前向传播完成，第一行完全变为蓝色圆圈，前面计算所占用的显存都不能释放。

等到反向传播开始（第二行从右到左），随着梯度的计算与完成应用，前向传播保留在显存中的张量才可以释放。

很明显，如果要降低显存占用，就要拿前向传播保存的中间计算结果开刀，这也正是 MegEngine 动态图显存优化的主要方向。

在这里插入图片描述

如上为梯度检查点技术原理示意，前向传播中第三个点为检查点，它会一直保存在显存中。第四个点在完成计算后即可释放显存，在反向传播中如果需要第四个点的值，可以从第三个点重新计算出第四个点的值。

DTR

问题

大部分场景下，网络训练的中间特征占用了绝大部分得到显存，模型权重这是占据了极小的一部分。
这些中间特征在反向传播的算子中还需要用到，所以不好被释放。

解决办法

引入LRU cache的机制，选择代价最低的tensor进行释放，在需要用到的时候进行重计算。

重计算的开销越小越好
释放的tensor占用显存越大越好
该tensor的访问间隔越长越好

LRU cache：距离上次访问时间间隔最长的，进行释放

另外，DTR 论文中还提出，除了重计算带来的开销之外，其他的额外开销主要用于寻找应该被释放掉的最优 tensor。因为在显存中，tensor 停留的时长是不断在变化的，所以只能在需要释放的时候现场计算最优的 tensor。

对此，论文中提出了两个运行时的优化技巧：

不考虑小的 tensor，当 tensor 大小小于候选集中的 tensor 的平均大小的 1% 时，不加入候选集；
每次在需要释放 tensor 的时候，随机采样 sqrt(N) 个 tensor 进行遍历（N 为目前可释放的 tensor 候选集的大小）

在这里插入图片描述

重计算的局限性

释放的显存不连续，形成的显存碎片，无法容纳新的tensor

例如，新的tensor需要100M的显存，为此释放了两个tensor，但是这两个tensor不是连续的，不能被使用，根据释放机制，就会一直释放下去，直到释放出一段连续可用的显存

在这里插入图片描述

在显存池里的排列并非最优，是按照计算顺序排列的

在这里插入图片描述

in-place op 失效

重计算实质上把in-place op变成了非in-place op

in-place op：模型权重会被修改利用，以此来节省显存和cache，原地修改会造成后续的值改变。

DTR：模型权重不可被修改，额外申请资源进行计算，生成新的tensor，分散在显存池中，很难形成连续大显存。

在这里插入图片描述