梯度裁剪与混合精度

梯度裁剪与混合精度

本节摘要:上一节的优化器与调度假设梯度是 sane 的。通常不是。一个坏批能让梯度范数尖峰三个数量级。混合精度训练在损失侧引入 FP16 溢出,放大了这点。本节构建两条生产训练离不开的安全带:把梯度裁剪到配置的全局 L2 范数,以及带 autocast 与 GradScaler 的混合精度循环——检测 NaN 与 Inf、干净地跳过该步、并日志缩放因子供取证。建侧的问题是两者接线要正确:裁剪在 unscale 前则阈值在缩放梯度上;裁剪在 unscale 后则 GradScaler 的操作顺序要紧。正序是 scaler.scale(loss).backward() → scaler.unscale_(optimizer) → clip_grad_norm_ → scaler.step(optimizer) → scaler.update(),任何其他顺序都产出静默坏的循环。

对应原课程:Phase 19 · Lesson 45 · gradient-clipping-amp(原英文 phases/19-capstone-projects/45-gradient-clipping-amp/docs/en.md)。本节属「预训练/分布式」赛道第四节。

学习目标

阅读完本节,你应当能够:

1. 算所有参数梯度的全局 L2 范数,超阈值时就地裁剪。
2. 把训练步包进 autocast + GradScaler,使 FP16 前向与反向扛住溢出。
3. 在损失或梯度里检测 NaN 与 Inf,跳过优化器步,日志跳过。
4. 每步报 GradScaler 的缩放因子,使一长串跳过立即可见。

一、问题与直觉

昨天跑干净的训练,步 8217 的损失曲线垂直起飞。罪魁是单个批,梯度范数 4200,前峰的二十倍。不裁剪,优化器迈一步把模型前一小时学的全抹了。全局 L2 裁剪在范数 1.0,同批只贡献单位范数更新,损失留在趋势线上,运行存活。

混合精度把吞吐推 2~3 倍,代价是 FP16 指数范围窄。典型梯度在 FP16 溢出评估为 Inf,经后续层传播成 NaN,下一步优化器把所有权重设成 NaN。PyTorch 的 GradScaler 通过在反向前把损失乘大缩放因子、优化器步前把梯度除回同因子解决。若有梯度在 unscale 时是 Inf 或 NaN,缩放器跳过该步并把缩放因子减半;若前 N 步干净,加倍。整个训练里因子找到 FP16 范围允许的最高值。

全局 L2 范数

全局 L2 范数是拼接梯度向量的欧氏范数,非每参数范数。PyTorch 实现为 torch.nn.utils.clip_grad_norm_(parameters, max_norm),函数返回裁剪前范数,使本节能日志自然值与裁剪值两者——「我们每步都在裁剪」的诊断需要它。

autocast 与 GradScaler

torch.amp.autocast(device_type) 是选择性把合格运算(多数 matmul 类)跑 FP16 的上下文管理器。torch.amp.GradScaler(device_type) 是反向前缩放损失、优化器步前反缩放梯度的助手。两者一起设计;用一不用另一是配置错误。本节用 CPU autocast(因 CI 跑 CPU),CUDA 上把 device_type="cpu" 改 "cuda" 即可。CPU 上 GradScaler 是桩(CPU autocast 默认跑 BF16 不需损失缩放),但本节保留调用点使接线与 GPU 循环相同。

NaN 与 Inf 检测

检测在两处。第一,损失本身在反向前用 torch.isfinite 检查;Inf/NaN 损失不产出有用梯度,不进优化器就跳过。第二,scaler.unscale_(optimizer) 后扫描反缩放梯度 has_non_finite_grad(...),任何 Inf/NaN 当跳过。两检测合起来覆盖前向与反向两种失败模式。

二、从零实现

code/main.py 实现:

• compute_global_grad_norm(params):所有参数梯度的全局 L2 范数。
• has_non_finite_grad(params):扫描梯度是否有 Inf/NaN。
• train_step_amp(model, opt, scaler, x, y, max_norm):正序五步,返回日志。
• run_amp_demo(model, loader, steps):跑混合精度训练,每步打缩放因子、范数、是否跳步。

正序五步骨架:

def train_step_amp(model, opt, scaler, x, y, max_norm=1.0):
    with torch.amp.autocast("cpu"):                      # 前向 FP16/BF16
        logits = model(x)
        loss = F.cross_entropy(logits.view(-1, V), y.view(-1))
    if not torch.isfinite(loss):                          # 损失侧检测
        return {"skipped": True, "reason": "non-finite loss"}
    scaler.scale(loss).backward()                         # ① 缩放反向
    scaler.unscale_(opt)                                  # ② 反缩放梯度
    grad_norm = compute_global_grad_norm(model.parameters())
    if has_non_finite_grad(model.parameters()):           # 梯度侧检测
        scaler.update(); opt.zero_grad()
        return {"skipped": True, "reason": "non-finite grad", "grad_norm": grad_norm}
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)   # ③ 裁剪
    scaler.step(opt)                                      # ④ 优化器步(内部检查)
    scaler.update()                                       # ⑤ 更新缩放因子
    opt.zero_grad()
    return {"skipped": False, "loss": loss.item(),
            "grad_norm": grad_norm, "scale": scaler.get_scale()}
​

设计要点:正序是不可妥协的。clip 在 unscale 前则阈值在缩放梯度上(可能永远不裁剪);step 在 unscale 前则优化器看到缩放梯度(迈巨步发散)。scaler.step 内部会再次检查梯度有限性,有 Inf/NaN 就跳过实际优化器步但仍调 update 减半缩放因子。CPU 上 GradScaler 是桩,但保留五步使代码与 GPU 循环逐字相同——这是迁移性的关键。

三、框架对比

HuggingFace Trainer 把 autocast、GradScaler、梯度裁剪全包进 TrainingArguments(fp16=True、max_grad_norm=1.0)。PyTorch Lightning 的 Trainer(precision="16-mixed") 同理。Megatron-LM、DeepSpeed 在分布式上自己管缩放因子与跨 rank 的梯度裁剪(all-reduce 范数后裁剪)。本节手写让你看清五步正序、两处检测、缩放因子的动态调。BF16 是 FP16 的替代——指数范围与 FP32 相同、不需损失缩放,现代 GPU(H100/A100)BF16 更稳,本节的 GradScaler 在 BF16 下可省略,但 autocast 仍需。

四、可复用产物

code/main.py:compute_global_grad_norm、has_non_finite_grad、train_step_amp 均可复用。demo 在 CPU 上几秒跑完,故意注入一个 Inf 梯度,确认跳步路径被触发、缩放因子减半、日志正确。train_step_amp 的五步是生产训练步的标准骨架——换到任何 autocast 兼容模型上,循环原样工作。

五、练习

1. 顺序破坏:故意把 clip 放在 unscale 前,确认阈值在缩放梯度上、裁剪失效。
2. Inf 注入:某步注入 Inf 梯度,确认跳步、缩放因子减半、下一步恢复正常。
3. 裁剪频率:日志每步是否裁剪(自然范数 > max_norm),讨论裁剪频率高的训练是否学习率过大。
4. BF16 对比:把 autocast dtype 换 BF16,去掉 GradScaler,对比稳定性与精度。
5. 缩放因子轨迹:画缩放因子随训练的轨迹,确认它在干净训练里单调爬升。

本节要点回顾

1. 坏批抹训练:单批梯度尖峰三个数量级,不裁剪就抹掉前一小时学习。
2. 全局 L2 裁剪:拼接梯度向量的欧氏范数,阈值 1.0 是常见默认。
3. 混合精度 2~3 倍吞吐:FP16 范围窄,溢出成 Inf→NaN,GradScaler 解决。
4. 五步正序:scale→backward→unscale→clip→step→update,任何其他顺序静默坏。
5. 两处检测:损失侧 isfinite、梯度侧 has_non_finite,覆盖前向与反向失败。
6. 缩放因子动态:干净加倍、Inf 减半,找到 FP16 范围允许的最高值。

下一节,我们做「梯度累积」——用微批一次一个、缩放损失累积梯度,训练你买不起的有效批。