混合精度训练：bf16、fp32 Master 与 Loss Scaling

训练基础 里的 16 B/param 账本默认了「bf16 混合精度 + AdamW」。这一页回答它的前提：一个 step 里到底哪些数据用 bf16 跑、哪些必须用 fp32 记，以及为什么 bf16 不需要 loss scaling 而 fp16 需要。

1. 为什么要混合精度

纯 fp32 训练浪费：Tensor Core 的 bf16 吞吐是 fp32 的数倍，且权重、激活的显存和 HBM 流量直接减半。纯 bf16 训练不行：精度不够的地方真的存在——梯度归约、优化器的小步更新、softmax/loss 的归约，都会在低精度下损失到影响收敛。混合精度就是按「谁对误差敏感」分配精度。

2. 精度地图：一个 Step 里谁用什么精度

数据 / 算子	精度	原因
fp32 master Θ	fp32	更新的累加基准，必须精确
bf16 权重	bf16	前向反向实际读取，省一半显存和带宽
前向 GEMM	bf16 输入，fp32 累加	Tensor Core 内部累加器本来就是 fp32
激活	bf16	显存大头，对误差不敏感
loss / softmax 归约	fp32	归约链长，低精度误差累积
梯度 ∇W	bf16 存储，fp32 累积	通信与存储 bf16，本地累积 fp32
optimizer 状态 m、v	fp32	二阶统计量对精度敏感

3. Master Weights：为什么更新必须在 fp32 上

AdamW 的单步更新量级通常是权重的 $10^{-3}$ 到 $10^{-4}$ 以下。bf16 只有 8 位有效数字（尾数 7 位 + 隐含位），当更新量小于权重的舍入间隔时，$\Theta + \Delta\Theta$ 会被舍入回 $\Theta$——权重原地不动，训练悄悄停滞。fp32 有 24 位有效数字，相对精度高约 $2^{16}$ 倍，更新不会丢失。

所以标准做法是：常驻一份 fp32 master 权重，optimizer 在 fp32 上做更新，每个 step 开头（或前向前）把 master cast 成 bf16 供算子读取。这就是 16 B/param 账里那 4P 的 master 项，也是混合精度里「fp32 负责记」的核心。

4. fp16 的 Loss Scaling，bf16 为什么不用

两种 16 位格式的位分配不同：

格式	指数位	尾数位	动态范围	精度
fp16	5	10	~±65504，最小正数 ~6e-8	高（11 位有效数字）
bf16	8	7	与 fp32 同范围	低（8 位有效数字）

梯度里有大量小幅值（$< 2^{-24}$ 不罕见），fp16 的指数范围装不下它们，直接变成 0，梯度消失。fp16 训练的补救是 loss scaling：前向结束前把 loss 乘一个大常数 S，梯度整体放大 S 倍进入 fp16 的安全区，optimizer 更新前再除回 S；GradScaler 还会动态调整 S——出现 inf/nan 就跳步并把 S 减半。

bf16 指数位与 fp32 相同，梯度的幅值范围原样装得下，不需要 loss scaling，训练循环更省心。代价是尾数只有 7 位，权重的相对精度更粗——这正是需要 fp32 master 的另一面原因。现代 LLM 训练因此几乎一律选 bf16。

5. 落地：PyTorch AMP 骨架

scaler = torch.amp.GradScaler()          # 仅 fp16 需要；bf16 可省略

for batch, labels in dataloader:
    optimizer.zero_grad()
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        loss = model(batch, labels=labels).loss
    scaler.scale(loss).backward()        # bf16: loss.backward()
    scaler.step(optimizer)               # bf16: optimizer.step()
    scaler.update()

autocast 决定哪些算子进 bf16（GEMM、卷积）、哪些保持 fp32（softmax、归一化、loss 归约），维护着一张框架级的黑白名单。注意 AMP 默认并不维护 master weights——那是 DeepSpeed / Megatron / torch.distributed.fsdp 混合精度配置做的事；裸 AMP + AdamW 的直接 bf16 更新在小 lr 后期会有可见损失。

相关页面

• 训练基础：一个 Step 的执行图16 B/param 账本与执行图
• Activation Checkpointing激活那份显存怎么省
• ZeRO 与 FSDP静态三份怎么按 rank 切片
• 量化：从 Scale 到 W4A8推理侧的低精度合同对照