【技术探秘】为什么 BF16 混合精度训练不需要 GradScaler？

在使用 PyTorch 进行混合精度训练加速时我们经常会碰到FP16和BF16这两个概念。但你有没有发现一个有趣的现象写代码时用 FP16 总是得小心翼翼地套上一个GradScaler而一旦换到 BF16就直接一把梭哈连GradScaler的影子都不用见了这是为什么呢今天我们就来扒一扒它们底层的秘密 核心差异动态范围决定命运虽然 FP16 和 BF16 都是 16 位的浮点数但它们把“技能点”点在了完全不同的地方FP16 (Float16)采用 5 位指数区10 位尾数区。它的特点是精度高但数值的表示范围比较窄大约 6e-5 到 65504。BF16 (Bfloat16)采用 8 位指数区7 位尾数区。它牺牲了一点点尾数精度但换取了极大的动态范围。它的表示范围和老大哥 FP32 是一模一样的大约 1e-38 到 3e38。 为什么 FP16 离不开 GradScaler在深度学习的反向传播过程中梯度的值往往非常非常小。因为 FP16 的下限不够低一旦计算出来的梯度小于 6e-5这个数字在 FP16 的世界里就会直接跌破下限变成 0。这就叫数值下溢 (Underflow)会导致你的模型完全无法更新学习。GradScaler就是为了拯救 FP16 而生的。它会在反向传播前偷偷把 Loss 乘以一个大倍数放大梯度让梯度膨胀到 FP16 能安全容纳的范围内。等到优化器要更新权重时它再除以这个倍数把真实的梯度还原回来。 为什么 BF16 可以“裸奔”BF16 可是专门为了深度学习量身定制的。得益于它那和 FP32 相同的 8 位指数区它的数值范围宽广得惊人。在反向传播时哪怕梯度再细微BF16 也能稳稳接住几乎不可能发生下溢既然没有下溢危机自然就不需要GradScaler这种补救措施来画蛇添足啦。代码写起来更清爽训练过程也更稳定不用担心缩放因子计算失误导致的 NaN (Not a Number) 报错。 混合精度训练实战 DemoFP16 vs BF16纸上得来终觉浅我们直接看 PyTorch 代码对比你会发现 BF16 的训练循环明显精简很多。importtorchimporttorch.nnasnn# # 0. 准备环境初始化模型和虚拟数据# modelnn.Linear(10,2).cuda()optimizertorch.optim.Adam(model.parameters(),lr1e-3)datatorch.randn(32,10).cuda()targettorch.randn(32,2).cuda()loss_fnnn.MSELoss()# # 场景 A使用 FP16 训练 (必须使用 GradScaler)# print(开始 FP16 训练循环...)# 必须额外初始化一个 GradScalerscalertorch.amp.GradScaler(cuda)optimizer.zero_grad()# 开启自动混合精度指定 dtype 为 float16withtorch.autocast(device_typecuda,dtypetorch.float16):outputmodel(data)lossloss_fn(output,target)# ⚠️ 关键不同需要放大 loss 再反向传播scaler.scale(loss).backward()# ⚠️ 关键不同使用 scaler 来更新参数scaler.step(optimizer)# ⚠️ 关键不同更新缩放因子scaler.update()# # 场景 B使用 BF16 训练 (代码极简无需 Scaler)# print(开始 BF16 训练循环...)# 注意你的硬件需要支持 BF16 (如 NVIDIA Ampere 架构及以上的 GPU)optimizer.zero_grad()# 开启自动混合精度直接指定 dtype 为 bfloat16withtorch.autocast(device_typecuda,dtypetorch.bfloat16):outputmodel(data)lossloss_fn(output,target)# ✅ 直接反向传播不需要任何 scaleloss.backward()# ✅ 直接更新参数optimizer.step()