Efficient GPU Memory Management

image.png

关键词：AI系统、非线性网络、GPU内存管理、资源管理；

一、概要

本文发现了现AI系统内针对非线性网络的GPU显存管理存在不足，于是设计了一套方案来高效管理，相比于vDNN以及SuperNeurons最终最高降低31%的内存占用。

二、动机

1 显存卸载与预取

现阶段的方案：

image.png

该方案在每一层计算结束后均会使用同步语句

而目前该方案存在的问题如下图所示：某些层的计算时间很短，通信时间很长，所以不能简单的只在末尾同步，Motivation如下：

image.png

由图所示，f5与f9就很明显不太适合之前的方案，所以下面的方案就比较合适：

image.png

即不考虑同步的问题，一路向下运行，然后在每一层开始的时候判断一下该层所需的数据有没有在内存中即可。不在的话等待一下。

2 GPU显存管理

现存的GPU显存管理方案不够高效，有如下两种：

image.png

b图中最左边的512释放后不能跟free很好的合并，所以碎片了。

image.png

在非线性网络环境下，中间的200(4)因为后期要用，所以暂时还不能释放，于是隔断了前面的数据，导致新来的400无法被分配。

虽然文章拿这两个内容做Motivation，不过我觉得第一个‘显存卸载与预取’的方案并不是首次提出，所以这里拿来略有一些牵强

三、贡献总结

1 本文发现了一些深度学习的特征，以及数据交换与分配的不足之处；

2 以DFS为方案设计了网络构建方案，构建网络；

3 设计了新型的内存池来放置管理GPU显存；

4 实验；

四、系统设计

1 使用DFS方案构建非线性网络的结构；

image.png

2 内存卸载以及预取的方案；

使用cudaMallocHost()建立CPU存储空间，并直接转移到CPU中。

预取时基于某些原则：需要使得在后向传播的时候不能有两个卷积操作同时在GPU显存中（因为卷积操作属于计算密集型，可以被用来做hiden，太早放进来损失了一些并行的机会）。并且反向传播的candidate也是使用DFS做搜索。

这里我认为并没有讲的很清楚，里面有很多问题没有交代，比如什么时候转移回GPU，如果转移设计的不好的时候，会不会造成性能损失？等等问题

3 内存管理方案

image.png

该方案用可用一句话总结，即：把Tensor分为三类，并根据不同的类设计不同的存放方法：

①Movable tensors：只有一个reference，如上图的前向传播的l1，l3等，用完立刻释放；

②Temporary tensors：存在时间也比较短，但是比①长一些，比如fork层，比如l0，要等到l4结束了l0才能释放；

③Reclaimed tensors：这些数据不进行转移操作，比如Dropout, Softmax, 以及FC，这些数据太小；（这里我认为也不一定，这个方案做的过于粗粒度，通用化并不是很好）

基于上述内容，设计了内存池方案：如下图，即①类型数据从左到右分配；②类型数据从右到左分配；③类型数据用cudamalloc分配，最右边；

image.png

好处：①释放分配比较频繁，所以从左到右可以专门用来存频繁分配的数据；

②Temporary数据存活时间稍微长一些，并且这些数据被分配的时候往往比较靠前，所以释放的时候比较晚，放最右边可以减少碎片的问题（不然都挤在一起，又不释放，就各种碎片了）

③该数据比较少，所以使用cudamalloc分配就可以，性能损失比较少，且这些数据不转移，所以都是顺序生成顺序释放，所以不存在碎片问题；