半精度训练的第一个坑是把"位数少"理解成"一定不准"。 真正要分清的是两件事: 指数位决定能表达多大/多小的数,也就是溢出和下溢; 尾数位决定相邻两个数之间的间隔,也就是精度。 FP16 和 BF16 都是 16 位,但性格完全不同。

一张图理解 FP16 vs BF16

为什么 FP16 会炸

import torch

x = torch.tensor([1e4], device="cuda", dtype=torch.float16)
print(x * x)       # tensor([inf], device='cuda:0', dtype=torch.float16)

g = torch.tensor([1e-8], device="cuda", dtype=torch.float16)
print(g)           # tensor([0.], device='cuda:0', dtype=torch.float16)  ← 梯度下溢

b = torch.tensor([1e4], device="cuda", dtype=torch.bfloat16)
print(b * b)       # 仍然能表达大数,但尾数更粗

autocast 不是把模型参数永久改成半精度。 它只是在一个上下文里拦截算子调用,根据算子的数值风险选择合适 dtype: 大矩阵乘和卷积适合降精度,归约和 softmax 更保守,有些算子会主动提升到 FP32。 这就是 AMP 的核心: Automatic Mixed Precision,自动混合,不是自动半精度。

现代 PyTorch 写法

import torch

device = "cuda"
model = MyModel().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

# CUDA + fp16 训练:需要 GradScaler
scaler = torch.amp.GradScaler(device)

for x, y in loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad(set_to_none=True)

    # 只包 forward + loss
    with torch.amp.autocast(device_type=device, dtype=torch.float16):
        logits = model(x)
        loss = torch.nn.functional.cross_entropy(logits, y)

    # scale 后 backward,避免 fp16 梯度下溢
    scaler.scale(loss).backward()

    # 如需 gradient clipping,必须先 unscale
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    scaler.step(optimizer)  # 若发现 inf/nan,自动跳过 optimizer.step()
    scaler.update()         # 动态调整 scale

BF16 写法更简单

# A100 / H100 上训练 LLM 的常见默认
for x, y in loader:
    optimizer.zero_grad(set_to_none=True)
    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
        loss = model(x).loss
    loss.backward()
    optimizer.step()

# BF16 的指数范围和 FP32 一样,通常不需要 GradScaler

需要手动接管的场景

FP16 最危险的是梯度下溢。解决思路非常朴素: loss 乘一个很大的 scale,反向传播得到的梯度也跟着变大, optimizer.step 前再除回去。 关键是 scale 不能太大,否则梯度又会 overflow。 GradScaler 做的就是这个动态平衡。

关键 API 顺序

常见错误

今天的练习目标不是把某个模型训到高分,而是建立一套固定的 AMP 验证手感: 同一个模型分别跑 FP32、FP16 AMP、BF16 AMP, 记录吞吐、显存峰值、loss 是否贴近。 以后你给任何训练脚本开 AMP,都照这个流程验。

# amp_bench.py — FP32 / FP16 AMP / BF16 AMP 对比骨架
import argparse, time, torch
import torch.nn as nn
import torch.nn.functional as F

class TinyMLP(nn.Module):
    def __init__(self, d=4096, layers=8):
        super().__init__()
        self.net = nn.Sequential(*[
            layer for _ in range(layers)
            for layer in (nn.Linear(d, d), nn.GELU())
        ], nn.Linear(d, 1000))

    def forward(self, x):
        return self.net(x)

def run(mode):
    device = "cuda"
    model = TinyMLP().to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    scaler = torch.amp.GradScaler(device, enabled=(mode == "fp16"))
    dtype = {"fp32": None, "fp16": torch.float16, "bf16": torch.bfloat16}[mode]

    torch.cuda.reset_peak_memory_stats()
    torch.cuda.synchronize()
    t0 = time.perf_counter()

    losses = []
    for step in range(80):
        x = torch.randn(32, 4096, device=device)
        y = torch.randint(0, 1000, (32,), device=device)
        opt.zero_grad(set_to_none=True)

        enabled = (dtype is not None)
        with torch.amp.autocast(device_type=device, dtype=dtype, enabled=enabled):
            loss = F.cross_entropy(model(x), y)

        if mode == "fp16":
            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()
        else:
            loss.backward()
            opt.step()

        losses.append(float(loss.detach()))

    torch.cuda.synchronize()
    dt = time.perf_counter() - t0
    mem = torch.cuda.max_memory_allocated() / 2**30
    print(f"{mode:4s}  time={dt:.2f}s  peak={mem:.2f}GiB  last_loss={losses[-1]:.4f}  scale={scaler.get_scale():.0f}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["fp32", "fp16", "bf16"], required=True)
    run(parser.parse_args().mode)

运行方式

python amp_bench.py --mode fp32
python amp_bench.py --mode fp16
python amp_bench.py --mode bf16

# 记录这三列:
# 1) time:吞吐收益
# 2) peak:显存峰值
# 3) last_loss / loss 曲线:数值是否明显偏离

你应该看到什么

生产检查项

FP8 不是"把 AMP 的 dtype 从 FP16 改成 FP8"这么简单。 8 位太小,必须额外维护 scale / amax history, 把真实数值范围映射进 FP8 可表示区间。 这也是为什么 FP8 在生产里通常通过 NVIDIA TransformerEngine 这类库使用, 而不是手写一堆 tensor.to(torch.float8_*)。

TransformerEngine 的使用形状

# 概念示例:真实项目需按 TransformerEngine 版本与硬件配置调整
import torch
import transformer_engine.pytorch as te
from transformer_engine.common import recipe

fp8_recipe = recipe.DelayedScaling(
    margin=0,
    fp8_format=recipe.Format.HYBRID,  # forward E4M3, backward E5M2
    amax_history_len=16,
    amax_compute_algo="max",
)

model = te.Linear(4096, 4096).cuda()

with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    y = model(x)
    loss = y.float().pow(2).mean()

1. 画出 FP32 / FP16 / BF16 的 sign、exponent、mantissa 位宽,解释为什么 FP16 需要 GradScaler 而 BF16 通常不需要。
2. 说明 autocast 为什么不能理解成"把所有 op 都转成 FP16",并举出三类算子的 dtype 策略:降精度、升 FP32、跟随最宽输入。
3. 写出 AMP 训练循环的 5 个关键步骤:scale、backward、unscale、step、update,并说明 gradient clipping 应该插在哪一步。
4. 如果 AMP 训练第 20 step 开始出现 nan,设计一个二分排查流程:看哪些日志、禁用哪些局部 autocast、如何和 FP32 baseline 比。
5. 解释 FP8 E4M3 / E5M2 的差异,以及为什么 FP8 必须配合 scale / amax history 才能训练。