| 作者：邓哲也 | 旷视 MegEngine 基础架构组工程师

在近年来的深度学习畛域，许多钻研机构和研究者通过增大模型的参数量来晋升模型的体现，获得了十分显著的成绩，一次次令业界称奇。这主观上使得“扩充模型的尺寸”简直一度成为各家竞相追赶的惟一指标。几年间，最先进的模型的参数量已减少了成千盈百倍，但每张 GPU 的显存大小却简直没有增长。这导致大模型的训练往往依赖于巨量的 GPU 卡数。于是，很多想法杰出、有钻研激情的研究者单纯因为资金不足，难以持续从事深度学习的钻研，而近年里的重要科研成果也简直都是被几家头部钻研机构所垄断。从长远看来，这种趋势未必有利于深度学习这门迷信的倒退提高。
作为深度学习训练框架的开发者，咱们除了帮忙用户在一个训练任务中利用更多的 GPU 卡（即分布式训练）之外，还采纳了各种技术手段，以减少每张 GPU 上显存的利用效率，升高研究者的资金老本。减少显存利用效率的常见办法有：

• 生命周期不重叠的算子共享显存；
• 通过额定的数据传输缩小显存占用；
• 通过额定的计算缩小显存占用。

目前已有的办法中大多都要求计算图是动态的，随着越来越多的框架反对动态图模式，是否在动态图训练时最大水平地利用无限的显存资源，成为了评估深度学习框架性能的重要指标。MegEngine 在近期公布的 v1.4 版本中，通过引入 DTR[^[1]](https://arxiv.org/abs/2006.09… 技术并进行进一步的工程优化，提供了一种通过额定计算缩小显存占用的路径，从而让小显存也能训练大模型，享受更大 batch size 所带来的训练收益。在 2080Ti 上，ResNet-50、ShuffleNet 等网络的最大 batch size 能够达到原来的 3 倍以上。本篇文章将从工程实现的角度重点介绍在 MegEngine 中如何应用 DTR 技术对动态图显存进行优化的。

一、背景介绍

1.1 计算图

在深度学习畛域，神经网络模型实质上都能够用一个计算图来示意。它的训练过程能够分为三个局部：前向流传，反向流传，参数更新。

以 y=wx+b 为例，它的前向计算过程为输出 x 和参数 w 首先通过乘法运算失去两头后果 p，接着 p 和参数 b 通过加法运算，失去右侧最终的输入 y。

反向流传需要求出 y 对于 w 与 b 的导数，首先求出 y 对于 p 的导数是 1，p 对于 w 的导数是 x，应用链式法则就能够失去 y 对于 w 的导数是 x。

留神，在反向流传的过程中，会用到前向流传的两头后果。因而当网络结构过大时，显存容量会显著地制约 batch size 的大小。

1.2 动态图显存优化

因而，对于大网络结构的训练场景，在动态图上的显存优化次要能够分为三个方向：

1. 动态内存调配。因为取得了整张计算图，所以能够去剖析每一个 tensor 和每个算子的生命周期。对于生命周期没有重叠的算子，它们是能够共享显存的。
2. 梯度检查点（用计算换显存）。设置一些梯度检查点，剩下的两头后果就先开释掉，如果未来在反向流传的过程中发现前向后果不在显存中，就找到最近的梯度检查点，复原出被开释的 tensor。
   
3. 内存替换（用带宽换显存）。把临时不必的数据从 GPU 上替换到 CPU 上，到了须要的时候，再把它替换回来。

二、动态图显存优化与 DTR 策略

2.1 动态图显存优化

对于动态图的显存优化，相比动态图，最显著的变动是，动态图无奈提前取得全局的计算图信息。因为无奈失去每个 tensor 的生命周期，所以动态显存调配不再可用；梯度检查点还是可行的，且仍然能够寻找最优的检查点；内存替换在动态图中依然也是可用的。因而动态图显存优化有两个方向：

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

上图是从 ResNet-50 这个网络中取出的三个 tensor，别离是卷积、BatchNorm 和 ReLu 的输入 tensor，比照了用重计算和通过带宽来替换它们的工夫开销。能够发现替换的耗时比重计算的耗时广泛大两个数量级左右。因为在 CPU 和 GPU 之间替换数据的耗时取决于 PCIe 的速度，而 8 张 2080Ti 显卡同时训练时，每张卡上分到的替换速度只有 3GB/s 左右。因而，能够确定在动态图中次要的优化方向依然是用计算去换显存。

为了达到用计算换显存的目标，MegEngine 采取以下三步来实现。

1. 实现基础设施：记录产生每个 tensor 的计算门路，使框架反对开释和复原 tensor；
2. 用户提供策略：提供开释 tensor 的接口，由用户显式地调用，框架不须要提供任何策略，只需依照用户的策略去执行每一步，在须要复原 tensor 时现场重计算；
3. 框架寻找策略：框架主动寻找策略并执行它，不须要用户的干涉，做到用户对显存优化齐全无感知。

在了解框架如何开释和复原 tensor 前，咱们须要先理解 tensor 的计算门路。在网络训练的过程中，每个 tensor 的起源只有两种状况：

1. 由内部数据加载进来，例如：输出数据；
2. 是某个算子的输入，例如：卷积层的输入。

对于算子的输入，咱们能够记录这个 tensor 的计算门路（Compute Path），构造体如下所示：

• 每个 tensor 都会有一个 producer，如果 producer 是空，就示意它是由内部数据加载进来的，否则它是一个计算门路，其中：

  • op 示意产生这个 tensor 的算子；
  • inputs 示意这个算子须要的输出 tensor；
  • outputs 示意这个算子产生的输入 tensor；
  • compute_time 示意这个算子理论的运行工夫；
• users 中存储的是所有依赖该 tensor 作为输出的计算门路；
• ref_cnt 示意依赖该 tensor 作为输出的 tensor 数量。

对于如何利用计算历史来开释和复原 tensor，来看一个具体的例子：

首先在 MegEngine 中定义两个 tensor a 和 b，计算 c=a+b。图中每个灰色的长方形都示意显存，假如显存里只能放下 3 个 tensor。这时正好有足够的空间来放下 c，并记录下 c 的计算门路（对应上图黄色框所示）。接着算 d=a*b，因为此时显存里曾经没有空间放下 d 了，须要先把 c 从显存中开释，开释 c 的时候，c 的计算门路依然是保留在 host 端的，然而 c 占用的显存能够被开释掉，此时就有闲暇的地位放 d 了（对应图中第一个绿色框）。如果此时用户想 print(c)，框架发现此时 c 不在显存中，须要立刻把它复原进去。复原之前，发现显存曾经满了，就得先把 d 开释掉，而后依据 c 的计算门路复原出 c，返回给用户（对应图中灰色框）。如果用户持续 print(d)，就先开释 c，复原出 d（对应图中最初的绿色框作）。

通过这个例子能够发现，用户对于它应用的 tensor 是否在显存中是没有感知的，当用户想拜访一个临时被开释的 tensor 时，框架会当场把它复原进去给用户，用户会认为他要拜访的 tensor 始终在显存里。

2.2 DTR 策略

为了使得框架可能主动计算策略，咱们在 MegEngine v1.4 中引入了 DTR——《动静 tensor 重造》这篇论文中的技术，它是齐全动静的启发式策略。它的外围就是当显存超过一个阈值的时候，动静地抉择一些 tensor 将其开释掉，直到显存低于阈值。抉择时会依据三方面对 tensor 进行估价：

1. 重计算的开销越小越好；
2. 占用的显存越大越好；
3. 在显存中停留的工夫越长越好。

另外，DTR 论文中还提出，除了重计算带来的开销之外，其余的额定开销次要用于寻找应该被开释掉的最优 tensor。因为在显存中，tensor 停留的时长是一直在变动的，所以只能在须要开释的时候现场计算最优的 tensor。

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

1. 不思考小的 tensor，当 tensor 大小小于候选集中的 tensor 的均匀大小的 1% 时，不退出候选集；
2. 每次在须要开释 tensor 的时候，随机采样 sqrt(N) 个 tensor 进行遍历（N 为目前可开释的 tensor 候选集的大小）

三、MegEngine 中的工程实现

3.1 动态图外围——Tensor Interpreter

在介绍 DTR 实现之前，首先介绍一下 MegEngine 动态图的外围——Tensor Interpreter（解释器），它会把 python 代码翻译成上面这四种根底操作，顺次解释执行：

• Put：把内部数据从 host 端加载进显存中，失去一个 tensor
• ApplyOp：执行一个算子，它的参数是 op（算子）和输出 tensor，返回输入 tensor
• Del：删除一个 tensor，开释它在显存中占用的空间
• GetValue：获取一个 tensor 的值，须要把数据从显存中加载到 host 端

3.2 开释和复原 tensor 的底层实现

在前文，咱们提到过用户并不知道他拜访的 tensor 以后是否在显存中，然而框架能保障当用户想取得 tensor 的内容时，就算它不在显存中，也能够立刻复原进去。

如上图，若框架要开释掉以后这个 tensor 的显存，reset 它的指针就能够把最底层的显存开释掉。为了未来可能复原出该 tensor，须要在 tensorInfo 中保护一些信息，如果应用 drop（用计算换显存）就须要记录计算历史；如果应用 swap（用带宽换显存），就须要把它先替换到 cpu 上记录一个 host tensor。未来如果用户拜访了该 tensor，框架会查看它对应的 tensorInfo，如果发现曾经不在显存上了，就依据计算历史或 host tensor 在显存中复原出 tensor 的内容返回给用户。

3.3 引入 DTR 后的算子执行

上图是 DTR 外围的伪代码，对于 ApplyOp 办法，以往只须要执行黄色的代码，示意对 input 输出执行 op 算子。

当初因为咱们引入了 DTR 技术，这些输出 tensor 有可能曾经不在显存中了。因而，执行前首先须要给它们打上标记，在这个算子执行完之前不能开释掉这些输出 tensor。而后调用 AutoEvict()，管制以后的显存占用不超过阈值，办法是查看以后的显存占用，如果始终超过阈值就一直地调用 FindBestTensor()算法，再依据启发式估价函数找出最优的 tensor 开释掉。

做完 AutoEvict() 之后，以后的显存占用曾经低于阈值了，此时查看输出的每个 tensor 是否在显存中，如果不在显存中就调用 Regenerate()把它复原进去，而后能力执行以后算子。Regenerate(x)的过程就是重计算 x 的过程，重计算的时候读取 x 的计算历史——op 和 inputs，而后递归调用 ApplyOp 就能够复原出 x。

3.4 tensor 的删除操作

当一个 tensor 不会再被用户和框架应用时，这个 tensor 就能够被删除，从而开释其占用的显存。MegEngine 通过援用计数来管制 tensor 的删除，当援用计数变为 0 的时候，这个 tensor 就会主动发一个删除的语句给解释器。这样带来的问题是，如果真的把这个 tensor 删除的话，它的确能够立刻节俭显存，但会让整体的策略变得十分局限。

比方上面这张图是某张计算图的子图，能够看到一个 9MB 的 tensor 通过一个卷积算子，失去了一个 25MB 的 tensor，再通过一个 Elemwise 算子，失去一个 25MB 的 tensor，再通过 BatchNorm 算子和 Elemwise 算子，失去的都是 25MB 的 tensor。

留神到，因为这里的 Elemwise 算子都是加法，所以它的输出（两个红色的 tensor）在求导的时候都不会被用到。因而，求导器不须要保留住两个红色的 tensor，在前向计算完之后它们实际上是会被立刻开释掉的。这样的益处是能够立刻节俭显存，但在引入 DTR 技术之后，如果真的删掉了这两个红色的 tensor，就会导致图中绿色的 tensor 永远不可能被开释，因为它们的计算源（红色 tensor）曾经失落了，一旦开释绿色的 tensor 就再也复原不进去了。解决方案是在前向的过程中用开释来代替删除，也就是“假删除”——保留 tensorInfo，只是开释掉 tensorInfo 上面对应的显存。这样只须要保留 9MB 的 tensor 就能够开释掉前面 4 个 25MB 的 tensor，并且能够在未来的任意时刻复原出它们。

上图就是 MegEngine 中对 tensor 的删除的伪代码实现，在解释器收到 Del 指令时，会对 tensorInfo 调用 Free()函数，依据以后的状态是否是前向计算来决定做真删除还是假删除。假删除的实现很简略，打上删除标记，开释掉 tensorInfo 治理的显存即可；真删除的实现比较复杂，首先更新产生该 tensor 的输出 tensor 的 ref_cnt，而后调用 RemoveDep()查看所有依赖该 tensor 作为输出的 tensor，如果它们不在显存中，必须当初调用 Regenerate 复原出它们，因为一旦以后 tensor 被真删除，这些 tensor 就复原不进去了。

做完了上述操作之后，就能够真正开释掉该 tensor 对应的 tensorInfo 了。开释完还须要递归地查看 x 的计算历史输出 tensor，如果这些 tensor 中有 ref_cnt=0 且被打上删除标记的，就能够执行真删除。

3.5 训练耗时比照

下图是 MegEngine 的 DTR 实现与原论文在 PyTorch 中的实现在 ResNet-1202 上的训练状况比照。请留神试验用的显卡不同，所以从数据上看 MegEngine 稍快一些。在显存治理上 MegEngine 要更好一些，因为在 11G 的显卡上依然能跑 batchsize=100 的训练。除了论文中尝试的最大 batchsize=140 之外，咱们尝试了更大的 batch size，也都是能够运行的。

上面是 MegEngine 框架上开启不同显存优化的训练耗时比照，baseline 是在动态图模式下不加任何显存优化运行的后果。首先是两个常见的模型——ResNet-50 和 ShuffleNet，能够发现开启 DTR 优化后极限 batch size 超过了动态图 Sublinear 和 baseline，且在 batch size 雷同时耗时和 Sublinear 持平。

下面的两个模型都是偏动态的，所以咱们能够用动态图的 Sublinear 显存优化来做比照，而上面这个 SPOS 网络就比拟非凡，它是一个从输出到输入有多条门路能够更新的大网络。在训练过程中，每一轮会随机采样去更新某一条门路，这就导致每轮执行的语句可能不雷同。对于这种网络，在动态图里实现会比拟天然。因而，这里只取了动态图 DTR 优化的后果与 Baseline 比拟。不论是单卡还是八卡，动态图的极限 batch size 都在 100，如果关上 DTR 能够跑到 250 甚至更大。

3.6 碎片问题和优化办法

在实现 DTR 的过程中，咱们发现随着 batch size 增大，每次产生的 tensor 的显存占用也会增大，tensor 大了之后，显存中能存下的 tensor 数量就会变少，重计算次数就会增多，tensor 的生成和开释会越来越频繁，导致碎片问题十分重大。例如：尽管以后闲暇显存有 1G，然而它们扩散在很多个小的闲暇块中，如果此时有一个 1G 的显存申请无奈满足，就会触发碎片整顿操作，对性能造成微小影响。

对于这个问题，咱们提出了三种可能的优化：

• 参数原地更新

之前在 MegEngine 中很少有 inplace 的操作，如果模型自身参数特地微小，每次更新参数就相当于挪动了一个微小的 tensor 的地位，可能产生出更多的碎片。解决方案是关上 INPLACE_UPDATE 的环境变量，原地更新这些参数，能够缩小局部碎片。

• 改良估价函数

咱们对 DTR 的启发式估价函数做了一个小小的改良，引入了一些碎片相干的信息，心愿换出的 tensor 除了本人占用的显存越大越好之外，还心愿它在显存中两端的闲暇显存块大小之和越大越好。

$$
f(t)=\frac{\text{计算耗时}^\alpha}{\text{停留时长}^\beta(\text{显存 + 闲暇段大小})^\gamma}\delta^{\text{重计算次数}}
$$

此外咱们还引入了重计算次数这一惩办系数，心愿每个算子被重算的次数尽量平均。以及对于函数中的四个属性，增设了一些超参数，这样咱们能够通过扭转这些超参数来使启发式策略侧重于不同的属性。

• 动态布局策略

一个更无效的办法就是对于每轮的执行序列都雷同的网络，咱们把它看作是动态图，把网络动态化，就能够使用动态的显存调配，不仅完全避免碎片整顿，还能够显著升高显存峰值，去尝试更大的 batch size。

比方下图在 ResNet-50 上，batchsize=400 时，动态分配显存的峰值为 9595MB，动态调配显存的峰值为 8549MB，升高了 10% 左右。一旦动态调配显存之后，碎片问题再也不会产生。

四、将来工作方向

如果想要彻底施展出重计算的能力——训练出更大的模型和尽可能大的 batch size，一种可能的办法是侧重于模型的动态化。因为动态图尽管十分好写，然而在 batch size 较大时，受碎片问题的影响比拟大；在动态图上，能够享受到所有动态图优化的益处，比方动态显存调配、图优化技术等等。

更宏观地，咱们心愿形象出一套同时实用于动态图和动态图的显存优化策略，如下图所示：

不论是动态图的 Sublinear 优化还是动态图的 DTR 优化，都能够看作是对执行序列做了一个 Seq2Seq 的变换——在序列中退出了一些 drop 和 recompute 语句。区别在于，动态图是取得了整个运行序列后，计算最优的开释和重计算序列；动态图则是在解释执行序列的过程中，当场插入 drop 和 recompute 语句。执行序列有两种形式：动态图 Imperative Runtime 解释执行和动态图 Computing Graph 编译执行。Profile 会在理论运行中，记录每个算子的运行工夫、每个 tensor 在显存中停留的时长等运行时信息，之后用户能够依据 profile 的后果去调整计算序列。这样用户不须要理解底层的执行逻辑，也不必批改框架的源代码就能够针对不同的模型定制策略。

参考文献：

[1] Kirisame M, Lyubomirsky S, Haan A, et al. Dynamic tensor rematerialization[J]. arXiv preprint arXiv:2006.09616, 2020.

附：

GitHub：MegEngine 天元
官网：MegEngine- 深度学习，简略开发
欢送小伙伴退出咱们 MegEngine 技术交换 QQ 群：1029741705