3、混合精度训练原理：前向传播、反向传播中的精度选择与计算图优化

混合精度训练并不是把整个网络“一刀切”地换成 FP16，而是在计算图的不同位置、不同阶段，按数值敏感度选择合适的数据类型。理解前向传播与反向传播中的精度选择逻辑，是正确使用 torch.cuda.amp 或 torch.amp 的前提。本节从计算图出发，说明哪些算子用低精度、哪些必须保留 FP32，以及计算图优化如何减少不必要的精度转换开销。

一、核心原理：计算图上的精度分区

一次训练迭代的计算图可以粗略分为三部分：前向计算、损失计算、反向梯度计算。混合精度的基本策略是：矩阵乘、卷积等计算密集型算子使用 FP16/BF16 执行，以利用 Tensor Core 的高吞吐；归约、归一化、损失、参数更新等数值敏感操作保留 FP32，避免溢出和舍入误差累积。

以典型的 Transformer 层为例，前向传播中 Linear、BMM、Conv 会以低精度执行，而 LayerNorm、Softmax、CrossEntropy 通常被列入 FP32 白名单或黑名单。反向传播中，梯度计算同样遵循“矩阵运算低精度、归约与累加高精度”的原则。参数本身始终以 FP32 保存，优化器更新也在 FP32 中完成，这就是常说的“FP32 主权重”机制。

关键结论：混合精度不是降低模型精度，而是把“计算”和“存储”解耦——计算用低精度提速，存储和更新用高精度保稳。
混合精度训练计算图：前向与反向的精度分区 前向传播 Forward 输入数据 → FP16 转换 cast to half，减少显存带宽 低精度计算区（Tensor Core） Linear / Conv / BMM / MatMul FP16 输入 → FP16 输出 FP32 敏感算子区 LayerNorm / Softmax / Loss FP16 输入 → FP32 计算 → FP16 输出 损失计算 Loss FP32 累加，输出标量 loss Loss Scaling loss × scale，防止梯度下溢 反向传播 Backward 梯度反传起点 dLoss/dOutput，FP16 梯度 低精度梯度计算区 dLinear / dConv / dBMM FP16 梯度 → FP16 梯度 FP32 梯度累加区 LayerNorm / Softmax 反向 FP16 梯度 → FP32 累加 梯度 Unscale grad / scale，恢复真实梯度 FP32 主权重更新 optimizer.step()，FP32 参数 计算图 反向遍历 绿色：低精度计算区 | 黄色：FP32 敏感算子区 | 紫色：FP32 主权重与更新 | 红色：Loss Scaling 与 Unscale

二、方法流程：一次 AMP 迭代的完整链路

以 PyTorch 的自动混合精度为例，一次标准迭代包含以下步骤。注意每一步的精度归属，这是排查 NaN 和梯度异常的关键线索。

scaler = torch.cuda.amp.GradScaler()

for input, target in dataloader:
    optimizer.zero_grad()

    # 前向：autocast 自动选择算子精度
    with torch.cuda.amp.autocast(dtype=torch.float16):
        output = model(input)
        loss = loss_fn(output, target)

    # 反向：scaler 先放大 loss，再反向传播
    scaler.scale(loss).backward()

    # 梯度裁剪前必须先 unscale
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    # 更新：scaler 检查 inf/nan，决定是否跳过
    scaler.step(optimizer)
    scaler.update()

这段代码中有三个容易出错的细节：第一，autocast 只应包裹前向计算，不应包裹 backward()；第二，clip_grad_norm_ 必须在 unscale_ 之后调用，否则裁剪的是被放大后的梯度；第三，scaler.step() 内部会检查梯度是否出现 inf/nan，若出现则跳过本次更新并降低 scale。

三、计算图优化：减少精度转换与冗余节点

混合精度训练的性能收益不仅来自 Tensor Core，还来自计算图层面的优化。框架在构建计算图时，会进行以下几类处理：

  • 算子融合：将 Conv + Bias + ReLU 等连续算子融合为一个低精度内核，减少中间结果的显存读写和精度转换。
  • 精度传播：根据算子白名单/黑名单，自动决定每个节点的输入输出精度，避免手动插入大量 .half() 和 .float()。
  • 梯度累加精度提升：在反向传播中，对梯度累加操作使用 FP32 累加器，防止 FP16 累加导致的精度丢失。
  • Cast 节点消除：当相邻算子精度一致时，消除不必要的 cast 节点，降低 kernel 启动开销。
在 NVIDIA A100 及更新架构上，BF16 的动态范围与 FP32 相同，通常不需要 Loss Scaling，计算图也更简洁。但在 V100 等仅支持 FP16 的硬件上，Loss Scaling 仍是必需环节。

四、案例与场景：Transformer 训练中的精度选择

以一个 12 层 Transformer 语言模型为例，开启 AMP 后各模块的精度分配如下表所示。该表可作为排查精度问题的对照清单。

模块/操作 前向精度 反向精度 说明
Embedding 查表 FP32 → FP16 FP16 梯度 查表本身不计算，输出转 FP16
Q/K/V Linear FP16 FP16 Tensor Core 加速重点
Attention Score FP16 MatMul FP16 Softmax 前需转 FP32
Softmax FP32 FP32 数值敏感，防溢出
LayerNorm FP32 FP32 均值/方差需高精度
FFN Linear FP16 FP16 计算密集，收益最大
CrossEntropy Loss FP32 FP32 log-sum-exp 需高精度
参数更新 — FP32 主权重始终 FP32

五、常见误区与复盘

误区一：认为开启 AMP 后所有计算都是 FP16。实际上框架会根据算子白名单自动回退到 FP32，强行全局 .half() 反而会导致 LayerNorm 溢出和 loss 爆炸。
误区二：在 autocast 区域内手动调用 .half()。这会破坏 autocast 的精度决策逻辑，导致部分算子收到非预期类型。应让 autocast 自动管理，仅在必要时用 torch.cuda.amp.custom_fwd/custom_bwd 自定义。
误区三：忽略梯度裁剪与 unscale 的顺序。先裁剪后 unscale 会裁剪放大后的梯度，等效于缩小了裁剪阈值，导致梯度被过度压制，训练变慢甚至不收敛。

复盘时建议重点检查三件事：一是 loss 曲线是否出现周期性尖峰（可能是 scale 调整过于激进）；二是梯度范数是否异常小（可能是 unscale 顺序错误）；三是显存占用是否真的下降（若未下降，可能是模型中存在大量 FP32 常驻张量）。掌握前向与反向的精度分区逻辑，才能让混合精度训练既快又稳。

← 上一章 📖 返回目录 下一章 →