神经网络减少GPU显存占用的技术

来源于阿里云的PAI平台使用的技术

1）激活检查点（Activation Checkpoint）

在神经网络中间设置若干个检查点(checkpoint)，检查点以外的中间结果全部舍弃，反向传播求导数的时间，需要某个中间结果就从最近的检查点开始计算，这样既节省了显存，又避免了从头计算的繁琐过程。

2）梯度累积 (Gradient Accumulation)

以batch_size=16为例，可以每次算16个样本的平均梯度，然后缓存累加起来，算够了4次之后，然后把总梯度除以4，然后才执行参数更新，这个效果等价于batch_size=64。这是一种有效的增加Batch Size的方法。通过该策略可以将每个step的batch size扩充到很大，结合LAMB优化器会提升收敛速度。

3）混合精度训练（Mixed Precision Training）

采用混合精度训练的好处主要有以下两点：

减少显存占用，由于FP16的内存占用只有FP32的一半，自然地就可以帮助训练过程节省一半的显存空间。
加快训练和推断的计算，FP16除了能节约内存，还能同时节省模型的训练时间。具体原理如下图所示，核心是在反向传播参数更新的时候需要维护一个FP32的备份来避免舍入误差，另外会通过Loss Scaling来缓解溢出错误。
4）即时编译JIT

当PyTorch在执行一系列element-wise的Tensor操作时，底层Kernel的实现需要反复地读写访存，但是只执行少量的计算，其中大部分时间开销并不在计算上，而在访存读写上。比如，实现一个带有N个元素的Tensor的乘/加法Kernel，需要N次加计算，2N次读和N次写访存操作。我们称计算少, 访存次数多的Kernel为访存Bound。为了避免这种反复的读写，以及降低Kernel Launch的开销，可以采用Kernel Fusion。访存Bound的Kernel Fusion的核心原理是通过访存的局部性原理，将多个element-wise的Kernel自动合并成一个Kernel，避免中间结果写到内存上，来提高访存的利用率；同时由于多个Kernel合并成一个Kernel，Kernel launch开销也减少到了1次。

5）3D并行

3D并行策略指的是：数据并行，模型并行，流水并行三种策略的混合运用，以达到快速训练百亿/千亿量级模型的目的。该项技术最早由DeepSpeed团队研发，可以加速大模型的训练。

6）CPU Offload

反向传播不在GPU上计算，而是在CPU上算，其中用到的中间变量全部存储在内存中，这样可以节省下GPU的显存占用，用时间换空间，以便能放到更大尺寸的模型。

7）Zero显存优化器

ZeRO（The Zero Redundancy Optimizer）是一种用于大规模分布式深度学习的新型内存优化技术。ZeRO具有三个主要的优化阶段：

优化器状态分区(Pos) ：减少了4倍的内存，通信容量与数据并行性相同；
增加梯度分区(Pos+g) ：8x内存减少，通信容量与数据并行性相同；
增加参数分区(Pos+g+p) ：内存减少与数据并行度和复杂度成线性关系。