显存不够、训练慢几乎是每个PyTorch玩家都会撞上的墙。我见过太多项目在batch size从32降到8之后精度没保住反而训崩了也见过有人为了跑一个稍大的模型把代码改得乱七八糟最后才发现问题出在数值精度上。PyTorch AMP混合精度训练就是解决这类问题的标准手段之一在不改变模型结构、不显著牺牲精度的前提下把显存占用降下来同时把训练吞吐提上去。这篇文章我会从底层原理讲到实战代码再到各种翻车现场的排查尽量把AMP这玩意讲透。文章适合谁如果你正在用PyTorch训练CV模型、NLP模型或者刚接触大模型微调被显存不足和训练速度折磨又或者你只是听说过混合精度、FP16、GradScaler这些词想系统搞明白这篇文章都能用得上。已经玩得很溜的大佬可以直接跳到第4节看踩坑记录。1. 混合精度到底在做什么先搞懂三个底层问题1.1 FP16为什么能省一半显存PyTorch默认用FP32单精度浮点数存储模型参数和中间计算结果每个数占4个字节。FP16半精度浮点数每个数只占2个字节显存占用直接砍半。举个例子一个7B参数的大模型FP32下光模型参数就要28GB左右显存换成FP16存储就是14GB这一下就差了14GB很多原本跑不动的模型就能塞进消费级显卡了。这里要注意一个细节训练时显存大头不只是模型参数还有优化器状态和中间激活值。优化器状态这一块Adam优化器默认会保存一阶动量FP32和二阶动量FP32加上参数本身的FP32副本总共是参数量的12倍字节开销。AMP混训通常采用的做法是模型权重保留一份FP32副本作为“主权重”前向和反向用FP16计算更新时用FP32主权重来做这样既能享受FP16的速度和显存优势又不会让权重更新精度崩掉。激活值的显存节省也很可观尤其是输入分辨率高、序列长度长的模型激活值往往比参数还占显存。1.2 混合精度的“混合”指什么很多新手会误解“混合精度”的意思以为就是把所有东西都转成FP16。不是的混合精度强调的是“混合”——该用FP16的地方用FP16该用FP32的地方必须保留FP32。为什么不能全FP16因为FP16的数值范围太窄了。FP32的动态范围大约是1e-38到3e38而FP16能表示的范围大约只有6e-5到65504小于6e-5的数会直接变成0下溢出大于65504的数会变成inf上溢出。深度学习训练过程中梯度值经常非常小比如1e-10这种量级如果全程FP16梯度直接归零模型根本没法收敛。所以我们需要一个机制计算主力用FP16但把容易出问题的环节梯度更新、损失计算、某些特殊op保留在FP32。这就是“混合”的真正含义。AMP靠一套自动机制来“翻译”每个操作对于已经注册了FP16内核的算子比如卷积、矩阵乘法、Embedding自动用FP16计算对于数值稳定性要求高的算子比如Softmax、LayerNorm、CrossEntropyLoss自动保持FP32。PyTorch自动选择你不需要手动改模型代码。1.3 速度提升来自哪里省显存是AMP的显性收益速度提升则是另一个核心卖点。当代GPUNVIDIA从Volta架构开始AMD从CDNA架构开始都内置了专用的FP16矩阵计算单元算力通常是FP32的2倍甚至更高。以NVIDIA A100为例FP32的峰值算力约19.5 TFLOPS而FP16 Tensor Core算力高达156 TFLOPS差了8倍。就算实际训练中吃不满峰值FP16的吞吐也明显高于FP32。还有一个容易被忽略的点FP16数据量减半意味着显存带宽的压力也减半。训练过程中参数、梯度、激活值在显存里来回读写数据量越小搬运速度越快这也是吞吐提升的重要来源。我自己实测过一个ResNet-50训练任务AMP开启后batch size翻倍单卡吞吐提升了35%到80%具体看模型和硬件。2. 核心组件与显存账本AMP三板斧拆开揉碎2.1 autocast前向传播的“动态翻译官”PyTorch AMP的核心API是torch.autocast。用法非常简洁在训练前向过程外包一层from torch.cuda.amp import autocast with autocast(device_typecuda, dtypetorch.float16): outputs model(inputs)关键是理解autocast的效果范围它只作用于前向传播和损失计算这个过程反向传播的梯度计算也会受益因为梯度是从FP16的中间结果反推出来的计算图里已经是FP16了。但autocast不会自动把模型参数转成FP16也不需要你手动转。模型参数依然是FP32只是在进入被优化的算子时AMP会为这个算子创建FP16的临时输入来计算计算完再释放。这里有个常见的认知误区有人以为用了autocast模型就变成“FP16模型”了。不是的。真正的“FP16模型”要额外调用model.half()这俩是完全不同的方案。model.half()是把所有参数永久存储为FP16简单粗暴但很容易因为数值溢出导致训练崩掉AMP方案则保留了FP32主权重数值稳定性好得多这也是AMP成为主流的原因。autocast还支持device_type参数在CUDA上训练用device_typecuda在CPU上也可以开启混合精度用device_typecpu但CPU的加速效果有限主要图个内存节省。2.2 GradScaler梯度缩放的守护者如果说autocast是“翻译官”那GradScaler就是“守护者”。它的核心任务只有一个防止梯度下溢出。前面说FP16能表示的最小正数是6e-5左右而训练中梯度值落到这个范围以下的概率非常高尤其是在模型深、数据噪声大的时候。一旦梯度变成0那部分参数永远得不到更新模型就像瘸了一条腿。GradScaler的解决方案很聪明在反向传播之前先把loss放大一个倍数比如1024倍这样反向传播算出来的梯度也会整体放大1024倍原本会下溢出的梯度值就被“顶”到了FP16的可表示范围内。梯度更新前再把梯度除以1024恢复原来的尺度。这个过程完全自动。为什么是动态缩放GradScaler会根据实际训练的梯度情况自动调整放大倍数。如果连续若干步都没有发生inf/NaN它会认为现在的倍数偏保守可以放大一点以提高精度如果发生了溢出它会调小倍数保证训练稳定。这个动态调整逻辑就是GradScaler.scale()和GradScaler.update()两个方法配合实现的。常见用法模板scaler torch.cuda.amp.GradScaler() for step, (inputs, targets) in enumerate(dataloader): optimizer.zero_grad() with autocast(device_typecuda): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()注意几个顺序上的坑scaler.scale(loss).backward()先放大loss再反向传播这一步必须放在autocast的上下文之外scaler.step(optimizer)内部会检查梯度是否溢出如果溢出就跳过这次更新所以optimizer.step()绝对不能自己再调用一次scaler.update()会在每次迭代末尾更新缩放系数。2.3 BN层与AMP的兼容性BatchNorm层在AMP下有些特殊值得单独拎出来说。BN层有batch维度上的统计计算在FP16下batch统计量的精度可能受影响所以AMP的autocast默认会让BN层保持在FP32下运行。这是PyTorch的默认行为通常不需要你干预。但你可能会遇到另一个问题在用torch.compile或某些第三方库时BN的表现会变得不可控。这时候建议把模型中的BN层手动换掉或者用原生实现。如果用的是SyncBN多卡同步BNAMP下保证多卡通信的精度一致性也得留意一下。另外如果你做的是NLP任务大部分模型没有BN基本不用担心这个问题。CV任务里带BN的模型多得留个心眼。3. 手把手实战从最小改动到完整训练流程3.1 最小改动方案5分钟让现有代码跑起来如果你有一个已经跑通的PyTorch训练脚本改成AMP版其实只需要动4个位置。先看最精简的改动对照改动前optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, targets) loss.backward() optimizer.step()改动后scaler torch.cuda.amp.GradScaler() optimizer.zero_grad() with torch.autocast(device_typecuda, dtypetorch.float16): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()就这几行其他部分都不用动。训练逻辑、dataloader、模型定义、验证逻辑全部保持原样。这种“外科手术式”的改动方式是我最推荐的起步方式——改得越少出问题的可能性越低。跑起来之后你可以观察两个指标显存占用是否下降、每秒处理的batch数是否提升。建议固定batch size做一轮对比再逐步增大batch size直到显存打满这样能把AMP的收益看清。3.2 完整训练循环模板可直接抄作业的版本下面给一个更完整的训练模板包含了梯度裁剪、学习率调度、日志记录和Checkpoint保存适合大部分训练任务直接改改就能用import torch import torch.nn as nn from torch.cuda.amp import autocast, GradScaler model MyModel().cuda() optimizer torch.optim.AdamW(model.parameters(), lr1e-4) criterion nn.CrossEntropyLoss() scaler GradScaler() scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxepochs) model.train() for epoch in range(epochs): for batch_idx, (inputs, targets) in enumerate(dataloader): inputs, targets inputs.cuda(), targets.cuda() optimizer.zero_grad() with autocast(device_typecuda): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() # 梯度裁剪需要先缩放 scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update() scheduler.step() if batch_idx % 50 0: print(fEpoch {epoch} | Batch {batch_idx} | Loss {loss.item():.4f}) state_dict { model: model.state_dict(), optimizer: optimizer.state_dict(), scaler: scaler.state_dict(), epoch: epoch, } torch.save(state_dict, fcheckpoint_epoch_{epoch}.pt)这里有两个细节要重点解释。第一个是梯度裁剪。AMP下梯度是被放大过scale的如果直接对model.parameters()的梯度做裁剪你裁剪的是放大后的梯度阈值就完全错乱了。正确做法是先调用scaler.unscale_(optimizer)把梯度恢复原尺度再做裁剪。很多新手在这里吃亏模型训练不稳定还以为AMP有问题其实是裁剪逻辑不对。第二个是Checkpoint保存。我建议把scaler的状态也保存下来这样断点续训时缩放系数是连续的。如果不保存重启后GradScaler又从头开始调整系数虽然最终也能收敛但前期会有一段适应期平白浪费一些训练时间。3.3 显存和吞吐的量化验证不对比不踏实光说“省显存、提吞吐”不够你得学会自己测量。我习惯在每个训练阶段都做三组对照实验FP32 baseline、AMP默认策略、AMP配合更大batch size。记录三个核心指标显存占用nvidia-smi的Memory Usage、每秒迭代数、有效吞吐每秒处理的样本数。以我自己跑过的一个BERT-base微调实验为例配置显存占用Batch Size每秒样本数FP3211.2GB32210AMP8.4GB32356AMP 加大batch10.8GB48392可以看到AMP在相同batch size下显存降了约25%吞吐提升约70%把省下来的显存拿来扩大batch size后一个8GB显存的卡就能跑到接近原11GB的batch规模吞吐进一步提升。这里要提醒一句实际收益因模型而异。卷积网络和Transformer的收益通常明显但有些算子是带宽密集型的FP16加速有限遇到瓶颈算子多的模型提升可能只有10%-20%这很正常不用失望。3.4 如何选择合适的初始缩放系数与参数GradScaler的init_scale参数默认是65536这个值适用于绝大多数情况。什么时候需要调整如果你的模型训练中频繁出现loss变成NaN可能是初始缩放太大导致上溢出可以把init_scale调低一点比如1024如果训练早期梯度太小导致loss长时间不变可以适当调高。另外两个GradScaler参数值得了解growth_factor和backoff_factor控制缩放系数调整步长默认分别是2.0和0.5一般不用动。growth_interval连续多少次迭代无溢出后就放大系数默认2000小数据集训练时建议调小到200加快系数收敛。这些参数本质上是经验值。我的建议是能不动就不动除非你明确观察到问题再逐个调整。4. 高频翻车现场AMP实战问题排查实录用了这么久的AMP我踩过的坑、帮别人排查过的问题汇总下来十个里有八个是下面这些。4.1 训练loss变成NaN或inf训练直接崩了怎么办这是最吓人也最常见的问题。按照经验优先级从高到低排查先看数据侧训练数据里有没有NaN值标签有没有异常很多输入预处理写的除法分母可能为0数据里混入NaN后loss算出NaN随后梯度全是NaN模型权重直接变成NaN。这锅不能甩给AMP先清洗数据试试。再看损失函数如果你自定义了loss检查里面有没有数值不稳定的操作比如直接对log概率做exp这种操作在FP32下也可能溢出只是FP16下更容易暴露。解决办法是把loss计算的数值敏感部分放到FP32通常做法是在模型输出后不要急着把logits转成FP16让loss的输入保持FP32。然后看缩放系数打印GradScaler当前的scale值如果scale一直在下降说明梯度频繁溢出可能是学习率偏大需要降低学习率或者调低init_scale。我遇到过一例模型是深层的Transformer初始学习率3e-4在FP32下没问题AMP下梯度上溢出频繁降到1e-4就稳定了。最后一个冷门检查点用的优化器是不是SGDSGD对梯度尺度敏感AMP下容易出现更新过大或过小的问题。换成AdamW这类自适应优化器会稳很多。4.2 收敛变慢、精度明显下降不一定是AMP的锅如果训练能跑通但loss不如FP32降得干净先别急着关AMP。优先怀疑两件事batch size变大带来的影响、学习率没重新调。AMP让你能用更大的batch size但batch size翻倍本身就会改变训练动力学。原来32的batch变成64收敛曲线不一样是正常的这属于“大batch训练”问题不是AMP的问题。此时应该考虑同步调整学习率比如线性缩放规则、warmup变长或者使用更大的学习率别一上来就怪AMP。如果同样的batch size下AMP的收敛确实变差了检查一下模型里的自定义算子有些第三方实现的自定义op不支持FP16autocast可能没覆盖到导致部分模块还在FP32反而拖慢速度。用profiler看一下各op的时间分布能找到瓶颈。还有一类任务是“精度特别敏感的”比如科学计算里某些回归任务需要精确到1e-6级别的loss这种场景AMP未必合适。但绝大多数视觉、NLP任务AMP都能做到精度几乎无损。4.3 多卡训练、分布式训练下AMP怎么配合DDP和AMP可以一起用但有一个明确的坑梯度同步的顺序。DDP在每次backward之后会自动做梯度all-reduce如果梯度还是放大状态就做all-reduce通信精度会出问题而且每个rank的scale可能不一样理论上同步训练下所有rank的scale应该一致但保险起见还是要小心。标准做法是在DDP中结合torch.nn.parameters和scaler时确保每个rank都调用scaler.step(optimizer)并且用model.no_sync()来控制梯度同步。另外如果是多机训练每台机器上GradScaler的更新要一致这样梯度同步才有意义。实际项目里我通常会把DDP的梯度平均拿掉因为batch变大了梯度平均的影响不大或者在backward后手动做一次all-reducescaler.scale(loss).backward() if is_distributed: for p in model.parameters(): if p.grad is not None: dist.all_reduce(p.grad, opdist.ReduceOp.SUM) p.grad.div_(world_size) scaler.step(optimizer) scaler.update()4.4 显存没降反而爆了检查你的设置用了AMP后显存一点没降甚至更高这种情况不多但确实存在原因通常是以下几个。第一模型里还有大量参数以FP32形式常驻显存。AMP模式下PyTorch会保留FP32主权重这是合理的保存策略它本身就会吃掉和模型参数等量的显存。如果你的模型里有一些未注册到autocast范围的模块比如自己手写的某些自定义层它们依然会以FP32存储中间结果显存优势就被抵消了。第二你用了torch.no_grad包住验证流程但忘了关闭推理模式的缓存。验证时虽然不存梯度但某些算子的workspace还是会被缓存长期下来显存碎片化。建议验证时也使用autocast保持一致的计算路径。第三也是最容易被忽略的——你开着CUDA graph或某些jitable优化时显存占用会额外预留缓冲区。使用torch.compile配合AMP时要留一些显存余量有时候编译过程本身就会把显存吃掉几百MB到几个GB。4.5 低显存场景与大模型微调显存不够时还能做什么这里结合一个很多人问的问题低显存显卡比如6GB、8GB能不能跑大模型微调AMP是第一步能把模型从“跑不动”变成“勉强跑”但还不够。你还需要配合几个手段。梯度累积是必须的。AMP帮你把batch size翻倍但如果显存极限只支持batch size4你再怎么省也变不出batch size16。梯度累积可以模拟大batch的效果代价只是训练时间变长。代码样例accumulation_steps 4 optimizer.zero_grad() for micro_step in range(accumulation_steps): outputs model(inputs[ micro_step]) loss criterion(outputs, targets[micro_step]) scaler.scale(loss).backward() if (step 1) % accumulation_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()混合精度梯度累积gradient checkpointing三件套组合下来8GB显存跑一个中等规模的Transformer微调是没问题的。另一个常见问题是“如何让模型在推理时预留显存”——这个问题在ComfyUI、WebUI这些推理场景里特别常见。核心思路和训练不同推理不需要保存梯度也不需要优化器状态所以显存大头是模型本身和激活值。做法包括用model.half()把模型参数转成FP16加载推理场景直接用half没问题不会有梯度溢出风险、开启torch.inference_mode()、手动调用torch.cuda.empty_cache()清理碎片、为其他程序预留显存可以设置环境变量。4.6 关于“让显卡调用内存做显存扩充”的提醒热度很高的“让显卡调用内存做显存扩充”本质上是显存耗尽后的次优解具体机制是利用统一内存或手写tensor offload。但我要说句实在话代价往往是训练速度掉一个数量级因为PCIe带宽只有显存带宽的几十分之一频繁来回搬运数据会成为绝对瓶颈。AMP能显著减少你需要搬运的数据量所以如果你万不得已必须用这个方案请先开AMP把数据量压到最小再配合梯度累积减少同步频率。优先级顺序永远是AMP - 梯度累积 - gradient checkpointing - offload最后一步是实在没招才用的。5. 进阶视角大模型时代的AMP与显存新解法5.1 MoE架构、激活参数与显存的关系如果你关注大模型应用一定听过“总参数和激活参数对显存的影响”这个话题。有一类架构叫MoE混合专家把模型拆成多个专家子网络每个token只激活其中的一部分专家。总参数可能几千亿但每次前向用到的参数激活参数只有几十亿。那么问题来了MoE架构需要全部参数进显存吗严格来说推理时如果使用离线量化、逐层加载不一定全部常驻显存但训练时由于反向传播需要计算每个专家参数的梯度理论上所有被访问过的专家都要参与更新因此训练场景下大部分专家参数还是要在显存或CPU内存中可访问。AMP在这里的价值是让单位显存能装下更多参数同时FP16 Tensor Core让矩阵乘更快这对MoE这种大计算量架构特别友好。对普通Transformer显存可以简单拆成三块账参数占的显存、优化器状态占的显存、激活值占的显存。AMP优化的是前两块参数以FP16参与计算 部分激活值。如果你用的是Adam优化器FP32下优化器状态占12字节/参数这是最大的一块。参数量越大AMP配合更高效的优化器如8bit AdamW带来的收益就越明显。5.2 量化、FP16推理与部署实践训练用AMP推理部署可以走更远的路线。AMP训练得到的模型权重虽然保留着FP32主权重但推理时完全可以加载FP16权重显存立刻减半速度还快。更极致的是INT8量化把权重压到1字节/参数但量化要做校准否则精度损失明显。这里也要提一个热词很多人在部署场景问“能不能用FP16跑全精度模型”。答案是可以的但注意区分两种做法的区别。model.half()是简单粗暴地把模型全转FP16适合纯推理而AMP是训练时期的机制。两者目标不同不要混用。部署时的推荐检查顺序是先试model.half()推理精度满足就直接用精度不够再考虑量化校准或部分层保持FP32。5.3 环境搭建时的几个常见坑在Windows的WSL2下跑PyTorch GPU版本很多人装完发现CUDA不可用DistributedDataParallel起不来。我建议优先用官方pytorch/pytorch镜像或者在WSL2里通过conda安装CUDA版本的PyTorch确认torch.cuda.is_available()返回True再开AMP。AMD显卡在WSL2下也有对应的ROCm方案但兼容性比NVIDIA差一些AMP的Tensor Core加速可能不完整需要实测确认收益。安装PyTorch时CUDA版本不匹配是最经典的翻车现场你装了一个CUDA 12.1的PyTorch但本机驱动只支持CUDA 11.8很矛盾建议统一用PyTorch官方站给的conda命令安装同时用nvidia-smi确认驱动版本支持的CUDA版本。环境变量可以设置export CUDA_VISIBLE_DEVICES0 export PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:256第二个环境变量对显存碎片化有缓解作用配合AMP使用显存利用率能再高一点。写在最后的个人经验AMP用了这几年我最深的体会是它不是一个“开关”而是一套需要理解的机制。上来就改代码、不管原因遇到NaN就慌看到显存没降就骂都会错过真正的坑位。我建议所有刚接触AMP的人第一次跑通后别急着加batch size先把scaler.get_scale()、loss.item()、显存占用打印出来跑几十步对比FP32和AMP的曲线体会到数值稳定性的差别再往后推进。如果你现在正好在低显存显卡上挣扎我的建议路径是先用AMP看显存余量再开梯度累积看收敛情况然后开gradient checkpointing看速度掉多少最后才考虑offload或者换卡。每一步都有清晰的代价和收益别跳着来。最后分享一个小的实操技巧验证模式下也别忘了用autocast只开训练不开验证会导致推理时的中间结果精度路径和训练时不一致有些模型会出现“训练好好的验证结果波动巨大”的灵异现象。把训练和验证的精度策略统一是你跑AMP项目时最值得养成的习惯。