moDNN: Memory Optimal DNN Traini

image.png

一、前言

这篇文章是2018年DATE上的一篇关于DNN中优化内存的论文，前面也总结了许多相关文章，这些问题的背景出发点均是由于随着模型越来越大，系统GPU内存不足以装下如此大规模的模型，从而使得训练失败。

如何能让GPU装下模型并不太影响训练的效率成为这些工作的研究重点。

本文基于数据释放、预取，动态卷积函数选择、动态Batch大小选择等方向来优化系统。与之前的一些方案相比（如vDNN），本文说其方案更为智能、并且考虑了batch大小。此外，本文没有运用重计算的思想，文章给的原因是因为这样会带来额外的开销（This approach incurs high performance degra- dation. ）。那我们就来看看这篇文章的主要思想是什么。

二、基本概念交代

在分析这些方案前，先交代一些基本背景：

• 1 卷积函数分快慢，慢的消耗空间少：general matrix multiplication (GEMM)；快的消耗空间多：the Winograd algorithm , fast Fourier transform；

• 2 weights (i.e., ΔWs and Ws)消耗比较少的内存，所以始终把他们放在GPU里。

image.png

• 3 要了解：减少batch这种做法，其实是一种trade off，大的batch size会提升性能（一次训练的比较多，并行度高），但是反而会使得空间不足，造成低效函数的使用，从而影响性能；小的batch size训练又比较慢，所以很难说这个训练的性能。

• 4 方案比较复杂，简单来说就是如果一个Tensor有N1种做法，那么就有下图那么多种方案。

image.png

爆炸增长。

所以文章声称：

-1 引入自动调整Batch size的方案以便最优化内存使用；

• 2 更聪明的去制定Tensor转移方案；

• 3 更合理的选择CONV函数，而不像vDNN一样无脑选择最快的；

三、核心思想

image.png

文章介绍到，如果把一个大batch分成许多小的，那么有个很严重的问题就是准确度，不过可以把小batch的梯度整合起来进行操作，从而不影响精度（the gradients are accumulated from all the sub-batches）。

第一个公式首先定义了对于某个Task，在sub-batch size下的最小内存占用。

image.png

输入Tensor+输出Tensor+卷积函数工作空间。

之后又定义出了最小的使用内存量：

image.png

于是：Sw（weights权重的内存量）+batch size为1时的Tensor输入输出大小。

之后是定义如何选择最合适的Batch size。

首先需要profile工作来对不同种类的卷积函数所需空间以及性能进行收集。然后迭代计算batch size，直到下列等式满足：

image.png

这个函数似乎写的不太清楚，没有体现迭代的过程，不过我理解的来说，max是为了取出我所有操作（这里t是值的操作）中最消耗内存的那个。M-Sw是指我内存总量-weights所需内存量。然后对后面的数进行除法操作意思是：当前内存能装下几倍的当前的batch消耗的空间。

由于Mmin与b不是完全的y=bx关系（即如果我b增加2倍，那么Mmin应该会增加的小于等于两倍）

所以导致这个公式在迭代过程中前后的两个b并不一致。逐步逼近直到左右两个b一致。

• 转移策略以及CONV算法选择

image.png

其中包括几个要点需要我们注意：1 将数据转移与计算最大化并行；2 减少额外释放操作需要的时间；3 更明智的预取数据；4 选择最合适的卷积函数；

该流程包括三个主要步骤：（1）对t操作涉及的输入与输出数据做准备；（2）选择合适的卷积函数；（3）选择预取的位置与数据；

• 首先，在3~8行是数据准备阶段，对于task t来说，对于不在GPU中的数据需要给其分配空间，包括输入与输出数据。

如果分配失败，则需要释放掉部分数据（4~5行），如果找不到可以释放的数据，则需要进行defragmentation操作，把所有数据都reload一遍（此时就肯定能装下，因为之前的batch操作保证了内存可以装下最大task中的所有数据）。

并且我们需要记录一下这个额外操作的时间。

• 第二步是对CONV函数进行操作。9~14行能得出，快的卷积函数有可能需要更多的内存释放操作，所以预取数据就会造成delay。所以需要针对“收益”进行最优化的方案选择（下面最后一个公式也于此相关）。此操作造成的延迟同样会被t记录下来。

• 15~17行是预取数据操作，在当前t阶段，要再进入一个循环去搜索t+1到n所有task，如果input(s)或者output(s)不能被分配，直接停止。因为有可能内存不够，需要reload（这个系统是不是有点太厉害？？？这都能预测出来？我不信，，，）

如果至少有一个tensor可以分配，那么去判断是否需要在当前t阶段执行转移（判断会不会造成延迟，不过这里没说清楚，我感觉造成延迟同样可以转移，），如果可以进行，那么就给其分配内存，如果内存分配失败（我也没懂这里为什么又会失败）

干脆这一段伪代码我们理解的时候就全部理解成功的情况，最优情况。不然里面细节不能多想，这代码也解释不清楚。

最后更新时间T。

文中提到，它的方案中一个数据的转移可以与多个计算阶段并行，而不用像vDNN那样，必须要当前计算和当前数据转移必须都完成：

image.png image.png

使用了链表来记录了当前内存中的空闲区域，有针对性的去选择空间把合适的数据放进来。不过有两种数据不能释放：1 下一个task中要用的；2 进来之后还没有用过的；

如果确实没有空间，那么就进行offload all data and then reload the required data for the next task (i.e., defragmentation)。

• 选择合适的预取时间

引入公式来判断当前时间是否合适：

image.png

即在t完成后再取s所需数据是否来得及

image.png

例如：这个图中其实就会引入delay，如果EST(s)正好是s-1结束，那么这个公式就不满足了。

• 卷积算法选择策略

对于算法alg的受益时间：

image.png

为：最慢基础算法执行时间-alg算法执行时间-为了运行快算法而需求腾空内存额外花费的时间-下一个阶段S的延迟时间

所以这个时间可以简单理解为：用快算法带来的受益-用快算法造成的损失。

这个值越大，就越好。

整体来说，这篇文章篇幅太短，所以很多东西写的很不清楚，理解有点小难。不过里面给的一些思想还是很好的，比如可以调整batch size来减少内存使用；其他方面有一点老套，估计在18年有可能算比较好的方案