PyTorch 怎样用 autocast 控制混合精度范围

如果你想让 PyTorch 的混合精度只覆盖模型前向和损失计算，核心写法就是把这两部分放进 torch.autocast，然后在上下文外执行反向传播和参数更新。这样既能让 autocast 自动为适合的算子选择低精度，也能把需要稳定精度的局部计算单独切回 float32。

官方地址：https://pytorch.org/

记住一个边界：autocast 管前向计算的精度选择，GradScaler 管训练梯度的缩放；backward() 通常放在 autocast 外，推理则只需要 autocast，不需要 GradScaler。

• 训练：autocast 包住前向与 loss，退出后再 backward、step、update。
• 推理：只包住模型前向和必要的后处理，不创建 GradScaler。
• 局部稳定性：在外层 autocast 内用 enabled=False 创建 float32 子域。
• 判断依据：显式指定 dtype 的算子、原地操作和 out= 形式不会按普通可自动转换算子处理。

先把 autocast 的边界放对

在一个常见的训练回合里，模型前向和损失计算是 autocast 最合适的作用域。它会根据算子类型选择低精度或 float32，而不是把整个 Python 函数强行转换成同一种 dtype。反向传播依赖前向期间形成的计算图，因此通常在退出 autocast 后调用。

import torch

# autocast 只包住前向和损失计算，退出后再做反向传播
with torch.autocast(device_type="cuda", dtype=torch.float16):
    logits = model(inputs)
    loss = loss_fn(logits, targets)

# 训练更新不放进 autocast，保持职责边界清楚
loss.backward()
optimizer.step()
optimizer.zero_grad(set_to_none=True)

这里的关键不是把上下文写得越大越好，而是让它围绕“需要自动选择算子精度的前向区域”。把数据读取、优化器更新或无关的日志逻辑塞进上下文，不会让训练更可靠，反而会让 dtype 归属变得难以判断。

设备和任务不同，dtype 也要明确

使用新式接口时，把设备类型写在 torch.autocast 参数中更直观。CUDA 训练常见的是 float16，CPU 场景常见的是 bfloat16；实际选择仍要结合设备支持和模型的数值稳定性。

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
inputs = inputs.to(device)

# CUDA 示例：只让当前前向区域采用混合精度
with torch.autocast(device_type="cuda", dtype=torch.float16):
    logits = model(inputs)

# CPU 示例：CPU 设备使用 bfloat16 时，单独写清作用域
with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    logits = model(inputs)

# 不要把两个设备示例同时放进同一个真实训练回合
print(logits.dtype)  # 用于观察示例张量的 dtype，不代表所有算子都相同

上面的两个上下文是互斥示例。实际代码应根据输入张量和模型所在设备选择其中一个，避免把设备字符串、输入位置和 autocast 配置写成三套互不相干的状态。

需要稳定计算时，用嵌套作用域切回 float32

有些敏感计算需要 float32 的动态范围。此时不用取消整个外层 autocast，而是在局部建立 enabled=False 的子域，并把传入该子域的浮点输入显式转为 float32。这样边界清楚，其他前向算子仍可继续使用外层的自动转换。

import torch

with torch.autocast(device_type="cuda", dtype=torch.float16):
    features = encoder(inputs)

    # 局部关闭 autocast，敏感算子使用 float32 输入和计算
    with torch.autocast(device_type="cuda", enabled=False):
        stable_features = features.float()
        stable_score = sensitive_op(stable_features)

    # 子域结束后回到外层 autocast 的前向区域
    loss = loss_fn(stable_score, targets)

# backward 仍在 autocast 外部执行
loss.backward()

注意，关闭 autocast 不会自动把已经得到的张量变成 float32，所以 features.float() 是这个边界中的重要动作。子域只应覆盖真正需要稳定精度的计算，范围过大就会失去混合精度的收益。

训练时怎样接上 GradScaler

对使用 float16 的训练，autocast 与 torch.amp.GradScaler 是两个独立组件：前者负责前向算子的 dtype 选择，后者负责缩放 loss 和梯度，降低小梯度下溢的风险。一个最小训练骨架如下。

import torch

# scaler 在一次完整训练过程中复用，不要每个 batch 都重新创建
scaler = torch.amp.GradScaler("cuda")

for inputs, targets in train_loader:
    inputs = inputs.cuda(non_blocking=True)
    targets = targets.cuda(non_blocking=True)
    optimizer.zero_grad(set_to_none=True)

    # 前向与 loss 在 autocast 内部计算
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        logits = model(inputs)
        loss = loss_fn(logits, targets)

    # 缩放、反向、更新各自承担不同职责
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

如果要在更新前裁剪梯度，先对对应优化器调用 scaler.unscale_(optimizer)，再执行裁剪；同一个优化器在一次 step 间隔内不要重复 unscale。梯度累积时，也要等一个有效 batch 的梯度积累完成后再统一 unscale、step 和 update。

推理、原地操作和编译场景的几个边界

推理或评估只需要用 autocast 包住前向，不需要 GradScaler。训练与推理可以共享模型，但作用域和更新步骤不能混写。

import torch

model.eval()
with torch.inference_mode():
    # 推理没有 backward 和 optimizer.step，因此不需要 GradScaler
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        predictions = model(inputs)

return predictions

另外，显式传入 dtype= 的算子不会按普通 autocast 规则重新选择 dtype；原地变体和带 out= 的调用也不走同样的自动转换路径。遇到 dtype 与预期不符时，先检查这三类写法，再确认是否把张量送到了正确设备。

如果把 AMP 和 torch.compile 组合使用，还要注意编译反向传播对 autocast 上下文的假设。保持“前向与 loss 在 autocast 内、backward 在外”的结构，并按 PyTorch 编译器文档确认 backward_pass_autocast 的配置与实际训练方式一致。

一张表记住作用域判断

常见问题

autocast 要不要包住 backward？

通常不要。推荐把前向和损失计算放入 autocast，退出后再 backward；反向算子会沿用前向阶段相关的 dtype 选择。

局部关闭 autocast 后为什么还要调用 float()？

enabled=False 只关闭该区域的自动转换，不会替已经是低精度的输入自动升为 float32。敏感算子需要稳定精度时，应显式转换传入张量。

推理也要创建 GradScaler 吗？

不需要。没有梯度和参数更新时，autocast 足够完成推理阶段的混合精度控制。

为什么某个算子没有按预期转换？

检查它是否显式指定了 dtype、是否使用原地调用或 out= 参数，再检查当前 autocast 的 device_type 和上下文是否覆盖了该调用。