185、AMP 自动混合精度训练的源码机制与 loss scale 策略:FP16 训练提速 30% 的工程实现
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