185、AMP 自动混合精度训练的源码机制与 loss scale 策略：FP16 训练提速 30% 的工程实现
一、从一次诡异的 loss 爆炸说起
去年年底我在调 YOLOv11 的蒸馏实验，batch size 拉到 64 之后显存直接爆了。换成 FP16 训练，前 200 个 iteration 一切正常，loss 曲线平滑得像教科书。结果第 201 个 iteration 突然 loss 变成 nan，梯度直接炸穿。我盯着终端看了五分钟，心想这锅得甩给 AMP。
后来翻 PyTorch 源码才发现，不是 AMP 的锅，是我没理解 loss scaling 的触发时机。GradScaler 默认的 growth_interval 是 2000，但我的 batch 里偶尔会出现极端梯度，scale factor 还没来得及调整就已经炸了。这个问题在 YOLOv11 的检测头里尤其常见，因为分类和回归分支的梯度量级差了两个数量级。
二、AMP 到底在干什么：一个被过度包装的机制
很多人以为 AMP 就是“自动把 FP32 换成 FP16”，实际上 PyTorch 的 AMP 做了三件事：
- 自动插入 cast 操作：在 forward 过程中，把某些操作（比如卷积、线性层）的输入自动转成 FP16，但保持 BN、LayerNorm 这些对精度敏感的操作在 FP32。
- 维护一个动态的 loss scale
