06 · 混合精度

⏱ 26 分钟 ｜ ⭐ 性价比最高的一个优化

🎯 一句话

用一半的位数存数据，速度快一倍、显存省一半。 听起来像作弊，但它能work是因为 —— 深度学习需要的是"动态范围"，而不是"小数点后多少位"。

🔢 一、五种数字格式

   FP32   [符号 1][指数  8][尾数 23]     ← 传统标准
   TF32   [符号 1][指数  8][尾数 10]     ← A100 的默认，硬件自动用
   FP16   [符号 1][指数  5][尾数 10]     ⚠️ 范围窄
   BF16   [符号 1][指数  8][尾数  7]     ⭐ 范围同 FP32，精度低
   FP8    [符号 1][指数 4-5][尾数 2-3]   ← H100 起
        ↑ 决定【范围】      ↑ 决定【精度】

🔑 BF16 vs FP16 是这一章最重要的对比： BF16 牺牲精度换范围，FP16 牺牲范围换精度。 而深度学习更怕"溢出"（范围不够）而不是"不精确" —— 梯度本来就是噪声估计，差几个百分点无所谓； 但一旦变成 inf 或 0，训练直接崩。

所以现在的默认是 BF16，不是 FP16。

💰 二、为什么快：三个来源

   ① Tensor Core（最主要）⭐
      FP32:  19.5 TFLOPS
      BF16: 312   TFLOPS    ← 16 倍
      → 半精度【才能】用上 Tensor Core（第 2 章）

   ② 带宽减半
      同样的显存带宽，一次能搬两倍的元素
      → 对带宽瓶颈的操作（第 3 章那一大堆）直接快一倍 ⭐

   ③ 显存占用减半
      → 能开更大的 batch → 算术强度上升 → 又快一点

💡 注意第 ② 点：即使某个操作没用 Tensor Core，减半的数据量也让它变快。 这就是为什么混合精度对整个模型都有收益，而不只是矩阵乘。

⚙️ 三、"混合"精度：哪些用低精度，哪些不能

   ✅ 用 BF16/FP16：
   ├─ 矩阵乘、卷积（占 90%+ 的计算量）
   ├─ 激活值（占大部分显存）
   └─ 前向和反向的中间结果

   ❌ 必须保持 FP32：
   ├─ ⭐ 主权重副本（master weights）
   ├─ 优化器状态（Adam 的 m、v）
   ├─ Softmax / LayerNorm 的【累加部分】
   └─ 损失的求和归约

⭐ 为什么必须有 FP32 主权重副本

   典型情况：权重 = 1.0，学习率 × 梯度 = 0.0001

   FP16 下：1.0 的最小可表示间隔约 0.001
   → 1.0 + 0.0001 = 1.0     ⭐ 更新【完全丢失】💀

   → 所以：
     · 用 FP32 保存主权重，更新累加在 FP32 上
     · 每步把它转成 BF16 拿去做前向/反向

🔑 这就是第 3 章那个 "12 字节/参数" 的来源： FP32 主权重 4 + Adam 的 m 4 + v 4 = 12。 省显存的是激活，不是优化器状态。

🩹 四、Loss Scaling：FP16 专属的补丁

   ⚠️ 问题：梯度普遍很小（1e-7 量级很常见）
      FP16 的最小正规值约 6e-5
      → 大量梯度【下溢成 0】💀

   ✅ 解法：
   ① loss × S（比如 S = 65536）再反向
   ② 梯度全部被放大 S 倍 → 落进 FP16 的可表示范围
   ③ optimizer.step() 之前把梯度 ÷ S 还原

   ⭐ 动态 loss scaling：
     · 没溢出 → 每隔 N 步把 S 翻倍（试探更大的范围）
     · 溢出（出现 inf/nan）→ S 减半，【跳过这一步更新】

⭐ BF16 不需要 loss scaling —— 它的指数位和 FP32 一样宽， 梯度不会下溢。这是 BF16 在工程上最大的省心之处。

# FP16：需要 GradScaler
import torch
scaler = torch.amp.GradScaler()
with torch.autocast("cuda", dtype=torch.float16):
    loss = model(x)
scaler.scale(loss).backward()
scaler.unscale_(opt)                              # ⭐ 裁剪前必须先还原
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(opt); scaler.update()

# BF16：干净得多 ⭐
with torch.autocast("cuda", dtype=torch.bfloat16):
    loss = model(x)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()

⚠️ scaler.unscale_() 这一步经常被漏掉： 不还原就做梯度裁剪，等于按放大后的值裁剪，裁剪阈值完全失效。

⚠️ 五、五个真实的坑

💥 一个高频的困惑："我开了 AMP，为什么没变快？" 排查顺序：

① 矩阵维度是 8 的倍数吗？（4095 → 4096）
② 模型是不是太小？（小 kernel 的启动开销占主导）
③ 瓶颈是不是根本不在 GPU？（第 4 章的 dataloader 实验）
④ 用 profiler 确认 Tensor Core kernel 真的被调用了

🚀 六、FP8：H100 之后的下一步

   两种格式：
   E4M3（指数4尾数3）→ 精度稍好，用于【前向的激活和权重】
   E5M2（指数5尾数2）→ 范围稍大，用于【反向的梯度】

   ⭐ 收益：算力再翻一倍（H100: 1979 TFLOPS），带宽再减半

   ⚠️ 代价：动态范围极窄，必须【per-tensor scaling】
      → 每个张量维护一个缩放因子，实时统计并调整
      → 框架（Transformer Engine）帮你做，但要理解它在做什么

💡 FP8 目前的实际状态： 推理上已经相当成熟（尤其是权重和 KV Cache）； 训练上仍需谨慎 —— 大模型预训练用 FP8 需要仔细的逐层策略， 通常敏感层（第一层、最后一层、LayerNorm）仍保持更高精度。

🔗 和站内其他章的关系

✅ 检查点

1. BF16 和 FP16 的区别是什么？为什么现在默认用 BF16？
2. 混合精度快的三个来源是什么？哪个最主要？
3. 为什么即使不用 Tensor Core 的操作也会变快？
4. 为什么必须保留 FP32 主权重副本？举个具体数字。
5. 混合精度省的是哪部分显存？不省哪部分？
6. Loss scaling 解决什么问题？动态版本怎么工作？
7. 为什么 BF16 不需要 loss scaling？
8. "开了 AMP 但没变快"的四步排查顺序？
9. FP8 的两种格式各用在哪？它的主要代价是什么？

🛑 可以停在这里

⚡ 走神救援

用一半位数存数据：快一倍、省一半显存。能work是因为深度学习需要的是动态范围而不是小数点后多少位。⭐⭐BF16 vs FP16 是核心对比：BF16 指数 8 位（范围同 FP32）尾数 7 位，FP16 指数 5 位（最大仅 65504）尾数 10 位 → ⭐深度学习更怕溢出而不是不精确（梯度本就是噪声估计，但变 inf/0 就崩）→ 所以默认是 BF16。快的三个来源：①⭐Tensor Core（FP32 19.5 vs BF16 312 TFLOPS，半精度才能用上）②带宽减半（所以即使不用 Tensor Core 的操作也变快）③显存减半能开大 batch。"混合"的含义：矩阵乘/激活用 BF16，但⭐主权重副本、Adam 的 m/v、LayerNorm/Softmax 的累加必须 FP32——因为权重 1.0 加更新 0.0001 时，FP16 的最小间隔约 0.001，更新完全丢失；⭐这就是第 3 章"12 字节/参数"的来源，所以混合精度省的是激活不是优化器状态。Loss scaling（FP16 专属）：梯度普遍 1e-7 会下溢，先 loss×65536 再反向、step 前还原；动态版没溢出就翻倍、溢出就减半并跳过这步。⭐BF16 不需要 loss scaling（指数位和 FP32 一样宽）——工程上最省心之处。⚠️五个坑：维度没对齐到 8 的倍数（改了精度速度没变）、在 autocast 里手动 .half()、梯度裁剪前忘了 unscale_()（裁剪阈值完全失效）、FP16 中期 NaN（换 BF16）。"开了 AMP 没变快"排查：维度对齐 → 模型太小 → 瓶颈不在 GPU → profiler 确认 Tensor Core 被调用。FP8（H100）：E4M3 前向、E5M2 反向，算力再翻倍但动态范围极窄需 per-tensor scaling；推理已成熟，训练仍需谨慎（敏感层保持高精度）。

下一节 👉 07-算子融合与编译.md