做深度学习训练，尤其是大模型微调或者CV任务，最让人抓狂的往往不是模型结构写不出来，而是同一套代码，别人8G显存跑得飞快，到你6G的卡上第一步就OOM。这时候很多人的第一反应是换显卡，或者疯狂砍batch size，其实还有一个性价比极高的思路被忽略了：PyTorch AMP混合精度训练。它能同时降低显存占用、提升训练吞吐，改动量小到只需要在训练循环里加两三行代码，而且不需要换卡、不需要改模型结构。这篇博文就围绕AMP实战展开，把原理、代码、调参和踩坑一次讲清楚，适合正在被显存不足困扰、想提升训练速度的PyTorch使用者。
1. 为什么大家都在用混合精度训练
1.1 先弄清楚显存到底被谁吃了
显存不是只装模型参数这么简单。跑一次训练，显存里同时放着四样东西：模型参数（weights）、优化器状态（optimizer states）、前向过程的激活值（activations）以及临时梯度（gradients）。以AdamW优化器为例，它除了模型参数本身，还要维护一阶动量m和二阶动量v，这两个状态跟参数同尺寸，意味着每个参数在显存里占用的字节数比想象中多得多。
很多人问“模型参数和显存到底什么关系”“MoE架构是不是所有参数都要进显存”，本质上都是在问同一个问题：到底哪些东西占了显存大头。答案分场景。小模型场景下，优化器状态和参数占大头，所以你会看到FP32训练一个1B模型，光参数加优化器状态就吃掉12GB以上。大模型或者长序列训练场景下，激活值才是吃显存的巨头，因为每个token、每层网络都要保留一份中间结果用于反向传播。AMP混合精度训练恰好对这两块都有作用：参数和激活值改用FP16存储，显存占用直接砍掉一大块。这就是“省显存”的第一层逻辑。
1.2 省显存的核心原理：FP16的“半字节”优势
浮点数格式决定了存储开销。FP32用4字节表示一个数，FP16只用2字节，Bit数直接减半。显存占用本质上就是字节数，同一条数据从FP32换成FP16，占用自然减半。举个例子，一个1亿参数模型，FP32权重占400MB，FP16权重只占200MB，如果前向激活值原本占2GB，换FP16后大概能压到1GB多一点。整体下来，很多模型能把峰值显存从12GB压到7-8GB，省出的这4GB足够你把batch size翻倍，或者把序列长度拉长。
但这里有个很多新手容易误解的点：AMP并不是把所有东西都变成FP16。它叫“混合精度”，核心思想是“该用FP16的地方用FP16，该保FP32的地方保FP32”。比如主权重通常会保留一份FP32副本用于参数更新，前向计算和梯度计算用FP16副本。为什么非要留FP32副本？因为FP16的数值范围太窄，用它直接做参数更新，学习率稍微大一点，权重更新量就会被舍入误差吃掉，训练直接不收敛。保留FP32 master weight是稳定性和精度之间的一个平衡方案。
1.3 提升吞吐的关键：Tensor Core与带宽减半
FP16带来的第二个收益是吞吐提升。这里有两个来源。第一个是GPU上的Tensor Core单元，它专门为半精度矩阵运算做了优化，FP16的矩阵乘法峰值算力远高于FP32普通计算。以T4为例，FP32算力大约8.1 TFLOPS，而FP16 Tensor Core算力能到65 TFLOPS左右，理论上差了近8倍。当然端到端不会真有8倍收益，因为你的模型里不是所有算子都能落到Tensor Core上，但1.5到3倍的综合提速在计算密集型任务里非常普遍。
第二个来源是显存带宽。训练大模型时，数据搬运往往比计算更耗时。FP16数据减半，意味着从显存读取相同“意义”的数据，耗时减半。像attention、embedding这类带宽敏感算子，收益尤其明显。再加上多卡训练时，梯度同步的通信字节数也能减半，整个训练吞吐自然就上去了。所以AMP不是“投机取巧”，而是从硬件架构层面把冗余的字节数和计算格式消掉。
2. AMP在PyTorch里到底是怎么工作的
2.1 autocast：不是把所有计算都切成FP16
PyTorch从1.6开始把AMP做进了官方API，核心组件是torch.autocast（老版本是torch.cuda.amp.autocast）。它做的事情不是简单地把所有tensor转成half，而是维护一张“算子策略表”：哪些算子适合FP16，哪些算子用FP32更稳，哪些算子无论输入是什么都强制FP32。比如Conv、Linear、MatMul这类计算密集型算子，在输入是FP16时就会以FP16计算；而Softmax、BatchNorm、LayerNorm这类对精度敏感的算子，即使你传入FP16张量，autocast也会在内部用FP32计算，计算完再转回去。
这个设计非常关键。如果你手动把所有输入都.half()喂给模型，大概率会出现数值不稳定，甚至直接训练发散。但用autocast包裹前向传播之后，你什么都不用管，PyTorch自动为每个算子选择合适精度。这也是为什么AMP的接入成本这么低——你不需要理解每一层的数值特性，框架帮你做了。
2.2 GradScaler：防止梯度消失的保险丝
FP16有个天然缺陷：可表示的数值范围只有大约±65504，而且越靠近0，精度越低。深度学习反向传播的梯度经常非常小，比如1e-7甚至更小，这种数值在FP16里会被表示成0，也就是梯度下溢（underflow）。一旦梯度变成0，网络权重就再也不更新了。
GradScaler就是用来解决这个问题的。它的做法很直接：在反向传播之前，把loss整体乘以一个缩放因子（默认初始值是2的16次方，即65536），让梯度进入FP16可表示的安全范围；反向传播结束后，在优化器更新参数之前，再把梯度除以同一个缩放因子。因为缩放方式对整个梯度张量是统一的，方向不变，数值又被拉回合理区间，所以训练精度基本不受影响。
而且GradScaler是动态的。它会周期性检查梯度或者loss是否出现inf/nan，如果连续若干步都没有问题，就会适当调大scale；如果某一步溢出了，就回退scale再跳过这一步的参数更新。整个机制像一个带保护阀的值班员，既保证梯度不消失，又防止放大过头。这也是为什么AMP实现里必须搭配GradScaler使用，光用autocast不配GradScaler，很多模型训练到一半就废了。
2.3 新旧API与设备支持情况
PyTorch的AMP API有过一次演进。早期版本走的是torch.cuda.amp.autocast和torch.cuda.amp.GradScaler，到了PyTorch 2.x，官方把API统一成了torch.autocast(device_type="cuda", dtype=torch.float16)和torch.amp.GradScaler("cuda")。功能上两者等价，旧写法也还能用，但新项目建议直接用新API，因为后续维护和扩展都以新API为准。
设备支持上也值得说清楚。AMP的收益建立在GPU硬件支持Tensor Core的基础上，NVIDIA从Volta架构（V100）开始支持FP16 Tensor Core，Turing（T4/RTX 20系）、Ampere（A100/RTX 30系）、Ada Lovelace（RTX 40系）都支持得很好，所以只要你不是老掉牙的GPU，都能吃到红利。AMD ROCm环境下，PyTorch的AMP也做了适配，但稳定性和性能优化程度不如NVIDIA平台，这一点要有心理预期。如果你在CPU上训练，就别指望AMP了，CPU端的half运算没有专门加速单元，强行.half()反而更慢；CPU场景一般用BF16，PyTorch新版也支持torch.autocast("cpu", dtype=torch.bfloat16)，但收益主要在内存占用而不是计算速度。
3. 从零把训练脚本改成AMP的完整实操
3.1 改造前先量化基线
我见过太多人一上来就改代码，改完发现“好像快了一点”但又说不清快了多少，显存省了也没记录下来，最后根本没法判断AMP到底值不值得用。正确做法是先跑一次基线，把下面三个数字记下来：单次epoch耗时、峰值显存、每秒训练样本数。
写一个简单的benchmark脚本，用torch.cuda.max_memory_allocated()统计显存峰值，用时间戳统计吞吐，不需要额外工具：
import torch import time def benchmark(model, loader, device, num_steps=50): model.train() start = time.time() sample_count = 0 for i, (x, y) in enumerate(loader): x, y = x.to(device), y.to(device) optimizer.zero_grad() loss = model(x, y) loss.backward() optimizer.step() sample_count += x.size(0) if i >= num_steps: break elapsed = time.time() - start throughput = sample_count / elapsed peak_mem = torch.cuda.max_memory_allocated() / 1024**3 print(f"throughput: {throughput:.2f} samples/s, peak mem: {peak_mem:.2f} GiB")跑完基线再动手改。这样后续对比才有依据。顺便说一句，max_memory_allocated统计的是PyTorch实际分配的张量内存，和nvidia-smi看到的进程显存不一样，后者还包括CUDA上下文、缓存分配器等开销，但记录训练趋势用前者就够准了。
3.2 最小改动接入AMP
假设你原来的训练循环长这样：前向算loss、loss.backward()、optimizer.step()。接入AMP只需要四步改造：初始化一个GradScaler、用autocast包裹前向、用scaler.scale(loss)替代直接backward、用scaler.step替代optimizer.step，最后别忘了scaler.update。
import torch device = "cuda" model = MyModel().to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) # 1. 初始化 scaler scaler = torch.amp.GradScaler("cuda") for epoch in range(epochs): for x, y in train_loader: x, y = x.to(device), y.to(device) optimizer.zero_grad() # 2. 前向传播放进 autocast with torch.autocast(device_type="cuda", dtype=torch.float16): loss = model(x, y) # 3. 反向传播用 scaler.scale 包裹 scaler.scale(loss).backward() # 4. 参数更新用 scaler.step scaler.step(optimizer) scaler.update()就这么简单。有几个细节必须注意。loss必须是一个标量tensor，如果你的模型返回dict或者多个loss，自己先合成一个再传给scaler。scaler.scale(loss).backward()这一步要放在autocast外面，不能放进with torch.autocast()里面，因为scaler放大的操作本身是FP32精度。还有前向输入不需要手动.half()，把原始FP32数据喂进autocast块即可，autocast会自动管理。
如果你用的是PyTorch 1.x，把torch.amp.GradScaler("cuda")换成torch.cuda.amp.GradScaler()，把torch.autocast换成torch.cuda.amp.autocast，其余逻辑一模一样。
3.3 跑通后如何进一步榨取吞吐
接上AMP只是第一步，省出来的显存和计算余量还能做三件进阶操作。
第一个操作是增大batch size。AMP省下来的显存，本质上是训练预算，最直接的用法就是把batch size调大。batch size增大后，GPU计算效率更高，吞吐还能再上一个台阶。我实测过一个BERT-like模型，FP32基线batch size开到8就到顶了，AMP后能开到16，吞吐从每秒约320样本涨到约580，接近翻倍。
第二个操作是开启梯度累积替代暴力增大batch。有些场景增大batch size会直接影响优化器行为，或者数据加载跟不上。这时可以用梯度累积，每N个batch累积梯度再更新一次。AMP省下的显存可以让你把N从2提到4或8，等效batch变大，训练更稳，吞吐也不降。
第三个操作是给优化器状态减肥。AMP主要动了参数和激活值，但AdamW的优化器状态还是FP32。如果想进一步压显存，可以配合8bit优化器（比如bitsandbytes库），把m和v降到8bit，这部分能再省好几GB。注意AMP和8bit优化器是正交的，两者可以同时用，我在6G显存的卡上跑LoRA微调大模型时，就是“AMP+8bit AdamW+梯度累积”三个一起上，效果立竿见影。
还有一个容易忽略的收益来源，数据加载。AMP把GPU计算时间缩短之后，原来被掩盖的CPU数据加载瓶颈会暴露出来，你会看到GPU利用率忽高忽低。这时候把DataLoader的num_workers调大、开pin_memory=True，让GPU等数据的时间降到最低。很多人改完AMP发现吞吐没变化，十有八九是卡在这里。
4. 实战中踩过的坑与排查速查表
4.1 典型报错与定位逻辑
AMP的报错不算多，但每一个都很经典。最常见的报错是类型不匹配的错误，类似于expected scalar type Half but found Float。这种通常是你的自定义loss函数、数据增强逻辑或者某些不在autocast策略表里的算子，显式地接收了FP16输入，却做了FP32运算。排队思路很简单：把所有不在autocast覆盖范围内的自定义前向逻辑，显式.float()转回去，或者把整个自定义函数移到autocast块外面用FP32算，再把结果转成FP16传回模型。
另一种情况是Dataloader环节没处理好。AMP对输入tensor的类型不敏感，但如果你的数据管线里有人手动的.half()转换，可能导致BN层或自定义算子的输入和内部期望类型不一致。我建议的做法是：数据增强和预处理的每个环节都用FP32，只在进入模型前由autocast统一接管。不要手动提前转half，除非你清楚自己在干什么。
4.2 loss变NaN的排查思路
如果你发现训练到一半loss变NaN，先别急着怪AMP。排查顺序应该是：确认基线FP32训练是否本身就NaN，如果FP32也NaN，那是模型结构或学习率的问题，跟AMP无关。如果FP32正常、AMP后NaN，再看三个地方：学习率是否过大，AMP下建议先用原学习率的0.8-1.0倍起步，不要激进加；检查GradScaler是否在正常工作，当你发现loss出现inf时，scaler会自动调小scale并跳过该步，你可以打印scaler.get_scale()看它是不是一直在下降；检查用的是不是FP16做loss计算，某些自定义loss在FP16下会溢出，把loss计算强制留在FP32就解决了。
还有一个隐蔽问题，梯度裁剪的顺序。混合精度下，梯度是被scale过的，所以torch.nn.utils.clip_grad_norm_必须改成scaler.unscale_(optimizer)之后再做，否则裁剪阈值和实际梯度尺度不匹配，训练会变成“薛定谔的收敛”。PyTorch官方推荐写法是：
scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update()这一步很容易漏，但漏掉的后果很严重，可能直接精度崩坏。
4.3 性能没提升时的定位方法
同样一套AMP代码，有人提速2倍，有人感觉不到变化。吞吐没提升的几个主要原因，按概率排序：第一，batch size太小、GPU没有跑满，AMP省了计算量但GPU利用率本来就不高，收益自然不明显，解决方案是增大batch size；第二，模型里回退到FP32的算子占比太高，比如大量使用BatchNorm、Softmax、LayerNorm，这些算子在autocast下依然跑FP32，如果你的模型几乎全是这类算子，收益肯定有限；第三，数据加载和预处理是瓶颈，GPU在等CPU，AMP再怎么加速也白搭，此时优先处理DataLoader；第四，GPU本身不支持Tensor Core，比如P100以下的旧卡，收益基本为0。
判断瓶颈在哪，最直接的工具是torch.profiler，它能把每个算子的耗时拉出来，看FP16算子和FP32算子各占多少时间。或者用nvidia-smi dmon看GPU利用率，如果利用率长期低于70%，说明瓶颈在别处。经验法则：计算密集型任务、大batch训练、长序列任务收益最大，小模型、小batch、IO密集任务的收益会打折，这是正常现象，不代表AMP没用。
4.4 低显存场景的实际心得
最后说点实操体会。我经常在6G、8G显存的小卡上跑实验，这种环境里AMP几乎是救命级别的操作。很多人问“低显存到底能不能跑这个模型”，我的标准答案是：先用AMP跑一遍，再谈其他。比如微调7B级别的大模型，FP32下6G显存连基础的LoRA都跑不动，AMP配合LoRA后，前向和反向的计算量明显降低，显存从OOM边缘拉到能跑完一个step。网上很多人讨论“某某模型是不是要全部参数进显存”，这里补充一个知识点：MoE这类架构，显存瓶颈主要在参数和优化器状态，AMP只帮你缓解一部分，真正要降低参数显存还得靠量化、CPU offload或者LoRA这类参数高效微调。AMP擅长的是帮你把激活值和计算精度压下来，这套组合拳打下来，小显存卡也能做不少事情。
如果你准备在自己的项目里落地AMP，我的建议是把它当成默认配置而不是特殊操作。今天的新模型训练脚本，我几乎都直接带上AMP，只有当模型本身极小、或者数值极其敏感的时候才会关掉它。AMP不是银弹，但它是成本最低、见效最快的一档优化手段，值得长期放在你的工具箱里。
