显存不够，框架来凑：两行代码显存翻倍，2080Ti也能当V100来用

　　机器之心原创

　　作者：思

2080Ti 竟然可以当 V100 来用，这个功能有点儿厉害。

　　自深度学习大潮兴起，模型就朝着越来越大、越来越 「深」 的方向发展。

　　2012 年，拥有 5 个卷积层的 AlexNet 第一次在视觉任务上展现出强大的能力。在此之后，基础模型就开始「深」化起来：2014 年的 VGG-Net 达到了 19 层；2015 年的 ResNet、2017 年的 DenseNet 更是将深度提升到了上百层。

　　模型大小的提升极大地提高了性能。因此，各大视觉任务都将 ResNet、DenseNet 等当作基本的 BackBone。但与此同时，模型的增大也意味着对显存的需求随之变高。

　　为什么 GPU 显存如此重要？

　　九年前，Hinton 等人率先用两张 3GB 显存的 GTX 580 GPU 高效训练 AlexNet。在此之后，显存需求与模型大小就一直同步增长。打比赛想要取到好成绩、做实验想要超越 State of the art 效果、做工程想要拟合庞大的业务数据等等，这些都离不开显存的加持。

　　模型加一层，显存涨一分

　　在深度学习模型中，占用显存的总是那些特别大的张量，比如各层的权重矩阵、计算出来的张量（激活值）、反向传播需要的张量等。在视觉任务中，占据绝大多数的是中间计算出来的张量。随着模型变得更深更大，每一层的激活值张量都需要保留在显存中。

　　以 ResNet50 为例，在模型的训练中，前向传播中 50 层的计算结果都需要保存在显存中，以便让反向传播利用这些张量计算梯度。如果使用 ResNet108，需要的显存就会比 ResNet50 多出一倍多。显存的增加，带来的当然是模型效果的提升。另一方面，如果内存不够，许多工作也必将无法实现。

　　显存不够，写论文、打比赛屡遭掣肘

　　在实验室跑模型、写论文的过程中，显存不够用也是常有的事。一般实验室的显卡都是大家共用的，可能分配到每个人的手上已经所剩无几。甚至于，随着顶尖模型越来越大，所有人都没有足够的算力、显存去复现终极实验，更不用说超越其 SOTA 结果。

　　遇到这种情况，学生无非只有两种选择：向导师申请新的 GPU 资源，或者缩减模型做一个 Mini 版的实验。前者并不总是能够成功，后者则可能会有种种不完美。如果能用有限的显存跑顶尖的大模型，做实验、写论文都会变得更加简单。

　　此外，无论是在学校还是在公司打比赛，算力不够、显存不足都是常有的事。顶尖竞争者的模型结构可能相差无几，区别就在于谁的模型更大、更有能力去处理复杂的样本。更直观地说，排行榜领先者的模型也许就只差十几层，但也正是因为显存受限少了那十几层，有些模型才与冠军失之交臂。

　　显存：约束算法工程师的瓶颈

　　再举一个常见的例子，企业中的算法工程师拥有足够的算力，显存没那么重要。然而，只使用并行策略分担显存，还是可能会出现显存足够、但每张 GPU 的计算负载又不足的情况。

　　4 张 V100，显存占满，而 GPU 利用率很低。

　　即使是 V100 这样强大的算力，训练大模型时也很容易占满 16GB 显存。然而由于批量不够大，上图每张 V100 GPU 的利用率只有 20% 到 30%。只有继续增大每次迭代的数据吞吐量，才能增加 GPU 的利用率。

　　MegEngine：显存需要优化

　　其实对于深度学习从业者来说，日常应用中出现的情况远不止上面三种。做深度学习，不论是研究还是工程，时不时就会遇到显存问题。但这个问题优化起来又很复杂，需要利用大量的工程实现来缓解。显然，这样的优化应该由深度学习框架来完成。不过，在实际应用中不难发现，TensorFlow、PyTorch 似乎都没有提供完善的官方解决方案。

　　但如果把目光投向新生势力，情况可能就不一样了。在旷视开源深度学习框架 MegEngine 最近发布的 1.4 版本中，该框架首次引入了动态图显存优化技术，大大降低了显存占用问题。

　　具体来说，通过复现并优化 ICLR 2021 Spotlight 论文《Dynamic Tensor Rematerialization》（以下简称 DTR），MegEngine 实现了「用计算换取更多显存」。有了这项技术的加持，模型的显存占用大大降低，同样的硬件可以训练更大的模型、承载更大的 BatchSize。如此一来，学生的小显卡也能开始训练大模型，而工程师们的服务器也经得起更充分的应用。

　　原本需要 16GB 显存的模型，优化后使用的显存峰值就降到了 4GB。

　　MegEngine 这种显存优化技术，让 1060 这样的入门级显卡也能训练原本 2080Ti 才能加载得上的模型；而 11GB 显存的 2080Ti，更能挑战原本 32GB V100 才能训练的模型。要知道，V100 的价格可是 2080Ti 的 9 倍还多。

　　两行代码，显存「翻倍」

　　如要需要自己去优化显存， 可能 99% 的算法工程师都会放弃。最好的办法是告诉深度学习框架，这次训练就分配多少显存，剩下的就交给框架自己去优化。MegEngine 的动态图显存优化就是基于这一逻辑。

　　通过两行代码，框架可以全自动地完成显存优化，将所有优化逻辑与复杂的工程实现都隐藏在 MegEngine 内部。

　　如上图所示，在动态计算图中导入 DTR 显存优化模块，并配置显存释放阈值为 5GB。训练时，因为显存已经「翻倍」了，Batch Size 翻四倍也能装到 GPU 中。

　　显存扩增带来的收益

　　很多时候，提高显存的利用率，最显著的作用就是能训练更大的模型。从一定程度上来说，参数量越大就意味着效果越好；而批大小越大，梯度更新方向就越准确，模型性能也就越优异。MegEngine 开发团队做了很多实验，以确保提高显存利用率的同时训练是优质的。

　　最简单的验证方法就是不断增加批大小，看看显卡到底能坚持到什么程度。下面两张表分别展示了在 PyTorch 及 MegEngine 上加载或不加载动态图显存优化（DTR）技术的效果。

　　如果不使用动态图显存优化技术，PyTorch 上的模型一次训练迭代最多只能处理 64 个样本，MegEngine 能处理 100 个样本。只要加上 DTR，PyTorch 模型一次迭代就能处理 140 个样本，MegEngine 能尝试处理 300 个样本。

　　如果换算成模型大小，加上动态图显存优化技术的 MegEngine，在相同的 GPU 及批大小情况下，能高效训练增大近乎 5 倍的模型。

　　MegEngine 动态图显存优化技术

　　深度学习模型的显存占用一般分为权重矩阵、前向传播的中间张量、反向传播的梯度矩阵（Adam 优化器）三部分。

　　权重矩阵和梯度矩阵占的内存很难优化，各个模型基本上都有一个定值。前向传播的中间计算结果则不然：随着 Batch Size 的增加以及模型层和数量的增加，显存必然跟着增加。如果模型比较大，中间计算结果将占据最主要的显存。

　　如上图所示，在前向传播中（第一行从左到右），蓝色圆圈表示模型的中间计算结果开始占用显存。一直到前向传播完成，第一行完全变为蓝色圆圈，前面计算所占用的显存都不能释放。

　　等到反向传播开始（第二行从右到左），随着梯度的计算与完成应用，前向传播保留在显存中的张量才可以释放。

　　很明显，如果要降低显存占用，就要拿前向传播保存的中间计算结果开刀，这也正是 MegEngine 动态图显存优化的主要方向。

　　用计算换显存

　　对于动态计算图，最直接的方法就是用计算或内存换显存。因此，MegEngine 首先要决定到底使用哪种技术。

　　MegEngine 团队通过实验发现，用计算耗时远比交换耗时少。例如从显存中节省 612.5MB 空间，用带宽换显存要比用计算换显存慢了几十上百倍。

　　因此很明确，动态计算图中也应该使用梯度检查点技术，用计算换显存。

　　如下为梯度检查点技术原理示意，前向传播中第三个点为检查点，它会一直保存在显存中。第四个点在完成计算后即可释放显存，在反向传播中如果需要第四个点的值，可以从第三个点重新计算出第四个点的值。

　　虽然大致原理不难理解，但具体怎么做还是比较复杂的，MegEngine 团队借鉴了论文《Dynamic Tensor Rematerialization》，将其优化并实现到 MegEngine 中。

　　DTR，最前沿的显存优化技术

　　DTR 是一种完全动态的启发式策略，核心思想是当显存超过某个阈值时，动态地释放一些合适的张量，直到显存低于阈值。一般而言，释放张量的标准有三个：重新计算出该张量的开销越小越好；占用的显存越大越好；在显存中停留的时间越长越好。

　　除去从检查点恢复前向传播结果张量带来的主要开销，DTR 的额外开销在于寻找应该被释放的最优张量，即计算上图张量 t 的 f(t)值。为了降低这一部分的计算量，MegEngine 还采用了两种运行时优化：

• 不考虑小的张量，它们不加入候选集
• 每次在需要释放张量的时候，随机采样并遍历少部分张量，以节省计算开销

　　最难的是工程实现

　　虽然 DTR 看上去原理也不复杂，但真正的难题在于提高易用性，即将所有细节都隐藏到框架的底层，只为开发者提供最简单的接口。

　　在此就用一个最简单的计算例子，跟着框架演算一遍，看看 MegEngine 是如何利用动态图的计算历史恢复与释放张量的。

　　现在假设输入有 a 和 b 两个张量，并希望计算 a*b 与 a+b，但是显存最大只能保存三个张量。在黄框计算 c=a+b 时，显存还能保留张量 c，然而在下一步绿框计算 d=a*b 时只能先释放 c 才能保存 d。

　　不巧的是，下一步灰框需要获取黄框的计算结果，然而为了节省显存，c 已经被释放了。所以，MegEngine 现在需要做的是重新运行灰框的计算图，计算 c=a+b，并加载到显存中。显然，这样做必然需要释放 d 的显存。

　　这样一来，鉴于显存的限制，MegEngine 就会自动选择合适的张量释放，并在需要时重新计算。如果需要重新计算某个张量的结果，例如上图的 d，就需要具体的历史计算信息（在这里就是 a+b 这样的计算路径），与此同时还需要知道 a 和 b 这两个输入张量。

　　所有这样的历史计算信息都由 MegEngine 自动获取与保存，MegEngine 的工程师已经在底层用 C++ 处理完毕，用户完全不需要考虑。

　　struct ComputePath {

　　 std::shared_ptr

　　 op;

　　 SmallVector

　　 inputs;

　　 SmallVector

　　 outputs;

　　 double compute_time = 0;

　　} *producer;

　　SmallVector

　　 users;

　　size_t ref_cnt = 0;

　　以上为 MegEngine 底层用于追踪计算路径信息的结构体。其中 op 表示产生该张量的算子；inputs 和 outputs 分别表示这个算子需要的输入与输出张量；compute_time 表示该算子实际的运行时间。

　　实际上，在使用 MegEngine 的过程中，全都是用 Python 接口创建张量，只不过框架会对应追踪每个张量的具体信息。每当需要访问张量，不用考虑张量是否在显存中时，没有也能立刻恢复出来。所有这些复杂的工程化的操作与运算逻辑都隐藏在了 MegEngine C++ 底层。

　　Python 代码会翻译成 C++ 底层实现，C++ 代码会通过指针管理显卡内存中真正的张量（右图绿色部分）。

　　幸好这样的复杂操作不需要算法工程师完成，都交给 MegEngine 好了。

　　MegEngine 能做的事情远不止于此，只不过大多是像动态图显存优化这种技术一样，润物细无声地把用户的实际问题解决于无形。2020 年 3 月开源的 MegEngine 在以肉眼可见的速度快速成长，从静态计算图到动态计算图，再到持续提升的训练能力、移动端推理性能优化、动态显存优化…… 这也许就是开源的魅力。只有不断优化和创新，才能吸引和满足「挑剔」的开发者。MegEngine 下一个推出的功能会是什么？让我们拭目以待。