"自动混合精度"这六个字，在很多训练脚本里可能只是几行配置，但真正把它吃透之后，你会发现它是目前提升训练效率最成熟、性价比最高的一招。简单说，自动混合精度（AMP）就是在训练过程中让模型自动在 FP16 和 FP32 两种精度之间切换：前向和反向计算主要用 FP16 跑，权重更新等关键环节留在 FP32，从而同时降低显存占用、拉高计算吞吐。这个技术解决的痛点非常直接：模型参数越来越多，一块卡装不下了；或者显存明明够，但训练速度被算力上限卡住。这篇文章我从底层原理讲到 PyTorch、TensorFlow 两套落地实现，再把训练中踩过的坑和排查思路一并整理出来，适合正在调大模型训练、想榨干显卡性能、或者第一次接触加速技术的同学。

1. 混合精度训练的底层逻辑：省显存只是表象，Tensor Core 才是核心

很多人以为混合精度就是"用更少位数的数字，省点显存"，这个认知只对了一半。省显存确实是直接收益，但更关键的是它把 GPU 上那颗专门为半精度设计的计算单元——Tensor Core 真正用起来了。不理解这一层，后面很多配置就只是照着抄，出了问题也不知道从哪排查。

1.1 FP16 到底省了多少显存：一个可以直接算的例子

先看一组基础数据：

有一个点值得多说一句：BF16 虽然也是 16 位，但它的指数位和 FP32 一样多，所以表示范围跟 FP32 相同，只是尾数精度更低。这意味着 BF16 在训练中不容易出现 FP16 那种"数值溢出直接变 inf"的问题，代价就是精度更低一点。从 A100 开始，新一代显卡基本都原生支持 BF16，如果你的卡支持，优先用 BF16 做混合精度是更稳的选择。

显存的节省算起来很简单：权重、梯度、优化器状态是训练显存的三座大山。以 Adam 优化器为例，它除了模型参数本身，还要额外保存一阶动量、二阶动量，每个值各一份。如果用 FP32 全精度训练，一套 FP32 权重在显存里要摊三份账；切换到混合精度后，模型权重的一份 FP32 主副本仍然保留，但前向计算使用的 FP16 副本和梯度可以直接按半精度存放，显存开销能再砍掉约 40% 到 50%。实测中，一个原本 12GB 显存刚好卡住的模型，切到 AMP 后往往能腾出 2 到 4GB 余量，直接把 batch size 提上去，或者把输入分辨率再加大一档。

1.2 Tensor Core：为什么半精度反而比全精度快这么多

如果只是"省显存"，那混合精度还不足以成为训练标配。真正让人无法拒绝的是速度。从 Volta 架构开始，NVIDIA 在 GPU 上加入了 Tensor Core，专门做矩阵乘法和卷积这类密集计算。Tensor Core 的输入要求就是 FP16，它内部做的是 FP16 乘法和 FP32 累加，再把结果转回 FP16 输出。

这里有个关键点：Tensor Core 的算力是普通 FP32 核心的好几倍甚至十几倍。拿常见的数据来说，在某些显卡上，FP16 矩阵运算的吞吐可以达到 FP32 的 8 倍。所以哪怕你的模型只有一部分算子能切到半精度，整体训练时间也能看到明显下降。我实际用某个经典图像分类模型做过对比，同样的 batch size、同样的 epoch，从全 FP32 切到 AMP 后，单卡训练时间大约缩短了 35% 到 45%。这不是玄学，是计算单元本身的工作方式决定了这个收益。

需要提醒的是，Tensor Core 对张量形状有内部切分的偏好，通常是按 8 的倍数、16 的倍数来切分矩阵。如果你的模型里某些维度不是这些数的整数倍，Tensor Core 的利用率会打折扣。实际操作中，把全连接层的 hidden size、卷积的 channel 数尽量对齐到 8 或 16 的倍数，是很多开源模型在结构设计时就在做的事。

1.3 FP16 的天生缺陷：范围和精度都很紧张

FP16 的数值范围最大只有 65504，这在深度学习里是个很尴尬的数字。训练时激活值、梯度如果超过这个范围，直接变成无穷大 inf，损失函数马上变 NaN，训练直接崩掉。反过来，如果数值太小，低于 FP16 能表示的最小规格化数，就会下溢变成 0。梯度在反向传播里往往是很小的数值，很多梯度的绝对值远小于 1，在 FP16 下直接归零，参数更新随之失效。

这就是为什么不能"一把梭"把整个训练都切成 FP16。有些算子对精度敏感，比如 BatchNorm 的统计量、某些损失函数内部的运算，切到 FP16 后训练稳定性会显著变差。所以混合精度的核心思想是：让大头计算跑 FP16 吃到速度和显存红利，把容易出问题的环节留在 FP32，再用额外机制兜住数值溢出的风险。

2. 自动混合精度的核心机制：它到底"自动"了什么

既然要混合，就得有人来决定哪些地方用 FP16、哪些地方用 FP32、梯度太小怎么办、数值溢出了怎么恢复。这套决策逻辑，就是自动混合精度的"自动"所在。搞懂这套机制，你才能在任何框架里都游刃有余，而不只是会抄 API。

2.1 三件套：Master 权重、Autocast 和梯度缩放

名字叫得玄乎，拆开看就三样东西，理解它们三个的分工，AMP 就懂了大半。

第一，Master 权重。哪怕前向和反向都用 FP16 计算，你必须保留一份 FP32 的模型参数副本，所有更新操作都在这个 FP32 副本上完成，然后再转成 FP16 参与下一次前向。原因很直接：模型权重在训练中变化幅度很小，往往在小数点后第四位、第五位这种量级。FP16 只保留大约 3 位有效数字，如果你在 FP16 上直接做参数更新，每次更新幅度可能直接落在 FP16 表示不了的那一位上，权重永远改不动。保留 FP32 主副本，本质上是给参数更新留了足够的"记数精度"。

第二，Autocast。它做的是算子级别的精度调度。你写模型代码的时候不需要手动去加 .float() 或 .half()，autocast 会在前向执行时自动判断每个算子应该跑在哪种精度下。判断依据是算子白名单：矩阵乘法、卷积这类 Tensor Core 友好算子，用 FP16；BatchNorm、Softmax 的某些分支、LayerNorm 这类容易精度敏感的算子，留在 FP32。我见过不少人误以为 autocast 会把整个模型切成 FP16，其实它是算子级别的细粒度调度，不是模块级别一刀切。

第三，梯度缩放（Grad Scaling）。这是对抗 FP16 下溢的核心手段。既然梯度数值普遍偏小，那就在反向传播之前，先把损失值乘上一个较大的缩放因子（比如初始 65536），这样反向传播算出来的梯度也跟着被放大，落回 FP16 能正常表示的区间。注意，梯度是被"搬"到这个范围内的，不是凭空变大，所以梯度数值本身的相对比例关系没变。

2.2 动态 Loss Scale：崩了自动降，稳了自动升

梯度缩放最关键的问题是：缩放因子设多少合适？设大了，梯度直接溢出变 inf；设小了，小梯度照样下溢成 0。固定值很难适配所有模型，所以训练框架普遍采用动态缩放策略。

动态策略的逻辑很朴素的：一开始把缩放因子设得大一点，赌一把看梯度会不会溢出；如果某次训练发现 FP16 梯度里出现了 inf 或 NaN，说明当前缩放因子太高了，就把因子乘以 0.5；如果连续很多步都没有溢出，说明还有放大空间，就把因子乘以 2 往上升。PyTorch 里的默认配置是初始缩放因子 65536，每隔 2000 步检查一次是否该往上调。

这个机制的巧妙之处在于，它让节点在"溢出"和"下溢"之间自动找到一个平衡位置，而且不依赖任何人工判断。唯一需要注意的是，梯度中出现一个 inf 就会被判定为溢出，所以当 loss 突然变成 NaN，第一反应不该是去怀疑数据，而是查一下这一步是不是触发了一次降缩放，看一下训练日志里 loss scale 有没有跳变。

2.3 哪些算子必须留在 FP32：一张可复用的对照表

很多人问我一个很难回答的问题：我能不能自己指定某些层用 FP32？答案是可以，但大多数时候没必要，因为框架内置的算子白名单已经把常见问题都兜住了。下面这张表是我整理出来的一份比较实用的参考：

有一个实战经验供参考：我自己处理过一个语言模型，在 Embedding 和输出层常驻 FP32 后，困惑度下降曲线明显更平滑了，而训练速度的影响其实只有几个百分点。所以如果你发现切了 AMP 之后模型效果退化明显，优先检查这两个位置，把它们留在 FP32 往往是最快的修复方案。

3. PyTorch 里落地 AMP：从零改造训练循环

PyTorch 的 AMP 经历了两个阶段，老版本里用 torch.cuda.amp，新版已经迁移到 torch.amp 这套统一接口。不管 API 怎么变，核心的 autocast 和 GradScaler 两个组件是不变的。下面我直接给一段完整的训练循环改造示例，你在自己的代码里对照着改就行。

3.1 一段可以直接抄的训练循环

先看改造前最普通的 PyTorch 训练循环：

optimizer.zero_grad() loss = criterion(model(images), labels) loss.backward() optimizer.step()

加上 AMP 之后，变成这样：

from torch.amp import autocast, GradScaler scaler = GradScaler() optimizer.zero_grad() with autocast(device_type="cuda", dtype=torch.float16): outputs = model(images) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

只有五个改动点，但每个都很关键：

第一，autocast 的作用范围。把前向计算和 loss 计算包进 with 块里。注意，反向传播不要包在 autocast 里面，它会在 backward 时自动按前向的精度执行，不需要你额外处理。

第二，用 scaler.scale(loss) 代替原来的 loss.backward()。这一步完成两件事：把 loss 乘以缩放因子，再执行反向传播。所以后面不要再单独调 backward 了。

第三，scaler.step(optimizer) 和 scaler.update() 的顺序不能反。step 负责把缩放后的梯度还原成真实梯度，再执行参数更新；如果检测到这次梯度里有 inf 或 NaN，step 会跳过本次优化，避免用坏梯度更新参数。update 在 step 之后调用，负责根据本次结果动态调整缩放因子。

还有一点需要特别说明：优化器最好在循环开始的时候就创建好，GradScaler 也在循环外创建。不要每步都 new 一个 scaler，那样动态缩放永远从初始值开始，等于是自断双臂。

3.2 GradScaler 和梯度裁剪的正确配合

用了 AMP 之后，梯度裁剪的顺序会坑很多人。常见的做法是这样：

scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) scaler.step(optimizer) scaler.update()

为什么必须先 unscale 再裁剪？因为梯度在 backward 之后还是放大过的状态，直接对放大后的梯度做范数裁剪，裁剪阈值等于被放大了同样倍数，裁了等于没裁。所以必须先用 unscale_ 把梯度还原成真实值，再做裁剪，最后再 step。整个顺序是：scale(loss) → backward → unscale → clip → step → update。

有一个细节值得记住：只要你在一个训练步里调用了 unscale_，这一步里就不能再多次调用 step，否则第二次 step 会因为梯度已经被还原而行为异常。如果因为特殊写法必须分多次 step，那就不应该提前 unscale，而是要分别对每个 optimizer 调用 unscale_。

3.3 分布式训练和混合精度的配套细节

DataParallel 或者 DistributedDataParallel 配合 AMP 时，上面这套流程基本不用改，但有几个隐患要提前排查。

第一个是每个进程各自维护一份 GradScaler。DDP 场景下每个进程的模型副本是同一份参数的切片，梯度通过 all-reduce 同步，但 GradScaler 的缩放因子是各进程本地维护的。好在溢出判断是全局同步过梯度的结果，实际运行中所有进程的缩放因子会趋于一致，很少出现一个进程降了另一个没降的情况。真遇到缩放因子频繁跳变且各卡不一致，优先检查各卡的数据加载是否打乱一致。

第二个是梯度裁剪在 DDP 下不要每个进程单独做，会导致梯度范数计算不一致。正确的做法是先用 dist.all_reduce 汇总梯度范数再做统一裁剪，这个已经是另一篇文章的量了，这里只提醒你不要把上一节的裁剪代码直接无脑套进 DDP 训练里。

第三个是 checkpoint 保存时要额外记录 scaler 的状态。我吃过这个亏：训练到中断点，只保存了模型权重和优化器状态，没保存 GradScaler，恢复训练后缩放因子回到初始值，前几千步一直在剧烈波动，损失曲线看起来像踩了弹簧。正确的保存方式是把 scaler.state_dict() 一并存下来，恢复时 load 回去。PyTorch 的 checkpoint 推荐做法就是把 model、optimizer、scaler 三个 state_dict 都打包，缺一个都容易在长训练中断后出问题。

4. TensorFlow / Keras 的混合精度配置：一行策略，注意自定义训练循环

TensorFlow 这边的 AMP 设计思路跟 PyTorch 不太一样。PyTorch 是"显式地包 autocast、显式地管理 scaler"，TensorFlow 则倾向于全局声明一个混合精度策略，然后把精度调度交给框架内部去处理。两种哲学没有谁绝对好，但如果你两个框架都写，很容易因为 API 习惯不同而踩坑。

4.1 全局策略：一行代码把所有层切成混合精度

在 TensorFlow 2 里，最省事的做法是：

tf.keras.mixed_precision.set_global_policy("mixed_float16")

这行代码放在模型定义之前执行，后续所有通过 Keras 层接口创建的层都会自动遵循混合精度策略。你不需要在每一层手动指定 dtype，框架会自动把卷积、全连接这类算子放到 FP16，把 BatchNorm、LayerNorm 留在 FP32。如果你用的是 model.fit 这种高层 API，loss scaling 也是框架自动处理的，训练循环里你几乎感觉不到 AMP 的存在。

不过这里有一个容易误会的点：set_global_policy 只对之后创建的层生效。如果你在设置策略之前已经创建了一批层，这些层很可能还是 FP32 的。实际排查过的问题里，有人把 set_global_policy 写在模型文件最开始，结果模型是用 load_model 从磁盘加载的，层的 dtype 在序列化时就已固定，策略没能生效，训练速度没任何变化。正确的做法是确保你的模型是在策略生效之后由代码构建的，而不是直接加载旧权重。旧权重加载也没事，模型结构本身要重新走一遍构建流程。

4.2 自定义训练循环里的 loss scale 管理

一旦你从 model.fit 切到自定义训练循环（我用过一段时间，主要是为了夹带一些自定义指标和复杂的梯度操作），loss scaling 就不会自动做了，得自己手动处理。典型的自定义循环是：

strategy = tf.keras.mixed_precision.LossScaleOptimizer(optimizer) with tf.GradientTape() as tape: predictions = model(images, training=True) loss = loss_fn(labels, predictions) scaled_loss = optimizer.get_scaled_loss(loss) scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables) gradients = optimizer.get_unscaled_gradients(scaled_gradients) optimizer.apply_gradients(zip(gradients, model.trainable_variables))

这里有两个关键的包装方法：get_scaled_loss 和 get_unscaled_gradients。它们的语义对应 PyTorch 里的 scaler.scale(loss) 和 unscale_，只是以对象方法的形式暴露出来。需要留意的是，LossScaleOptimizer 包装的是优化器本身，所以你原本传给 apply_gradients 的优化器对象要换成这个包装后的实例。如果你在这种自定义循环里用 get_scaled_loss 但忘了 get_unscaled_gradients，梯度会以放大的状态直接应用，参数更新的步长被放大几千倍，损失曲线会直接起飞到火星去。

另外，TensorFlow 默认使用动态 loss scaling，但也可以手动设为固定值：

optimizer = tf.keras.optimizers.Adam() optimizer = tf.keras.mixed_precision.LossScaleOptimizer( optimizer, dynamic=False, initial_scale=1024 )

固定缩放不是常规推荐，但在某些调试场景下很有用：比如你想确认某个模型在特定缩放因子下的数值行为，固定了才好做对比。生产训练还是用动态策略，省心得多。

5. 实战中常见的坑与排查记录：从 NaN 到"没加速"的逐个拆解

AMP 的坑不在 API 本身，而在各种看起来毫不相关的场景里冒出来。我把这几年遇到、以及帮别人排查过的问题整理成几类，每一类都附上排查思路。训练脚本报错不可怕，可怕的是报错信息让你完全摸不着头脑。

5.1 loss 变成 NaN：三种典型原因的快速定位

这是 AMP 训练里最常见的求救信号。很多人遇到 NaN 第一反应是调学习率、换初始化，但在 AMP 场景下，顺序应该反过来。

先检查梯度是否溢出。看日志里 GradScaler 的缩放因子是不是刚刚发生过下降，如果是，说明 FP16 梯度触及了上限。处理办法：把初始缩放因子调低，比如从 65536 降到 4096 试试；或者检查模型里是否有超大输出值的层，比如没有归一化的 logits 直接参与了大数运算。

再检查数据里是否有异常值。输入中出现 NaN 或 inf，会一路污染到梯度里，但这一步在非 AMP 训练里也会崩，不具备 AMP 特异性。通常我会用 torch.isnan 检查输入张量，快速排除这一项。

最后检查优化器和模型的数值敏感层。如果上面两个都没问题，大概率是某个算子切到 FP16 后精度不足导致的梯度异常。排查方法是逐个把 LayerNorm、损失函数这些候选层固定到 FP32。我遇到过的一个实际案例是模型里用了自定义的余弦相似度损失，内部做了归一化，这个算子对精度非常敏感，切 FP16 后训练频繁崩溃，固定到 FP32 后问题消失，训练时间只增加了约 8%。

5.2 切了 AMP 但速度毫无变化：先查这四个位置

这是第二常见的抱怨。速度没有提升，原因往往不在 AMP 本身，而是瓶颈根本不在这。

第一个要查的是 GPU 利用率。用 nvidia-smi 看显卡的利用率，如果你发现 GPU 一直不满，说明你的训练是数据加载瓶颈型而不是算力瓶颈型。混合精度再怎么优化计算，也救不了一个在等着 CPU 喂数据的训练流程。解决办法是开 DataLoader 的多进程、加大 prefetch_factor、用内存映射加载数据。

第二个要查的是模型是否太"瘦"。AMP 的加速主要靠 Tensor Core，如果你的模型以单元素逐张量操作为主、没有多少大矩阵运算，那收益非常有限。一个只有几个全连接层的小模型切 AMP，速度提升可能只有 5% 以内，但显存倒是实打实省了一截。这种场景，AMP 依然有价值，只是价值在显存而不在速度。

第三个要查的是输入管线和计算之间的重叠。即便 GPU 利用率看起来不错，每次反向后如果设备同步点太多，也会白白损失吞吐。PyTorch 里减少 .item() 的调用、把评估指标的计算挪到训练循环之外，都是常见优化。

第四个位置比较隐蔽：检查你的卷积输出通道数。Tensor Core 对通道数不是任意值都高效，落到不是 8 的倍数的维度上，内部要做 padding 之类的处理，利用率会掉。我曾把一个模型中间的卷积通道从 96 调到 128，在显存几乎不变的情况下，FP16 下的训练速度又提了一截。

我把这些排查项和优先级做成了速查表，方便你对照操作：

5.3 哪些模型不适合 AMP：提前甄别可以少走弯路

AMP 不是万能的，有些场景切了反而添乱。根据我的经验，下面几类情况要谨慎。

第一个是训练量级非常小、模型只有几百万参数的场景。这种模型对算力要求本来就不高，切 AMP 后速度提升不明显，反而多了数值风险和调试成本。收益和成本不成正比，不如保持全 FP32。

第二个是数值范围分布极广的任务。某些科学计算类模型，比如流体模拟里的神经网络代理模型，中间张量经常横跨好几个数量级。FP16 表示范围只有 65504，这类模型切半精度后即使有 autocast 兜底，也容易出现训练崩溃或精度损失。对这种任务，优先考虑 BF16，表示范围与 FP32 一致，风险小很多。

第三类是已经在使用某种自定义精度策略的模型。如果你的代码里早就手动做了 transfer learning 层的精度管理，是冻结某些层、对某些层设 .half()，再叠加 AMP 后两者的精度策略可能互相打架。排查起来非常痛苦，建议二选一，不要叠加。

5.4 两个容易被忽略但很实用的技巧

最后分享两个不怎么出现在文档里、但实际非常有用的小经验。

第一个，监控 loss scale 的走势本身就是一个诊断信号。如果缩放因子长期不增长、持续低迷，说明模型里的梯度一直偏大，这往往暗示模型的初始化或学习率不合理，哪怕全精度训练能正常收敛，这样的配置在 AMP 下也会更危险。所以我把 loss scale 曲线的周期性波动当作训练健康度的参考指标之一。

第二个，在做超参数迁移时，AMP 训练的学习率通常可以直接沿用全精度配置，不用特意调小。因为主权重始终是 FP32，参数更新粒度没变。但如果发现切了 AMP 后收敛变慢，优先检查是不是缩放因子经常触发跳过步——当梯度频繁溢出时，step 会被跳过，等于学习率被隐性调小了。这时候把初始缩放因子调低，往往比调学习率更有效。

自动混合精度这套技术，说复杂也复杂，说简单也简单。复杂在背后的数值机制和各种边界情况，简单在最终落地就是几个 API、一行配置、一个正确的调用顺序。我个人的体会是，AMP 不应该被当成一个"开一下就行"的加速开关，而应该是你训练工具箱里一个需要理解其脾气、知道它怕什么的常规工具。先把第 2 节里那三件套的机制弄清楚，再照着第 3 节、第 4 节的代码落地，遇到问题回到第 5 节的排查表一个个过一遍，大部分坑其实都能在半小时内定位。最后再补一句实操心得：每次新项目，我会先在原模型上完整跑一个全 FP32 的 baseline，记录稳定的 loss 曲线、显存峰值和吞吐数据，再切 AMP。没有这个对照，你后面排查任何问题都少了一个最重要的参照系。这个习惯帮我省掉了大量"看起来正常但实际在退化"的隐性坑，建议你也试试。