PyTorch 混合精度训练:速度提升之前,先监控数值稳定性

📅 2026/7/4 21:29:35
PyTorch 混合精度训练:速度提升之前,先监控数值稳定性
PyTorch 混合精度训练速度提升之前先监控数值稳定性一、混合精度不是无风险加速开关PyTorch 混合精度训练可以降低显存占用提高吞吐。很多模型打开 AMP 后速度会提升。但混合精度改变了数值计算路径可能带来梯度下溢、loss 不稳定和精度退化。因此混合精度实验不能只看每秒样本数。必须同时监控 loss 曲线、梯度范围、验证集指标和最终收敛。速度提升如果换来不可控的数值误差就不算优化成功。二、训练链路要监控关键量flowchart TD A[前向计算] -- B[autocast] B -- C[loss] C -- D[GradScaler] D -- E[反向传播] E -- F[参数更新] F -- G[指标监控]AMP 通常使用 autocast 和 GradScaler。autocast 控制部分算子使用低精度GradScaler 通过缩放 loss 减少梯度下溢。两者配合能提高稳定性但不是保证。监控项至少包括 loss、scale 值、是否发生 overflow、梯度范数和验证指标。如果 scale 频繁下降说明数值稳定性可能存在问题。三、代码要保留可对照实验scaler torch.cuda.amp.GradScaler() for batch in loader: optimizer.zero_grad(set_to_noneTrue) with torch.cuda.amp.autocast(): output model(batch[input]) loss criterion(output, batch[label]) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度实验要和 FP32 基线对照。相同数据、相同随机种子、相同训练步数比较吞吐、显存、loss 和验证指标。没有基线无法判断收益和风险。amp_report: throughput_gain: 1.34 memory_reduction: 0.28 val_metric_delta: -0.001 overflow_steps: 7报告里不要只写“速度提升 34%”。验证指标变化、overflow 次数和稳定性也要写。这样别人才能判断是否适合复用。四、部分算子可能需要排除不是所有计算都适合低精度。归一化、softmax、某些损失函数和小数值累积可能更敏感。若发现训练不稳定可以局部关闭 autocast 或强制使用 FP32。还要关注硬件差异。不同 GPU 对 FP16、BF16 的支持不同。BF16 动态范围更大某些场景更稳定但不是所有硬件都高效支持。实验报告应写明硬件环境。梯度裁剪也要重新验证。FP32 下稳定的裁剪阈值在混合精度下未必仍然合适。若出现 loss spike可以同时观察梯度范数和 scaler 状态判断是优化器问题还是数值溢出问题。分布式训练中AMP 还会和梯度同步交互。某个 rank 发生 overflow所有 rank 的更新都应保持一致。框架通常会处理这些细节但实验报告仍应记录 world size、通信后端和 AMP 配置。保存 checkpoint 时也要记录 scaler 状态。训练中断后恢复如果只恢复模型和优化器不恢复 scaler前几个 step 的数值行为可能和原实验不同。可复现实验需要保存完整训练状态。最后混合精度收益要按模型类型报告。CNN、Transformer、推荐模型和小型 MLP 的瓶颈不同AMP 收益不能简单迁移。五、总结PyTorch 混合精度训练要同时评估吞吐、显存、loss 稳定性、overflow 和验证指标。AMP 不是无风险开关。训练加速只有在数值稳定和最终指标不退化的前提下才是真正的优化。