n - dag优化_DNN显存优化的终点？Checkmate论文总结

我承认标题有些夸张（但既然你点进来了，就说明它有效果。。。）

对于神经网络的显存优化在近几年越来越热门，因为模型越来越大了，一个显卡直接放不下，所以需要各种各样的方法训练大模型。

最近比较集大成的一个大模型训练的框架就是DeepSpeed，他有以下几种优化手段

1. offload一些不用的tensor（比如前向传播的activation）到cpu memory，等需要他的时候（比如反向传播算梯度的时候）再把它从host传到device
2. 利用model parallel让多张卡一起存一个模型
3. 利用checkpoint机制，丢掉一些activation，反向传播需要的时候再重新算一遍

个人认为第一种和第三种本质上都差不多，都是用时间换空间，不过如果第一个做得好，就成了先进的prefetch，也许就可以完全省掉任何代价。

目前来说研究员最喜闻乐见的显存优化还是方式3，比较简单直观，一键搞定。

不过3需要考虑一个问题，就是扔掉哪些节点：如果扔掉relu的output，那么重新计算的代价就很小，但如果扔掉conv的output，那就需要再算一遍卷积，代价就比较大了。

除此之外，每个节点被需要的次数也不一样，比如如果是sigmoid求导，那么根本不需要output，直接用input算就行，如果是relu，那么本来它的input就被扔掉了（一般relu就是in-place操作），如果再把output扔掉，那可能压根就算不出来了。而且如果网络有分支结构，即一个节点的输出是多个节