PyTorch AMP混合精度训练实战：省显存、加速与踩坑指南

做深度学习训练尤其是大模型微调或者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权重占400MBFP16权重只占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不是把所有计算都切成FP16PyTorch从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_typecuda, dtypetorch.float16)和torch.amp.GradScaler(cuda)。功能上两者等价旧写法也还能用但新项目建议直接用新API因为后续维护和扩展都以新API为准。设备支持上也值得说清楚。AMP的收益建立在GPU硬件支持Tensor Core的基础上NVIDIA从Volta架构V100开始支持FP16 Tensor CoreTuringT4/RTX 20系、AmpereA100/RTX 30系、Ada LovelaceRTX 40系都支持得很好所以只要你不是老掉牙的GPU都能吃到红利。AMD ROCm环境下PyTorch的AMP也做了适配但稳定性和性能优化程度不如NVIDIA平台这一点要有心理预期。如果你在CPU上训练就别指望AMP了CPU端的half运算没有专门加速单元强行.half()反而更慢CPU场景一般用BF16PyTorch新版也支持torch.autocast(cpu, dtypetorch.bfloat16)但收益主要在内存占用而不是计算速度。3. 从零把训练脚本改成AMP的完整实操3.1 改造前先量化基线我见过太多人一上来就改代码改完发现“好像快了一点”但又说不清快了多少显存省了也没记录下来最后根本没法判断AMP到底值不值得用。正确做法是先跑一次基线把下面三个数字记下来单次epoch耗时、峰值显存、每秒训练样本数。写一个简单的benchmark脚本用torch.cuda.max_memory_allocated()统计显存峰值用时间戳统计吞吐不需要额外工具import torch import time def benchmark(model, loader, device, num_steps50): 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(fthroughput: {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(), lr1e-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_typecuda, dtypetorch.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微调大模型时就是“AMP8bit AdamW梯度累积”三个一起上效果立竿见影。还有一个容易忽略的收益来源数据加载。AMP把GPU计算时间缩短之后原来被掩盖的CPU数据加载瓶颈会暴露出来你会看到GPU利用率忽高忽低。这时候把DataLoader的num_workers调大、开pin_memoryTrue让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_norm1.0) scaler.step(optimizer) scaler.update()这一步很容易漏但漏掉的后果很严重可能直接精度崩坏。4.3 性能没提升时的定位方法同样一套AMP代码有人提速2倍有人感觉不到变化。吞吐没提升的几个主要原因按概率排序第一batch size太小、GPU没有跑满AMP省了计算量但GPU利用率本来就不高收益自然不明显解决方案是增大batch size第二模型里回退到FP32的算子占比太高比如大量使用BatchNorm、Softmax、LayerNorm这些算子在autocast下依然跑FP32如果你的模型几乎全是这类算子收益肯定有限第三数据加载和预处理是瓶颈GPU在等CPUAMP再怎么加速也白搭此时优先处理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不是银弹但它是成本最低、见效最快的一档优化手段值得长期放在你的工具箱里。