使用GPU训练大神经网络之 显存控制策略与提升

本文是“[PPoPP18] SuperNeurons:Dynamic GPU Memory Management for Training Deep Neural Networks”的论文阅读笔记，论文、talk视频PPT等详见作者主页：https://linnanwang.github.io/
参考资料：
https://linnanwang.github.io/
How to Optimize Data Transfers in CUDA C/C++
How to Overlap Data Transfers in CUDA C/C++
Memory Allocation Overhead
cudamalloc slow
vDNN: Virtualized Deep Neural Networks for Scalable, Memory-Efficient Neural Network Design
Training Deep Nets with Sublinear Memory Cost
Arbitrary 2D convolution
Fast Algorithms for Convolutional Neural Networks

---

GPU显存控制策略与提升

该论文对时下各个框架中显存控制的思想做了汇总，并做了一些技术上的提升。整体上分为三个部分：
1. Live Analysis：即资源回收，回收无用tensor的显存。
2. offloading&prefetching：将装不下的tensor暂存到内存等其他设备中。
3. Recomputation：对于占用显存多、计算消耗小的tensor，不进行存储，需要使用时重新计算。

1. Live Analysis

其思想主要是资源回收，对于后续操作不依赖的tensor进行回收，释放相应的显存空间。
关于如何分析依赖关系，笔者猜测应该有许多基于计算图的算法。论文中给出了一个简单的$O(N^2)$算法（$N$是网络层数）：维护一个需要保留的tensor集合，每层计算结束后加入当前层输出，并判断集合中的tensor是否被后续计算所依赖，如不被依赖则删除。
性能提升：简单的新增和删除tensor涉及到许多的cudaMalloc和cudaFree操作，会占用较多时间。可以一次性申请足够多的空间，然后程序自身维护空间内的赋值，标记tensor的活跃度等操作。详情请参考：Memory Allocation Overhead ，How to Optimize Data Transfers in CUDA C/C++。

2. Offloading & Prefetching

其思想主要是将GPU中装不下的tensor暂存到内存等其他设备中。论文中的Unified Tensor Pool (UTP)说的是 其可以应用于各种与Local GPU可以交互的设备之间，但实际实现时该runtime仅对CPU内存和GPU显存之间做传输。
有效的GPU相关数据传输技巧详见：How to Optimize Data Transfers in CUDA C/C++，How to Overlap Data Transfers in CUDA C/C++。简单说，对于CPU不使用的数据应用pinned memory；将数据计算和数据传输实现在不同的流中，使得计算和通信可以overlap同步。
论文中讨论了对哪些tensor可以暂存，标准是通信时间不会远大于计算时间，从而不至于拖慢训练速度；并且该tensor显存占比大，使得传输后可以释放充足的显存空间。最终论文决定仅对卷积层实现如上操作。

如果简单的将 所有计算过并且对当前层计算无用的卷积层输出 暂存到CPU中，会涉及许多设备间通信，极有可能拖慢训练速度。论文中采用了类似缓存的思想，定义了Caching tensors on GPU DRAM。即只有在空间不足的情况下才将tensor暂存，选取暂存tensor的策略为Least Recent Used (LRU)，即最长时间未被使用的tensor将被暂存到CPU内存中。论文中使用LRU的motivation是BP算法中最新使用的tensor最早用于梯度值计算。

暂存的思想在vDNN: Virtualized Deep Neural Networks for Scalable, Memory-Efficient Neural Network Design已经提出，但其中没有对应用的网络层做限制，存在计算时间远小于通信时间，难以覆盖从而拖慢训练速度的情况（据论文）

3. Recomputation

对于POOL，ACT，LRN和BN等层，其输出占用较多的显存，但是计算起来却很便宜。可以在每次使用之前重新计算相应层的输出，从而节约显存。
因为不是对所有层均进行重计算，所以可以对需要重计算的层汇总为一个个区块，称为recomputation segment（据论文）。在每个区块内，存在两种重计算的策略：speed-centric和memory -centric，即速度优先和显存优先。速度优先存储每一层重计算的结果，直至不再需要或者区块结束，从而仅需要对每一层进行一次重计算，计算复杂度为$O(N)$，多存储的tensor数为$N$（其中$N$为该recomputation segment内的神经层数）。memory优先则是每次仅存储重计算的输出，而不存储中间结果；若中间结果在后续计算中被依赖，则在需求时重新计算；其计算复杂度为$O(N^2)$，多存储的tensor数为$1$。

论文中提出了一种cost-aware的重计算策略，即判断是否有足够的显存完成speed-centric，如有则使用speed-centric，否则使用memory-centric。笔者预估大多数情况下均可满足speed-centric的条件，从而保证了足够的训练速度；同时对于显存占用大的网络层，存储其中间结果的显存牺牲也很大，从而使用memory-centric收益较多。下图是论文中的测试结果，可以看到cost-aware方法保留了speed-centric和memory -centric的优点，是一个很好的折中策略。

重计算的思想在Training Deep Nets with Sublinear Memory Cost已经提出，但其中忽略了各个网络层显存占比不均匀的情况，对所有区块均采用speed-centric策略，从而最高显存占用增高（据论文）。

已有框架显存控制策略的对比与分析

论文中和已有框架（Caffe, Tensorflow, MXNet, Torch）做对比，分析已有框架显存控制策略的优缺点，这里汇总如下（需要注意的是，论文中没有提及所比较各个框架的版本数，各数据来源的详细信息等）。
Caffe和Torch通过重用前向计算的tensor存储后向传播的梯度值，从而节约显存空间。对于非线性结构网络，其需要额外的空间从而满足其复杂的依赖关系。重用tensor属于Live Analysis的变形，本质上是重复利用无用tensor的空间。
MXNet实现了基于计算图的资源回收，即Live Analysis；还实现了speed-centric的recomputation，但忽略了各层显存占比分布不均的情况，对所有层均使用speed-centric。
Tensorflow也实现了基于计算图的资源回收，即Live Analysis；还实现了将long-lived数据暂存于CPU内存的策略，但是没有最优化两者之间的数据通信（例如没有使用pinned memory），导致了至少50%（据论文）的通讯速度损失。
所有框架均使用static的显存控制策略，而不是动态的进行显存控制；所有框架也没有使用conv workspace或者对其做优化。
最终的各个框架对比结果如下：

Brief Overview of Paper

Motivation简述如下：
1. 当前深度学习模型（如ResNet, DenseNet, Inception v4）需要大量的显存支持，用于网络存储计算。
2. 深度神经网络具有 能力随深度和广度增加而增强 的趋势（详见ResNet论文）。
3. GPU的DRAM有限，目前最大的是24GB，已经难以训练batch_size=32的inception v4 (44.3GB)。
4. GPU可以极大的提升神经网络训练效率（DL-driven novel architecture designs and massive parallelism)。
综上，我们需要一个可以使得GPU训练大神经网络模型的方法策略。

已有方法（按照论文中介绍）分为三种：Model Parallelism，Data Parallelism，和Parameter Reduction
1. Parameter Reduction是对参数做剪枝压缩，对inference阶段更有效。但训练过程中参数所占显存不多，对于大模型GPU训练效果提升不明显。
2. Model Parallelism对模型进行拆分，使得不同模块运行在不同设备中（e.g. AlexNet)。这种方法不需要更改训练过程，但需要特别的模型设计，或者大量的设备间通信，最终极大地降低了效率。
3. Data Parallelism对数据进行拆分（即拆分batch），需要特别设计的异步优化算法。这种方法仅要求少量的设备间通信（梯度值传输），是目前主流常用的方法。
即便存在Data Parallelism可以解决当前内存不足的问题，其仍然无法训练过大的神经网络（例如batch size=1的情况下，模型大小仍然大于显存的模型训练）。

论文中讨论的是：在Data Parallelism框架下，对单GPU的显存控制做优化，使得单GPU可以训练更深更广或者batch size更大的神经网络模型。
论文中claim：最终的runtime可以使得最大显存占用为maximum memory usage among layers（这也是layer粒度所能达到的最优结果），同时保持足够好的训练速度。下图是论文中的结果图，可以看到速度是所有架构中的先进水平，并且能够支持最大的batch size训练：

论文中还强调其定义的runtime是动态的(dynamic)，区别于现有架构中的static。Motivation是
1. 当代的神经网络模型有许多不遵循传统的线性结构（即如传统的AlexNet满足每一层的入度出度均为1，而当代的DenseNet等每层的入度出度不定），网络结构复杂，固定的显存控制策略难以应对复杂的情况。

2. 各层之间的显存占比分布不均，对各层使用相同的策略并非最优解。
所以，需要一个针对特定情况自适应的动态显存控制策略。

笔者对论文的疑问：
1. 虽然题目以及Introduction等部分强调了runtime是动态的，当代神经网络模型是非线性不可预测的，论文中分析以及evaluation的部分却没有对此进行充分的验证（关于非线性模型的分析多基于上诉结果图），所验证的非线性模型也是如ResNet以及Inception等较规律经典的模型，没有看到对DenseNet或者其他复杂情况的验证。
2. 论文中大量的分析比较都是基于经典简单的线性架构模型，例如Live Analysis使得显存近乎减半的分析对DenseNet就不成立；论文中分析显存占用的公式中默认 处理过的网络层forward输出和gradient无需存储，这对非线性架构网络也不成立。
3. 在与其他架构比较时，没有注明架构的版本号以及详细的数据获取方式。
4. runtime的速度提升来源于对conv space的使用，而其他架构并没有使用conv space（据论文）。论文甚至benchmark了不同可用显存情况下，不同的cudnn卷积算法的速度，从而选取最快的那个（详见Fast Algorithms for Convolutional Neural Networks）。所以上诉结果中的速度比较对笔者来说，并不能说明相应的显存控制策略可以保留较好的训练速度。
5. 论文中的runtime对RNN没有提供支持，conv space对于RNN也没有提升，猜测该runtime应该对一般的神经网络模型贡献有限。
论文作者在talk中说会在年中放代码，希望相关疑问能够从中得以解决。

综上，笔者仅总结论文中关于 显存控制策略以及提升 部分的思想，而不关注其中性能分析、结果介绍的部分。