自动混合精度（AMP）

概念更新 2026-08-02置信度 high#概念#基础#长青#训练#显存

AMP 根据算子数值特性在低精度与高精度之间自动选择，并用梯度缩放降低 FP16 梯度下溢风险，以节省显存和提高吞吐。

两个核心组件

`autocast`

在指定上下文中，让矩阵乘等适合低精度的算子使用 FP16/BF16，而归一化、归约等敏感操作按框架策略保留更高精度。它不是把整个模型永久转换成 .half()。

`GradScaler`

FP16 的微小梯度可能下溢为 0。动态 loss scaling 先计算：

$L_{scaled}=S\cdot L,$

反向传播得到放大的梯度；在更新前 unscale。如果检测到 inf/NaN，则跳过该步并降低 $S$ ；稳定时可逐步增大 $S$ 。

典型训练循环

scaler = torch.cuda.amp.GradScaler()

for x, y in loader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        logits = model(x)
        loss = criterion(logits, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

在梯度累积场景中，scale、unscale、clip 和 optimizer.step() 的位置必须与 accumulation 边界一致，不能每个 micro-batch 都更新参数。

AMP 能节省什么

部分激活和中间张量用 16 bit 存储；
支持 Tensor Core/低精度内核的算子可能提高吞吐；
因单步显存降低，特定环境可容纳更大的 micro-batch。

但 FP32 master weights、优化器状态、部分算子和缓存仍可能占较大内存。课程中“batch size 显著提高”的结果是讲者环境实测，不是跨模型、跨 GPU 的固定倍数。

平台与版本边界

课程 notebook 依赖 CUDA AMP；macOS/MPS 不能原样复现 torch.cuda.amp。现代 PyTorch 的命名和推荐调用可能随版本变化，概念应与当时 notebook 的 API 示例分开看。

常见误解

AMP = 所有张量 FP16：错误；混合精度的关键正是按算子保留必要的高精度。
有 autocast 就不需要 scaler：BF16 常不需要动态 scaling，但 FP16 训练中 scaler 仍很重要。
AMP 只省权重：主要收益还包括激活、带宽与低精度算力。
AMP 会自动解决数值问题：异常 loss、错误归一化或不稳定优化仍需诊断。