PyTorch AMP 怎样区分 autocast 与 GradScaler 的职责

我在给训练脚本加 AMP 时，最容易混淆的不是 API 名字，而是把“选择计算精度”和“保护梯度更新”当成了同一件事。PyTorch 官方文档把它们拆成两个可以独立使用的模块：autocast 负责在选定区域内为算子选择合适的精度，GradScaler 负责梯度缩放、反缩放和溢出处理。先分清这条边界，训练循环就不容易越写越乱。

官方地址：https://docs.pytorch.org/docs/stable/amp.html

一句话判断：autocast 管“前向怎么计算”，GradScaler 管“反向得到的梯度怎样安全地交给优化器”。float16 训练通常把两者配合起来；bfloat16 是否需要缩放，则要看设备、模型和数值表现，不能因为用了 autocast 就机械创建 scaler。

先把 autocast 和 GradScaler 拆成两条线

AMP 不是一个只能整包开启的开关。把训练过程画成三段更容易理解：前向算子产生输出，loss 进入反向传播，优化器读取梯度更新参数。

因此，看到 with torch.autocast(...) 时，先问“这段前向计算想怎样分配精度”；看到 scaler.scale(loss) 时，问的是“这次反向梯度是否需要数值保护”。两者可以组合，但不是互相替代。

autocast 应该包住哪些计算

常见训练循环把模型前向和 loss 放进 autocast 上下文，把 backward 放在上下文外。这样做的重点不是强制所有计算降精度，而是交给 AMP 的算子策略选择：适合低精度的计算可以更快，数值敏感的计算仍可能保留更高精度。

import torch

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

for inputs, targets in loader:
    inputs = inputs.to(device)
    targets = targets.to(device)
    optimizer.zero_grad(set_to_none=True)

    # autocast 只描述前向与 loss 的精度选择范围。
    with torch.autocast(device_type=device, dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

    # backward 不需要因为前向用了 autocast 就继续写在上下文里。
    loss.backward()
    optimizer.step()

这段代码只展示 autocast 的职责，适合用来理解推理或不需要梯度缩放的场景。它不等于完整的 float16 AMP 训练方案，因为 loss 仍然可能让很小的梯度在反向过程中下溢。

GradScaler 怎样接管梯度更新

需要梯度缩放时，先创建一个 scaler。每轮的顺序是：用 scaler 放大 loss，再对放大后的 loss 调用 backward；调用 scaler.step(optimizer) 时，scaler 会先检查并还原梯度，只有没有发现 inf 或 NaN 时才让优化器真正更新；最后用 scaler.update() 调整下一轮的缩放因子。

scaler = torch.amp.GradScaler("cuda")

for inputs, targets in loader:
    inputs = inputs.to(device)
    targets = targets.to(device)
    optimizer.zero_grad(set_to_none=True)

    # autocast 决定前向算子的计算精度。
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

    # scale 只放大反向使用的 loss，不改变模型参数的真实目标。
    scaler.scale(loss).backward()

    # step 会检查 inf/NaN；异常时跳过本次 optimizer.step()。
    scaler.step(optimizer)

    # 根据本轮是否稳定，调整后续迭代的 scale。
    scaler.update()

这里的“跳过更新”很重要：它不是训练循环报错，而是 scaler 在发现当前梯度不适合更新时保护参数。不要用 loss 的数值变化直接推断 optimizer 一定执行了更新；要把 scaler 的 step 视为一次带检查的更新请求。

需要梯度裁剪时，先 unscale 再裁剪

scaler.scale(loss).backward() 之后，参数的 .grad 还是缩放后的梯度。如果直接调用梯度裁剪，阈值也会被缩放，结果就不是原本想要的阈值。正确顺序是对目标优化器调用一次 unscale_，再按普通梯度执行裁剪，最后继续交给 scaler.step。

scaler.scale(loss).backward()

# 先把 optimizer 管理的梯度还原到真实尺度，再进行裁剪。
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# unscale_ 已经做过还原，step 会继续负责 inf/NaN 检查。
scaler.step(optimizer)
scaler.update()

同一个 optimizer 在一次 step 之间不要重复调用 unscale_。如果有多个 optimizer，需要分别对需要检查或修改的 optimizer 处理，但 scaler.update() 通常在这一轮所有 optimizer 都完成 step 后调用一次。

float16、bfloat16 和梯度累积不要混为一谈

autocast 与 GradScaler 是模块化的，所以组合方式取决于目标 dtype 和训练策略。

• float16：低精度范围较窄，梯度下溢风险更值得关注，常见做法是 autocast 配合 GradScaler。
• bfloat16：指数范围更接近 float32，很多场景下不需要同样的梯度缩放；但是否启用 scaler 仍要看设备支持、模型数值稳定性和实际训练表现。
• 梯度累积：一个有效 batch 的累积期间应保持相同 scale，不要在累积尚未完成时 unscale 或 update；完成一次有效 batch 后再 step 和 update。

for index, (inputs, targets) in enumerate(loader):
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        loss = loss_fn(model(inputs.to(device)), targets.to(device))
        # 除以累积步数，保持有效 batch 的梯度尺度一致。
        loss = loss / accumulate_steps

    # 累积期间保持 scaled gradients，不提前 unscale。
    scaler.scale(loss).backward()

    if (index + 1) % accumulate_steps == 0:
        # 有效 batch 完成后才提交更新并调整 scale。
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

如果只记住一条边界，可以这样记：autocast 是“算子级的前向计算策略”，GradScaler 是“优化器级的反向更新保护”。不要用扩大 batch、改变学习率或强行把所有层 cast 成 float16 来代替它们各自的职责。

我会怎样排查 AMP 训练循环

1. 先看 autocast 的作用域，确认它覆盖了 forward 和 loss，而不是把一大段无关的参数更新代码都包进去。
2. 再看 loss 是否通过 scaler.scale(loss).backward() 进入反向；如果使用 scaler，却仍直接 loss.backward()，缩放链路就断了。
3. 若做梯度裁剪、梯度范数读取或自定义梯度检查，确认先对对应 optimizer 调用了 unscale_。
4. 最后检查 step() 和 update() 的粒度：梯度累积要在有效 batch 完成后更新，多优化器要让每个 optimizer 都完成自己的 step。

这套排查顺序比“看到 loss 变成 NaN 就把所有地方改成 float32”更有针对性：先定位是算子精度选择、梯度缩放时机，还是优化器更新边界出了问题。

常见问题

只用 autocast 可以训练吗？

可以，autocast 和 GradScaler 是模块化的；但使用 float16 训练时，是否需要梯度缩放要根据模型和设备的数值表现判断。autocast 本身不会完成梯度溢出保护。

GradScaler 会把所有参数永久改成 float16 吗？

不会。GradScaler 只围绕 loss 和梯度更新做数值尺度处理，不负责决定模型参数或前向算子的 dtype。

为什么 scaler.step 没有让参数更新？

如果 scaler 在梯度中发现 inf 或 NaN，它可能跳过这次 optimizer.step()，随后通过 update 调整缩放因子。排查时要把这种保护分支和真正的程序异常区分开。

梯度裁剪能放在 scaler.step 之后吗？

通常不行。需要按真实梯度裁剪时，应在 scaler.step 前对对应 optimizer 调用 unscale_，再执行裁剪。

参考：PyTorch 官方 AMP 文档、Automatic Mixed Precision examples 与 AMP recipe。本文的图片均为围绕职责关系制作的原创静态说明图，不是软件截图或运行证据。