机器之心原创

作者：思

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）技术的效果。