深度学习框架的亚线性显存优化技术解析

基于梯度检查点的亚线性显存优化方法[1]由于较高的计算/显存性价比受到关注。MegEngine经过工程扩展和优化，发展出一套行之有效的加强版亚线性显存优化技术，既可在计算存储资源受限的条件下，轻松训练更深的模型，又可使用更大batch size，进一步提升模型性能，稳定batchwise算子。使用MegEngine训练ResNet18/ResNet50，显存占用分别最高降低23%/40%；在更大的Bert模型上，降幅更是高达75%，而额外的计算开销几乎不变。该技术已在MegEngine开源，欢迎大家上手使用： https://github.com/MegEngine。

深度神经网络训练是一件复杂的事情，它体现为模型的时间复杂度和空间复杂度，分别对应着计算和内存；而训练时内存占用问题是漂浮在深度学习社区上空的一块乌云，如何拨云见日，最大降低神经网络训练的内存占用，是一个绕不开的课题。

GPU显卡等硬件为深度学习提供了必需的算力，但硬件自身有限的存储，限制了可训练模型的尺寸，尤其是大型深度网络，由此诞生出一系列相关技术，比如亚线性显存优化、梯度累加、混合精度训练、分布式训练，进行GPU显存优化。

其中，亚线性显存优化方法[1]由于较高的计算/显存性价比备受关注；旷视基于此，经过工程扩展和优化，发展出加强版的MegEngine亚线性显存优化技术，轻松把大模型甚至超大模型装进显存，也可以毫无压力使用大batch训练模型。

这里将围绕着深度学习框架MegEngine亚线性显存优化技术的工程实现和实验数据，从技术背景、原理、使用、展望等多个方面进行首次深入解读。

背景

在深度学习领域中，随着训练数据的增加，需要相应增加模型的尺寸和复杂度，进行模型「扩容」；而ResNet [2] 等技术的出现在算法层面扫清了训练深度模型的障碍。不断增加的数据和持续创新的算法给深度学习框架带来了新挑战，能否在模型训练时有效利用有限的计算存储资源，尤其是减少GPU显存占用，是评估深度学习框架性能的重要指标。

在计算存储资源一定的情况下，深度学习框架有几种降低显存占用的常用方法，其示例如下：

• 通过合适的梯度定义，让算子的梯度计算不再依赖于前向计算作为输入，从而in-place地完成算子的前向计算，比如Sigmoid、Relu等；
• 在生命周期没有重叠的算子之间共享显存；
• 通过额外的计算减少显存占用，比如利用梯度检查点重新计算中间结果的亚线性显存优化方法[1]；
• 通过额外的数据传输减少显存占用，比如把暂时不用的数据从GPU交换到CPU，需要时再从CPU交换回来。

上述显存优化技术在MegEngine中皆有不同程度的实现，这里重点讨论基于梯度检查点的亚线性显存优化技术。

原理

一个神经网络模型所占用的显存空间大体分为两个方面：1）模型本身的参数，2）模型训练临时占用的空间，包括参数的梯度、特征图等。其中最大占比是 2）中以特征图形式存在的中间结果，比如，从示例[1]可知，根据实现的不同，从70%到90%以上的显存用来存储特征图。这里的训练过程又可分为前向计算，反向计算和优化三个方面，其中前向计算的中间结果最占显存，