这几年在社区里被问得最多的训练问题十个有八个跟混合精度有关。不是不知道用 AMP而是用了之后时不时冒出loss is nan、模型不收敛、显存没降多少这些现象最后要么关掉 AMP 回到 FP32 硬撑要么照葫芦画瓢把scaler代码抄进去却不知道它在干什么。这篇文章就围绕自动混合精度AMP梯度缩放器展开把它的设计思路、工作原理、实操写法以及各种踩坑经验完整说清楚适合正在用 PyTorch 训练模型、想提显存和速度却又担心精度失控的人阅读。文中所有示例都以 PyTorch 的torch.cuda.amp为主线这套机制也是目前业界最常用、最值得彻底搞懂的方案。1. 为什么需要梯度缩放器自动混合精度的核心设计思路1.1 半精度训练的诱惑与陷阱AMP 全称 Automatic Mixed Precision字面意思是“自动混合精度”它并不要求整个模型都跑在半精度 FP16 下而是让算子在 FP16 和 FP32 之间自动选择最合适的精度。传统 FP32 训练占用显存大、计算慢但如果直接把模型参数和中间激活全部改成 FP16绝大多数模型根本训不动。原因在于半精度浮点数的表示范围很窄FP16 能表示的最大有限值是 65504最小正常数大约是 6.1e-5次正规数可以更小但精度损失极其严重。训练过程中梯度数值常常低于 1e-4随便一个下溢就可能让梯度变成 0参数纹丝不动loss 看起来像一条水平线。我见过不少新手的做法是手动把模型model.half()然后所有输入也转成.half()训练几轮后发现 loss 不降要么直接梯度爆炸变成 inf要么干脆一死到底。这就是典型的不理解 FP16 数值范围导致的问题。自动混合精度的思路不是盲目全半精度而是让卷积、矩阵乘法、线性层这类对精度不敏感的算子走 FP16同时维持 FP32 的 master weight再靠损失缩放来解决梯度下溢。1.2 梯度缩放器在哪里起作用GradScaler也就是我们常说的梯度缩放器是 AMP 体系中专门解决“梯度下溢”和“梯度溢出”的组件。在 PyTorch 的标准流程里它负责将 loss 放大一定倍数再反传让梯度从很小的数值范围抬升到 FP16 能够安全表示的区间完成反向传播后再在优化器更新前把梯度“缩小”回来。注意这里的关键缩放发生在 loss 上而不是手动去缩放每个梯度。因为链式法则loss 乘以一个标量 S所有梯度也会同步乘以 S。原本落在 FP16 表达范围之外的微小梯度乘上 S 后落到正常范围这样反向传播过程中梯度数值就不会归零。同时如果梯度本身过大也会在计算图中直接被 FP16 截断检测成 inf这时缩放器需要进行一次“安全处理”。GradScaler在 PyTorch 中被封装成一个类内部维护一个动态缩放的 scale 值并根据每次optimizer.step()前的梯度检查结果来自动调整。这样做的好处是几乎没有额外的手动调参成本坏处则是很多人只调用了 API 却不理解内部逻辑遇到问题完全无从下手。下面我们把它按原理拆开看。2. 梯度缩放器原理拆解动态缩放、溢出检测与恢复2.1 loss 缩放是怎么保护梯度的为了直观理解可以看一组典型数值。假设反向传播过程中某个权重梯度真实值是g 2.5e-5FP32 下这不是问题但 FP16 的正常表示范围最小只到 6.1e-5 左右所以这个梯度在 FP16 里直接变成 0。若我们先把 loss 放大 1024 倍梯度也随之变成0.0256这个数在 FP16 范围内非常安全反传时就不会丢。拿到放大后的梯度后再在优化器真正更新前除以 1024就还原出真实梯度。但缩放不是越大越好。如果把缩放系数设成 65536原本梯度为 3 的数值就变成 196608远超 FP16 的 65504 上限立刻溢出成 inf。因此梯度缩放器要解决两个相反方向的危险下溢导致梯度消失上溢导致梯度爆炸。它不可能固定选一个缩放系数训练过程中梯度分布是动态变化的不同层、不同 batch、不同学习率阶段都会差异很大所以就需要动态调整。PyTorch 的GradScaler默认初始缩放系数是2**16也就是 65536之后每连续growth_interval步没有出现 inf/nan就乘growth_factor默认 2.0尝试放大一旦检测到某一步出现 inf/nan就乘以backoff_factor默认 0.5缩小缩放系数。这就是动态缩放策略的核心逻辑。它本质上是做控制理论中的反馈调节观察输出梯度状态反向调节输入信号幅度。2.2scaler.scale(loss).backward()到scaler.update()的内部逻辑标准 AMP 训练循环中有三步操作顺序绝对不能乱scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()第一步scaler.scale(loss).backward()做的事情是将当前缩放系数 S 乘到 loss 上再触发反向传播。这里 S 是一个 FP32 的 Python float乘到 FP32 loss 上后整个反向计算图里梯度都变成“被放大 S 倍”的 FP32 梯度。由于 PyTorch 反向传播过程中每个算子拿到的是 FP32 梯度它内部再自动转成 FP16 参与计算时数值下溢概率就大大降低。第二步scaler.step(optimizer)内部是一个非常关键的“检测点”。它不会盲目调用optimizer.step()而会先反缩放梯度。准确地说是先对所有梯度调用unscale_操作也就是把梯度整体除以 S同时检查是否存在 inf 或 nan。如果有任何一个参数的梯度发生溢出就跳过本次optimizer.step()也就是说本轮干脆不更新参数然后返回。如果没有异常就正常执行优化器更新。第三步scaler.update()用来更新缩放系数。它在每次 step 之后调用。如果刚才step顺利执行并且梯度没有异常就累加一个计数器达到growth_interval后放大 S如果刚才检测到异常就立即缩小 S并重置计数器。这种先放大、反向传播、再缩小、异常跳过的设计保证了最终生效的梯度与真实梯度在数值上等价同时不会因为一次异常梯度炸穿模型。2.3 关键参数和 API 细节GradScaler最常用的构造参数有init_scale、growth_factor、backoff_factor、growth_interval、enabled。参数默认值作用说明init_scale65536.0初始缩放系数决定起步时的梯度放大倍数growth_factor2.0连续正常步数达标后放大缩放的倍率backoff_factor0.5检测到溢出后缩小缩放的倍率growth_interval2000连续多少次 step 无溢出才放大一次缩放enabledTrue设为 False 时完全关闭缩放逻辑等效于普通训练默认参数对绝大多数 CV、NLP 模型都够用。不过有几个值得注意的点init_scale过大可能导致后续动态调整经常触发回退训练前期波动大但设置过小又会让小梯度直接消失。我习惯保持默认除非遇到频繁溢出再单独调backoff_factor或init_scale。此外PyTorch 还提供了torch.cuda.amp.autocast上下文管理器通常写作with torch.autocast(device_typecuda, dtypetorch.float16):。它负责算子层面的精度自动选择跟 GradScaler 解决的是完全不同的问题。所以 AMP 训练的完整代码永远是“autocast GradScaler”组合autocast 管前向和损失计算精度GradScaler 管反向传播和优化器更新。两者缺一不可。3. 实操过程PyTorch AMP 训练完整流程3.1 单卡训练标准写法我直接给出一套最常用的单卡训练模板这是我在多个检测、分割、生成模型上验证过的标准写法。以分类任务为例伪代码结构如下import torch import torch.nn as nn from torch.cuda.amp import GradScaler # 模型和优化器照常定义 model ResNet18().cuda() optimizer torch.optim.SGD(model.parameters(), lr0.1, momentum0.9) criterion nn.CrossEntropyLoss() # 创建梯度缩放器 scaler GradScaler() for epoch in range(epochs): for images, labels in train_loader: images images.cuda() labels labels.cuda() optimizer.zero_grad() # 前向过程用 autocast with torch.autocast(device_typecuda, dtypetorch.float16): outputs model(images) loss criterion(outputs, labels) # 反向传播使用 scaler scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这段代码看着简单但每一行都有讲究。optimizer.zero_grad()必须放在 autocast 外面还是里面其实都可以但为了统一习惯我在前向和反向前清零梯度。autocast 只包住前向和 loss 计算不代表反向也要包。反向通过scaler.scale(loss).backward()触发计算图本身受 autocast 上下文影响的是“算子执行时的 dtype”而不是“是否允许梯度反传”。关键点在于optimizer.step()与scaler.step()不能写反。如果没有梯度溢出scaler.step内部完成梯度反缩放后执行真正的优化器更新所以此时scaler.step(optimizer)等于optimizer.step()。如果写了optimizer.step()在前梯度是放大后的梯度相当于直接用了被放大几万倍的梯度去更新参数模型必炸无疑。3.2 代码逐步解释和 scale 变化观测为了让大家能看到缩放器的实际工作状态我建议在训练日志里打印两个值scaler.get_scale()和当前 loss。代码片段如下current_scale scaler.get_scale() print(fepoch {epoch}, batch {batch}, loss {loss.item():.4f}, scale {current_scale:.1f})get_scale()返回的是当前缩放系数。正常训练时这个值会在一段时间内保持某个量级然后突然跳到两倍再过一段时间再翻倍直到遇到溢出回退。如果模型收敛稳定scale 往往保持在 65536 或继续涨到 131072说明梯度整体并不大偏小需要通过放大 scale 防止梯度下溢。另一种观测方式是查看optimizer.param_groups[0][lr]和scaler.get_scale()的变化曲线。如果 scale 快速回退到 2048 或 1024并且反复跳动这说明梯度经常溢出。此时先别急着改学习率先看是不是模型里有某些层对 FP16 不友好后面会详细讲。3.3 多卡 DDP 训练中的注意点分布式数据并行DDP与 AMP 组合时流程大体一致但有一点容易被忽略DDP 的梯度同步默认在反向传播结束时通过 AllReduce 完成。使用 GradScaler 时scaler.scale(loss).backward()传出的梯度是放大后的梯度DDP 同步的也是放大后的梯度。由于梯度对每个 rank 来说都放大了同样倍数AllReduce 求平均后再反缩放数学上和“反缩放再求平均”等价所以不会错。真正需要注意的是scaler.step()必须在所有 rank 上同步执行也就是所有进程都完成反向传播后再统一调用否则 DDP 的梯度同步可能卡住。from torch.nn.parallel import DistributedDataParallel as DDP model DDP(model, device_ids[local_rank]) for images, labels in train_loader: optimizer.zero_grad() with torch.autocast(device_typecuda, dtypetorch.float16): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() # 所有 rank 执行到这一步再 step scaler.step(optimizer) scaler.update()这里不需要手动设置什么分布式相关的缩放参数GradScaler 本身不会在跨卡之间做特殊通信它只感知本卡上的梯度状态。如果某一卡出现 inf该卡会跳过optimizer.step但其他卡可能没有跳过这是因为 DDP 中梯度 AllReduce 之后每张卡看到的梯度不一致所以可能出现不同步。稳妥的做法是在所有 rank 上让scaler.scale(loss).backward()产出相同的 loss 和梯度分布换句话说数据分布、随机种子一致时基本可控或者显式在scaler.step()前对梯度做全部检查。实际工程中梯度出现 inf 通常意味着模型本身有问题而不是缩放器的问题所以很少遇到跨卡跳步造成的严重问题但心里要有这根弦。4. 常见问题与排查技巧实录4.1 loss 突然变成 nan 或 inf 的常规排查nan和inf是 AMP 训练里最经典的故障。我总结出一套固定排查路径按顺序执行绝大多数问题都能快速定位。第一步关闭 autocast 和 GradScaler用纯 FP32 跑同样数据。如果 FP32 下也出现 nan说明问题不在 AMP而是模型、学习率、数据本身有 bug。如果 FP32 正常才进入下一步。第二步保留 autocast暂时手动把scaler GradScaler(enabledFalse)跑一次。如果 FP32 版本的 autocast 正常但半精度版本出问题多半是模型中某些算子对 FP16 过于敏感。常见嫌疑对象有 BatchNorm 的 running stats、LayerNorm 的 epsilon 太小、注意力 softmax 后接的某些操作、损失函数内部累计顺序等。第三步在训练循环的scaler.step(optimizer)前显式调用scaler.unscale_(optimizer)然后打印optimizer.param_groups[0][params][0].grad的范数。如果梯度中存在 inf就能快速定位是哪一层参数产生的。日志片段参考scaler.unscale_(optimizer) for name, param in model.named_parameters(): if param.grad is not None and not torch.isfinite(param.grad).all(): print(fFound non-finite grad in {name})第四步检查 learning rate。AMP 场景下模型对学习率的敏感度通常比 FP32 更高尤其当 loss 缩放后梯度范围被改变即使有动态调控一次性过大的lr也可能让权重更新直接迈过合理区间。遇到“训练前几十步正常某一步突然 nan”的情况优先把 lr 调低 10 倍再试。这也是最常被忽略的原因。4.2 显存没有如预期下降很多人以为用了 AMP 显存一定减半实际通常只下降 10% 到 30%。原因在于 PyTorch 的显存缓存分配器不会在每次反向传播后自动将显存归还给系统而是保留在缓存中复用。即使模型参数和激活改成了 FP16之前在 FP32 模式下分配的缓存块仍可能被复用不会立刻收缩。解决方式通常是在训练脚本启动前设置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True让显存段可以动态扩展和收缩配合 AMP 能看到更好的显存收益。另一种做法是减小 batch size观察显存变化如果确认激活显存占用显著减少说明 AMP 确实生效只是 PyTorch 的缓存机制掩盖了效果。需要区分的是AMP 节省的是“模型中间激活”和“梯度”的显存但 master weight 仍然以 FP32 保存优化器状态如 Adam 的 momentum 和 variance也仍是 FP32。所以它省的不是全部而是部分。工程上常见做法是将 AMP 与梯度累积、混合精度优化器如 bitsandbytes 的 8-bit Adam结合才能进一步压显存。4.3 精度不达标或收敛变慢模型收敛速度变慢通常不是缩放器本身的问题而是梯度精度被截断到 FP16 后带来的一点点误差在累积。PyTorch 的 autocast 默认让 conv、linear 等核心算子走 FP16但某些结构性算子仍保持 FP32所以大多数模型效果不受影响。真正容易出问题的是两个地方。一是 BatchNorm。原版 BatchNorm 在半精度下计算均值和方差时数值稳定性较差尤其是 batch 比较小时。如果发现用 AMP 后模型验证指标比 FP32 差一截可以尝试把网络中的 BatchNorm 固定为 FP32。PyTorch 中可以通过torch.autocast的disable局部控制也可以把 BN 的forward用torch.cuda.amp.custom_fwd(cast_inputstorch.float32)强制输入保持 FP32。二是梯度裁剪。使用 AMP 训练 RNN、Transformer 类模型时经常需要梯度裁剪。正确写法有两种要么在scaler.scale(loss).backward()之后直接对“缩放后的梯度”做裁剪但裁剪阈值也要乘 scale要么先scaler.unscale_(optimizer)再按原始梯度做裁剪。第二种更直观scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update()这里scaler.step内部发现梯度已经反缩放后就不会再重复 unscale所以可以直接正常执行。如果不先 unscale剪裁阈值就会错误导致几乎所有梯度被限制在错误范围内训练缓慢。4.4 一个容易混淆的概念嵌入式开发里的 AMP在嵌入式异构多核开发领域也会遇到“AMP”这个词但它指 Asymmetric Multi-Processing非对称多处理比如 RK3506 这类芯片上多核异构系统中一个核跑 Linux、另一个核跑裸机实时任务核间通过共享内存和中断进行通信。这种“AMP”跟深度学习训练里的 Automatic Mixed Precision 完全是两回事。我做嵌入式交叉编译时看到板子资料里写“AMP 中断实例”就误以为是半精度训练结果发现是异构核间通信当时确实有点尴尬。如果你同时接触两个方向务必先看上下文别被同一个缩写带偏。5. 实操心得总结与扩展建议5.1 我反复踩过的几个坑第一永远不要在scaler.scale(loss).backward()和scaler.step(optimizer)之间手动修改梯度倍数。我自己刚开始调试时为了做梯度累积尝试把梯度手动除以累积步数结果和 scaler 的自动缩放叠加后数值混乱模型收敛异常。正确的梯度累积写法是让每个 mini-batch 都正常调用scaler.scale(loss).backward()梯度会自然累加最后一步scaler.step(optimizer)不会重复累积因此不需要额外除以步数。第二scaler.update()的位置不要放错。它必须在optimizer.step()尝试完成之后再调用用于更新下一次循环要用的 scale 值。如果放在 step 之前当前这一步的溢出状态还没被读取scale 更新就会滞后一步导致回退不敏感。第三autocast和GradScaler的启停要配套。如果只加 autocast 不用 GradScaler部分梯度会直接在反向传播中消失loss 看起来一直在降但模型参数基本不变这种“假训练”状态特别隐蔽。如果只加 GradScaler 不用 autocast等于把所有算子还是 FP32 跑梯度被放大再缩小白白浪费显存和速度。5.2 扩展到推理和部署场景AMP 不只是训练专用。推理阶段把模型权重转成 FP16 通常能明显降低显存占用并提高吞吐。PyTorch 里可以用torch.autocast包住推理前向也可以直接用model.half()torch.jit.trace导出 TorchScript 模型。但要注意区分训练时的 GradScaler 只服务于反向传播推理时不需要也不应该创建 scaler。推理前如果做过model.half()所有输入也需要转成 FP16否则数值类型不匹配会直接报错。如果你导出到 TensorRT 或 ONNX混合精度图的处理方式又不一样。TensorRT 会在构建 engine 时用校准数据自动选择各层精度不需要手动插入 scaler。这点和 PyTorch 训练完全是两套逻辑不要混为一谈。5.3 后续扩展方向AMP 的下一步是 BF16 训练。BF16 的指数位和 FP32 一样数值范围比 FP16 大得多但尾数位少所以在部分场景下更适合直接替代 FP32且几乎不需要损失缩放。Ampere 及更新架构 GPU 上PyTorch 只需要把 autocast 的dtype改成torch.bfloat16GradScaler 甚至可以直接省略。不过 BF16 在显存节省上不如 FP16 明显具体选哪种要看硬件、算子和任务容忍度。如果你想深入理解 GradScaler 的底层实现可以直接读 PyTorch 源码中torch/cuda/amp/grad_scaler.py和自动混合精度模块torch/cuda/amp/autocast_mode.py。源码里_scale_update和_unscale_grads_的逻辑非常清晰看完之后对“溢出检测、跳过优化器步骤、动态调整系数”这三个核心操作会有更直观的认识。在我实际使用中AMP 梯度缩放器现在已经是训练脚本的标配但最忌讳的是无脑套模板。只有把 loss 缩放、反缩放、动态回退、autocast 边界这些环节都理解清楚碰到 nan 或者收敛问题时才能快速定位。这套思路换个框架比如 TensorFlow 的tf.keras.mixed_precision.set_global_policy、PaddlePaddle 的 AMP也殊途同归只是 API 名称不同背后的数值逻辑完全一致。