结论先行:所有GPU训练都应该开混合精度(AMP),显存省一半、速度快30%、精度几乎无损。不开AMP就是浪费。
一、AMP原理
前向计算用FP16省显存,梯度缩放防止下溢,权重更新仍用FP32保持稳定。PyTorch用torch.cuda.amp自动管理。
二、怎么用
scaler = GradScaler(),前向with autocast(),反向scaler.scale(loss).backward(),优化器scaler.step()。四行代码搞定。
三、注意事项
不支持FP16的操作会自动用FP32。loss scaling溢出时自动跳过step。BF16硬件(A100以上)不用loss scaling更简单。
四、收益有多大
显存占用减半,batch size可以翻倍。训练速度提升20%到40%,因为Tensor Core对FP16有硬件加速。
补充:A100及以上用BF16不需要GradScaler,直接autocast(dtype=torch.bfloat16)即可,更简单也更稳定。4090用FP16需要GradScaler。
混合精度训练的GPU实例可在GPU算力平台选择。




