PyTorch 源码解读之 torch.cuda.amp: 自动混合精度详解

PyTorch 源码解读之 torch.cuda.amp: 自动混合精度详解

Nvidia 在 Volta 架构中引入 Tensor Core 单元，来支持 FP32 和 FP16 混合精度计算。也在 2018 年提出一个 PyTorch 拓展 apex，来支持模型参数自动混合精度训练。自动混合精度（Automatic Mixed Precision, AMP)训练，是在训练一个数值精度 FP32 的模型，一部分算子的操作时，数值精度为 FP16，其余算子的操作精度是 FP32，而具体哪些算子用 FP16，哪些用 FP32，不需要用户关心，amp 自动给它们都安排好了。这样在不改变模型、不降低模型训练精度的前提下，可以缩短训练时间，降低存储需求，因而能支持更多的 batch size、更大模型和尺寸更大的输入进行训练。PyTorch 从 1.6 以后（在此之前 OpenMMLab 已经支持混合精度训练，即 Fp16OptimizerHook），开始原生支持 amp，即torch.cuda.amp module。2020 ECCV，英伟达官方做了一个 tutorial 推广 amp。从官方各种文档网页 claim 的结果来看，amp 在分类、检测、图像生成、3D CNNs、LSTM，以及 NLP 中机器翻译、语义识别等应用中，都在没有降低模型训练精度都前提下，加速了模型的训练速度。

本文是对torch.cuda.amp工作机制，和 module 中接口使用方法介绍，以及在算法角度上对 amp 不掉点原因进行分析，最后补充一点对 amp 存储消耗的解释。

1. 混合精度训练机制

torch.cuda.amp 给用户提供了较为方便的混合精度训练机制，“方便”体现在两个方面：

用户不需要手动对模型参数 dtype 转换，amp 会自动为算子选择合适的数值精度
对于反向传播的时候，FP16 的梯度数值溢出的问题，amp 提供了梯度 scaling 操作，而且在优化器更新参数前，会自动对梯度 unscaling，所以，对用于模型优化的超参数不会有任何影响
以上两点，分别是通过使用amp.autocast和amp.GradScaler来实现的。

autocast可以作为 Python 上下文管理器和装饰器来使用，用来指定脚本中某个区域、或者某些函数，按照自动混合精度来运行。混合精度在操作的时候，是先将 FP32 的模型的参数拷贝一份，拷贝的参数转换成 FP16，而 amp 规定了的 FP16 的算子（例如卷积、全连接），对 FP16 的数值进行操作；FP32 的算子（例如涉及 reduction 的算子，BatchNormalize，softmax…），输入和输出是 FP16，计算的精度是 FP32。在反向传播时，依然是混合精度计算，得到数值精度为 FP16 的梯度。最后，由于 GPU 中的 Tensor Core 天然支持 FP16 乘积的结果与 FP32 的累加（Tensor Core math），优化器的操作是利用 FP16 的梯度对 FP32 的参数进行更新。

对于 FP16 不可避免的问题就是：表示的范围较窄，如下图所示，大量非 0 梯度会遇到溢出问题。解决办法是：对梯度乘一个 2**N 的系数，称为 scale factor，把梯度 shift 到 FP16 的表示范围。

GradScaler的工作就是在反向传播前给 loss 乘一个 scale factor，所以之后反向传播得到的梯度都乘了相同的 scale factor。并且为了不影响学习率，在梯度更新前将梯度unscale。总结amp的基本训练流程：

1. 维护一个 FP32 数值精度模型的副本；

2. 在每个iteration。

   1. 拷贝并且转换成 FP16 模型；
   2. 前向传播（FP16 的模型参数）；
   3. loss 乘 scale factor s；
   4. 反向传播（FP16 的模型参数和参数梯度）；
   5. 参数梯度乘 1/s；
   6. 利用 FP16 的梯度更新 FP32 的模型参数。
      但是，这里会有一个问题，scale factor 应该如何选取？选一个常量显然是不合适的，因为 loss 和梯度的数值在变，scale factor 需要跟随 loss 动态变化。健康的 loss 是振荡中下降，因此GradScaler设计的 scale factor 每隔 N 个 iteration 乘一个大于 1 的系数，再 scale loss；并且每次更新前检查溢出问题（检查梯度中有没有inf和nan），如果有，scale factor 乘一个小于 1 的系数并跳过该 iteration 的参数更新环节，如果没有，就正常更新参数。动态更新 scale factor 是 amp 实际操作中的流程。总结 amp 动态 scale factor 的训练流程：

3. 维护一个 FP32 数值精度模型的副本；

4. 初始化 s；

5. 在每个 iteration + a 拷贝并且转换成FP16模型 + b 前向传播（FP16 的模型参数） + c loss 乘 scale factor s + d 反向传播（FP16 的模型参数和参数梯度） + e 检查有没有inf或者nan的参数梯度 + 如果有：降低 s，回到步骤a + f 参数梯度乘 1/s + g 利用 FP16 的梯度更新 FP32 的模型参数。

2. amp模块的API

用户使用混合精度训练基本操作：

#amp依赖Tensor core架构，所以model参数必须是cuda tensor类型
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
#GradScaler对象用来自动做梯度缩放
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        # 在autocast enable 区域运行forward
        with autocast():
            # model做一个FP16的副本，forward
            output = model(input)
            loss = loss_fn(output, target)
        # 用scaler，scale loss(FP16)，backward得到scaled的梯度(FP16)
        scaler.scale(loss).backward()
        # scaler 更新参数，会先自动unscale梯度
        # 如果有nan或inf，自动跳过
        scaler.step(optimizer)
        # scaler factor更新
        scaler.update()

2.1 autocast类

｀`autocast(enable=True)`` 可以作为上下文管理器和装饰器来使用，给算子自动安排按照 FP16 或者 FP32 的数值精度来操作。

2.1.1 autocast算子

PyTorch中，只有 CUDA 算子有资格被 autocast，而且只有 “out-of-place” 才可以被 autocast，例如：a.addmm(b, c)是可以被 autocast，但是a.addmm_(b, c)和a.addmm(b, c, out=d)不可以 autocast。amp autocast 成 FP16 的算子有：

autocast 成 FP32 的算子：
剩下没有列出的算子，像dot,add,cat…都是按数据中较大的数值精度，进行操作，即有 FP32 参与计算，就按 FP32，全是 FP16 参与计算，就是 FP16。

2.1.2 MisMatch error

作为上下文管理器使用时，混合精度计算 enable 区域得到的 FP16 数值精度的变量在 enable 区域外需要显式的转成 FP32：

# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand