Mixed Precision Training.

混合精度训练（Mixed Precision Training）是指在训练深度学习模型的过程中，同时使用不同的数值精度（如半精度浮点数float16和单精度浮点数float32）进行计算，以提高计算速度并降低内存占用。混合精度训练的难点在于如何在获得性能提升的同时，确保训练过程的稳定性和模型的最终精度。

1. 深度学习中的数值精度

浮点数类型的数据在内存中以二进制方式存储，由三部分组成：

在计算机中，任何一个数都可以表示为二进制的科学计数法的形式（$1.xxx \times 2^n$或$0.xxx\times 2^n$）；其中指数位$n$决定了该数据类型的数值动态范围，指数位越多，可表示的数值范围越大；尾数位$x$决定了该数据类型的数值精度：尾数位越多，可表示的数值精度越高。

例如二进制的科学计数法$1.001001 \times 2^3$表示二进制$1001.001$，二进制$1001$对应整数$9$，二进制$0.001$对应小数$0.125$；所以该数据表示浮点数$9.125$。

根据不同的指数位和尾数位数量，常用的数值精度包括：

（1）FP32 (Single Precision or float32)

单精度浮点数FP32是深度学习中最常见的浮点数格式，遵循 IEEE 754 标准，使用 32 位来表示一个数字，包括 1 位符号位、8 位指数位和 23 位小数位。

FP32的计算规则如下：

\[(-1)^{\text{sign}} \times 2^{-126} \times \left(0+\frac{\text{fraction}}{2^{23}}\right)\] \[(-1)^{\text{sign}} \times 2^{\text{exponent}-127} \times \left(1+\frac{\text{fraction}}{2^{23}}\right)\]

FP32可以表示的数据范围为$[-3.4 \times 10^{38},3.4 \times 10^{38}]$，可以表示的精度为$10^{-6}$，即两个不同 FP32 数值的最小间隔为 $0.000001$。

（2）FP16 (Half Precision or float16)

FP16 同样遵循 IEEE 754 标准，但仅使用 $16$ 位，包括 $1$ 位符号位、$5$ 位指数位和 $10$ 位小数位。相比 FP32，FP16 减半了存储需求和计算带宽，非常适合于深度学习的训练和推理，尤其是在对内存和计算速度有严格要求的场景。

FP16的计算规则如下：

\[(-1)^{\text{sign}} \times 2^{-14} \times \left(0+\frac{\text{fraction}}{2^{10}}\right)\] \[(-1)^{\text{sign}} \times 2^{\text{exponent}-15} \times \left(1+\frac{\text{fraction}}{2^{10}}\right)\]

FP16可以表示的数据范围为$[-65504,66504]$，可以表示的精度为$2^{-24}/10^{-3}$，即两个不同 FP16 数值的最小间隔为 $0.001$。

（3）BF16 (Brain Floating Point 16 or Bfloat16)

BF16 是 Google 为加速机器学习而设计的非标准浮点格式，也使用 $16$ 位，但不同于 FP16，它拥有 $1$ 位符号位、$8$ 位指数位和 $7$ 位小数位，与 FP32 相同的指数范围。这样设计使得 BF16 在保持与 FP32 相似的数值范围的同时，牺牲一些精度以换取更高的计算效率。

在深度学习领域，数值范围的作用远高于数值精度；即数据类型的指数位的作用大于尾数位的作用。BF16 在不改变内存占用的情况下，用 $1/10$ 倍的精度换取了 $10^{34}$ 倍的值域。

2. 混合精度训练的原理

直接在低精度环境下训练会减少内存占用，但会引入精度损失。由于低精度数据的有效动态范围较小，在训练过程中会产生数据上溢(overflow)或下溢(underflow)的问题。

比如在FP16环境下，模型权重量级为$2^{-3}=0.125$，梯度量级为$2^{-14}=0.000061$，但是FP16在$[2^{-3},2^{-2}]$之间的固定间隔为$2^{-13}$，即比$2^{-3}$大的下一个数为$2^{-3}+2^{-13}$，因此本次梯度更新会失效，导致舍入误差。

论文Mixed Precision Training提出了同时使用高精度表示(如FP32)和低精度表示(如FP16/BF16)进行训练的混合精度训练方法，在减少占用内存的同时起到了加速训练的效果。

混合精度训练的主要思路是：权重从高精度表示转成低精度表示进行前向计算，得到损失之后用低精度表示计算梯度，再转成高精度表示更新到权重上。

（1）FP32 权重备份

为了避免舍入误差问题，采用FP32 权重备份，即备份一份高精度表示下的权重用于梯度更新。

尽管FP32额外拷贝一份权重，新增加了训练时存储的占用。 但在训练过程中，内存中占据大部分的是激活值，只要激活值基本都使用FP16/BF16来进行存储，最终模型的内存占用也能够减半。

（2）损失缩放 Loss Scale

在网络的训练后期，梯度（特别是激活函数平滑段的梯度）会特别小，低精度表示容易产生 underflow 现象。 为了解决梯度过小的问题，对计算出来的损失值进行scale，由于链式法则，损失上的scale也会作用在梯度上。scale过后的梯度会平移到FP16/BF16有效的存储范围内，在更新FP32的权重之前再把梯度unscale回去。

（3）改进算术方式

在进行大型累加时（batch-norm、softmax），为防止溢出，将低精度表示的矩阵相乘后和高精度表示的的矩阵进行加法运算，写入内存时再转回低精度表示可以大幅减少混合精度训练的精度损失。这是因为加法主要被内存带宽限制，对运算速度不敏感，因此不会降低训练速度。

3. 自动混合精度

自动混合精度（Automatic Mixed Precision, AMP）训练是一种深度学习训练技术，可以在训练过程中动态地选择使用浮点数的精度。对混合精度训练提供支持的GPU架构包括：

自动混合精度训练的基本思想是，根据计算的需求和成本，自动地在单精度（FP32）和半精度（FP16/BF16）之间切换。具体来说，AMP 训练会识别那些对精度要求不高的计算（如权重矩阵的乘法），并将这些计算转换为半精度计算，以减少梯度计算中的数值误差。而对于那些对精度要求较高的计算（如激活函数的计算），AMP训练仍然使用单精度计算，以保持模型的准确性和响应性。

（1）PyTorch中的AMP用法

使用 torch.cuda.amp 模块中的 autocast 类可以实现AMP训练。当进入 autocast 上下文后，支持 AMP 的 CUDA 算子会把 Tensor dtype 转换为 FP16/BF16，从而在不损失训练精度的情况下加快运算。autocast 应该仅包装网络的前向推理，包括损失计算，但不包括反向传播。

class torch.cuda.amp.autocast(
    device_type,       # 要使用的设备类型：cuda,cpu,xpu,hpu...
    dtype=None,        # 是否使用 torch.float16 或 torch.bfloat16
    enabled=True,      # 区域内是否应启用 autocast。默认值：True
    cache_enabled=None # 是否应启用 autocast 内部的权重缓存。默认值：True
)

在AMP的训练开始之前需要实例化一个 torch.cuda.amp.GradScaler 对象。通过放大损失的值，从而防止梯度的 underflow。

class torch.cuda.amp.GradScaler(
    init_scale=65536.0,   # 初始缩放因子
    growth_factor=2.0,    # 如果在 growth_interval 连续迭代中没有出现 inf/NaN 梯度，则放大缩放因子
    backoff_factor=0.5,   # 如果在迭代中出现 inf/NaN 梯度，则减小缩放因子
    growth_interval=2000, # 必须在没有 inf/NaN 梯度的连续迭代中发生的次数
    enabled=True          # 是否启用缩放
)

典型的混合精度训练流程如下：

# 创建模型和优化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# 在训练开始时创建一个 GradScaler 实例
scaler = torch.cuda.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()  # 清空历史梯度

        # 使用 autocast 运行前向推理
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)

        # 对缩放后的损失调用 backward() 以创建缩放后的梯度
        scaler.scale(loss).backward()

        # 安全地取消缩放梯度并调用 optimizer.step()
        scaler.step(optimizer)

        # 更新下一次迭代的缩放因子
        scaler.update()

⚠由于GradScaler对梯度进行了缩放，因此如果要引入额外的梯度操作，如梯度裁剪，需要在这些操作之前进行unscale：

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()  

        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)  # 在梯度裁剪之前进行unscale
        grad_norm = torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), 
                                                    max_norm=2.0,
                                                    norm_type=2,
                                                    error_if_nonfinite=False,)
        scaler.step(optimizer)
        scaler.update()