神经网络训练无法显存_深度解析MegEngine亚线性显存优化技术

作者 | 旷视研究院 编辑 | Linda 基于梯度检查点的亚线性显存优化方法 [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)模型本身的参数。