Pytorch 混合精度训练（Automatic Mixed Precision）原理解析

Pytorch 混合精度训练（Automatic Mixed Precision）原理解析

1. Overview

默认情况下，大多数深度学习框架（比如 pytorch）都采用 32 位浮点算法进行训练。Automatic Mixed Precision（AMP, 自动混合精度）可以在神经网络训练过程中，针对不同的层，采用不同的数据精度进行计算，从而实现节省显存和加快速度的目的。

Pytorch AMP 是从 1.6.0 版本开始的，在此之前借助 NVIDIA 的 apex 可以实现 amp 功能。Pytorch 的 AMP 其实是从 apex 简化而来的，和 apex 的 O1 相当。

AMP 里面的 Mixed 的方式很多，但是这里仅仅讨论 Fp16 和 Fp32 的混合。另外 pytorch 支持 cpu gpu 等不同设备上的 AMP，这里仅仅讨论 GPU 的 AMP（也就是 torch.cuda.amp）。

AMP 的使用非常简单，这里重点介绍 AMP 的原理。如果你只是想知道怎么使用 AMP 可以移步 官方文档。

2. 原理

2.1 Fp16 V.S. Fp32

通过上面的对比不难发现，Fp16 相对于 Fp32 主要有以下优势：

• 对于 memory-limited 算子， Fp16 相对于 Fp32 可以减小一半的访存，能提升算子性能
• 减小模型显存占用，相同条件下，可以使用更大的 batch size，或者是更复杂的模型
• 对于计算密集的算子，比如 linear， conv 等通过 Tensor Cores 可以进行加速

而劣势主要在于：

• 表示的动态区间更小
• 精度不如 Fp32

2.2 FP16 替换 FP32 的问题

针对 2.1 中的两点劣势，主要有两方面的问题：

• 溢出，由于 FP16 的范围更小，超出范围的时候就会溢出，具体来说包括上溢和下溢
• 舍入误差

溢出很好理解，就是超过其最大最小值，这里主要看下舍入误差。

如上面的例子，在模型训练的过程中，更新梯度的时候如果梯度较小（实际上就是很小）有可能会被舍弃，也就是没法起到更新的作用，显然这是不希望的。

2.3 解决方案

理论基础是 2018 年英伟达和百度在 Mixed Precision Training 中提出的，论文比较短，感兴趣的可以看看。

2.3.1 FP32 MASTER COPY OF WEIGHTS

训练的时候保留一份 FP32 的权重，最后更新的时候更新到 FP32 的权重上面，这样可以解决上面的舍入误差的问题。

2.3.2 LOSS SCALING

如下图，激活值的梯度只占用了 FP16 表示范围的偏左的一部分，而且还有相当一部分会发生下溢。因此如果对梯度进行放大，可以减少下溢。需要注意的是，一旦发生下溢将无法挽回，因为下溢后会变成 0，而 0 是一个合法的数，但是上溢后会出现 NaN 或 INF,这个是可以发现。

2.3.3 Operation Categories

作者对 op 进行了分类，哪些算子适合 FP16，就是 FP16 对精度没有影响或者影响很小，哪些只能用 FP32 计算。

3. Pytorch 中的 AMP

Pytorch 的 AMP 使用起来非常简单，原因就在于其实现了自动化，所以不需要用户设置太多参数。

3.1 autocast

autocast 主要就是实现了对算子进行分类，不同的算子用不同的精度进行计算。下面是具体的分类名单。

autocast 是通过修改 dispatch 的分发路径来实现对不同的算子的精度控制。

3.2 Scaler 调整策略

torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, 
growth_interval=2000, enabled=True)

- init_scale: 起始 scale
- growth_factor: update scale 的因子
- backoff_factor:如果出现 infs/NaNs，回退 scale 的因子
- growth_interval: 如果没有出现 infs/NaNs的，迭代 growth_interval 次后更新 scale
- Enable：控制是否对 grad 进行 scale```

Scaler 的大小在每次迭代中动态的估计，为了尽可能的减少梯度 underflow，scaler 应该更大；但是如果太大的话，半精度浮点型的 tensor 又容易 overflow（变成 inf 或者 NaN ）。所以动态估计的原理就是在不出现 inf 或者 NaN 梯度值的情况下尽可能的增大 scaler 的值——在每次scaler.step(optimizer) 中，都会检查是否有 inf 或 NaN 的梯度出现：

• 如果出现了 inf 或者 NaN，scaler.step(optimizer) 会忽略此次的权重更新（optimizer.step() )，并且将 scaler 的大小缩小（乘上 backoff_factor）；

• 如果没有出现 inf 或者 NaN，那么权重正常更新，并且当连续多次（growth_interval 指定）没有出现 inf 或者 NaN，则 scaler.update() 会将scaler 的大小增加（乘上 growth_factor）。

这里需要注意的是为什么要不停的调整 scaler?

• Scaler 的目标是尽量减少梯度 underflow， 因此应该尽可能的让 scaler 大，但是又要保证不会出现 overflow。
• 在训练中正常的 loss 通常是震荡中下降，所以调整策略也是隔 growth_interval 个 iteration 放大 growth_factor（大于1） 倍。

3.3 AMP 的过程

4 使用 AMP 的条件

显然首先你的硬件需要支持 Fp16。如果没有 FP16 Tensor Cores 的话，上面 2.1 中的优势除了第三点不存在外，其他的两点依然存在。

参考

• https://zhuanlan.zhihu.com/p/75753718)
• https://zhuanlan.zhihu.com/p/348554267
• https://pytorch.org/docs/stable/_modules/torch/cuda/amp/grad_scaler.html#GradScaler
• https://pytorch.org/docs/master/amp.html#autocast-op-reference
• https://developer.nvidia.com/automatic-mixed-precision