训模型训到显存告急、一个batch都塞不进去的时候大多数人第一反应是换更大的卡或者砍batch size。但我个人更推荐你先花半小时试试自动混合精度AMP——它不像换硬件那样要花钱也不会像砍batch size那样明显影响收敛而且能同时解决两个痛点显存占用和训练速度。自动混合精度说白了就是给训练过程配一个“动态调度员”让矩阵乘法这类吃计算量的算子跑在FP16上同时把权重更新、归一化统计这类对数值敏感的部分留在FP32。它的核心收益来自NVIDIA从Volta架构开始引入的Tensor Core以及半精度数据天然少一半的显存占用。这篇文章我会把AMP的原理、PyTorch和Apex的工程落地、以及我在训练过程中踩过的坑完整梳理一遍适合正在用PyTorch训练CV或NLP模型、想提速但不太确定混合精度会不会影响收敛的读者。1. 半精度训练的收益到底是什么量级先聊一个很多人忽略的事实自动混合精度不是让整个模型都跑在FP16里而是在同一份训练中同时使用FP16和FP32。之所以能提速主要是靠硬件里的Tensor Core。Tensor Core是NVIDIA GPU里专门做低精度矩阵乘法的单元它的计算吞吐量远高于普通FP32核心。从Volta之后的V100到Turing的T4、RTX 20系列再到Ampere的A100、RTX 30/40系列都有Tensor Core。不过注意只有Ampere架构和更新的硬件支持BF16V100和Turing虽然也能跑AMP但对BF16的支持并不完整需要用FP16混。从实测数据看收益大致可以分为三块收益项典型幅度原因训练吞吐量提升1.5-2.5倍矩阵乘等核心算子转入Tensor Core FP16路径时钟周期和带宽压力同时下降显存占用显存峰值下降约30%-50%激活值、部分梯度以FP16存储参数量大时收益更明显通信开销多卡场景尤其是数据并行下明显变小FP16梯度传输量是FP32的一半这里最容易踩的误区是大家在对比速度时只看“有没有开AMP”却没注意batch size是否保持一致。开AMP后显存省下来了很多人顺手把batch size翻倍速度收益和收敛效果都会和原来不具可比性。我自己实测过一个ResNet-50训练任务在同样batch size下AMP让每个step时间从0.42秒降到0.26秒吞吐提升大约1.6倍而如果开了AMP之后把batch size抬到显存允许的上限吞吐可以到2倍以上但BN的统计量、学习率对应的batch规模全都变了需要重新调参。还有一类任务要单独说瓶颈不在矩阵乘法的任务比如大量小算子、频繁的CPU-GPU数据拷贝、小batch的推理式训练AMP的收益会明显缩水。原因是这些任务本身没有足够的Tensor Core计算密集段可供加速省下的只是显存带宽。所以做目标检测里的一些后处理逻辑、或者强化学习里频繁交互采样时别指望AMP能救一切显存收益依然在但时间收益可能只有10%-20%。对大多数CV和NLP的训练任务我的判断是第一个晚上先跑一版AMP能接受精度小幅波动就长期用如果精度掉了再按后面的排查思路看是哪里出了问题而不是马上关掉AMP回到FP32。2. FP16 的两大陷阱范围小与精度低既然AMP能把速度提起来那为什么不能直接全模型用FP16训因为FP16有两个非常实际的问题这也是混合精度而不是纯半精度存在的根本原因。2.1 数值范围问题梯度过小直接变0FP16的数值结构是1位符号位、5位指数位、10位尾数位。它能表示的最大值大约是65504这个很好理解但很多人忽略的是最小正常值大约6.1e-5。这比FP32的最小正常值约1.18e-38小了许多个数量级。这意味着什么在深网络里反向传播算出的梯度经常远小于1e-4尤其越靠近输入层经过链式法则连乘之后数值越小。当梯度小到FP16可表示范围之外时它会直接下溢成0。梯度变0那一层的权重就再也不更新了——而且这不是收敛是“假死”loss看上去还在降因为浅层大参数还在动但深层的特征提取能力已经冻结了。解决方案就是后面要讲到的损失缩放在反向传播前把loss放大等价于把梯度放大让它落在FP16的合理范围内然后在更新权重前缩回去。2.2 精度不足问题损失对更新量不敏感另一个问题是FP16只有10位尾数也就是大约3位十进制有效数字。这带来的后果不是梯度直接变0而是“精度没有足够分辨力去呈现非常小的变化”。举个例子。一个参数的当前值是1.0梯度算下来是1e-4理论上更新量是1e-4量级。可当权重和更新量都要放进FP16时1.0附近的最小分辨率约是1e-31e-4的增量在FP16视角里根本不足以改变尾数于是权重保持原样。训练后期loss下降缓慢甚至停滞常常就是这个原因。有人会问那更新量都这么小了FP32不也差不多吗差别在于FP32的尾数有23位1.0附近的分辨率约是1.2e-7完全可以容纳1e-4这个量级的变化。FP16做不到。2.3 所以“混合”才成立顺着这两条思路就会明白计算密集、对精度相对宽容的算子典型是卷积、矩阵乘放到FP16能拿到吞吐和显存收益而权重更新、梯度累积、BatchNorm的统计量这些“数值安全区”留在FP32用来保住收敛质量。自动混合精度的“自动”二字核心就是自动做这个分配不需要你手工指定每个算子用哪种精度。3. 自动混合精度的三大保命机制理解AMP不能只看API。它在设计上围绕上面两个陷阱内置了三套机制主权重副本、损失缩放、算子精度分配。我把它们拆开讲。3.1 主权重副本权重更新永远在FP32这是最容易被忽略的一点。很多人以为AMP只是“把模型参数转成FP16继续训”实际上训练过程中的权重更新一直保留一份FP32副本。具体流程是将FP32权重作为主副本前向时将权重临时转为FP16参与计算反向计算出FP16梯度后转换回FP32FP32主权重副本接收FP32梯度并完成更新下一轮前向再把更新后的FP32权重转成FP16。这样做的好处是更新精度不被FP16限制累积误差也不会越滚越大。在代码层面PyTorch原生AMP并不是改参数dtype而是在前向时用autocast上下文自动把部分算子包装成低精度唯一真正“常驻FP16”的是参与计算的临时张量。3.2 损失缩放把梯度放大到安全区间前面说FP16的梯度下溢问题损失缩放就是应对方案。思路非常朴素既然梯度太小容易落到FP16的不可表示区间那就把loss在反向传播前乘以一个缩放因子比如1024相当于把梯度整体放大让它们回到FP16能表达的范围反向传播完成后真正更新前再除以这个缩放因子保证数学上等价。这里的难点是因子选多大。选小了照样下溢选大了梯度放大后可能超过FP16的上限65504变成inf。所以现代AMP都实现了动态缩放每N步比如PyTorch里的每隔2000步尝试把缩放因子调大一些如果连续很多步都没出现inf/NaN说明还有放大空间就把scale翻倍一旦某一步检测到梯度或loss异常自动把scale除以2或降得更多然后跳过这一步的权重更新。这个机制虽然简单但非常有效。PyTorch的GradScaler默认初始scale为2的16次方即65536之所以这么大就是考虑到要吞下极端梯度范围。实际训练里最常见的异常就是loss突然变成NaN尤其前几步排查时第一件事就是看scaler状态和scale值是不是已经掉下去了。3.3 算子精度分配黑白名单机制AMP并非把每个算子都塞进FP16。它内部维护了一套策略基本上分三类类别典型算子策略白名单强制FP16conv、matmul、linear、conv_transpose优先用Tensor Core FP16路径灰名单自动按输入决定elementwise类、softmax、sigmoid、tanh输入是FP16则输出FP16保持链条连续黑名单强制FP32BatchNorm、LayerNorm、损失函数相关计算数值稳定性优先保持FP32BatchNorm强制FP32的原因很容易理解它对batch内统计量敏感而且每层都在做归一化一旦统计量精度受损整个网络的数值稳定性会迅速恶化。LayerNorm同理Transformer模型在使用AMP时LayerNorm始终在FP32里跑。另外即使某个算子本身在白名单里如果输入张量是FP32AMP也不会强制转FP16而是尽量保持。这也是为什么有些算子“漏”回FP32会影响整体收益——一个FP32输入连着一串算子后面所有算子都可能跟着留在FP32。实践中如果发现AMP收益不如预期优先检查是哪个环节引入了FP32张量比如数据集加载时就转成了FP32并且没在合适的位置切回。4. 工程落地PyTorch原生AMP与Apex怎么选技术原理说清楚了接下来是实操。目前最常见的两套方案是PyTorch原生APItorch.cuda.amp和NVIDIA Apex。4.1 PyTorch原生AMP的正确用法PyTorch从1.6开始把AMP作为正式特性核心就两个东西torch.autocast和torch.cuda.amp.GradScaler。先给一份可以直接改到自己训练脚本里的最小模板import torch model MyModel().cuda() optimizer torch.optim.AdamW(model.parameters(), lr1e-4) scaler torch.cuda.amp.GradScaler() for epoch in range(epochs): for batch in dataloader: inputs, targets batch inputs inputs.cuda() optimizer.zero_grad() # 前向和反向都在 autocast 上下文里自动做精度分配 with torch.autocast(device_typecuda, dtypetorch.float16): loss model(inputs, targets) # 反向传播时先做 loss 缩放 scaler.scale(loss).backward() # 梯度裁剪要注意先 unscale scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # 更新权重内部会处理 scale 的还原 scaler.step(optimizer) scaler.update() if scaler.is_enabled(): cur_scale scaler.get_scale() # 这里可以加日志观察缩放因子的变化趋势这里的顺序非常关键必须用scaler.scale(loss).backward()而不是直接loss.backward()否则损失缩放不会生效。梯度裁剪前要先scaler.unscale_(optimizer)因为此时优化器内部的梯度还是被缩放过的直接按阈值裁剪会在错误尺度上操作。还有一个容易踩的细节如果某一步检测到inf/NaNscaler.step(optimizer)会自动跳过这步的权重更新所以你不必手动判断梯度是否合法。有些读者会发现scaler.scale(loss)之后loss.backward()里计算出来的梯度已经是scale过的FP32梯度那为什么还需要unscale_因为在scaler.step之前梯度以缩放形式存在参数上裁剪是数值操作必须用真实尺度。另外新版PyTorch还提供了torch.autocast装饰器方式可以只包model的forward部分其他代码留在FP32torch.autocast(device_typecuda, dtypetorch.float16) def forward(self, x): return self.net(x)这个用法的好处是侵入小、改起来快缺点是不好精细控制loss计算等位置的精度。论推荐程度我在项目里更建议用显式的with语句能看清每个计算段的精度策略。4.2 Apex 的 O1、O2 到底差在哪NVIDIA Apex是老牌混合精度库虽然PyTorch 1.6之后原生AMP已经够用但很多老项目还挂着Apex而且有些卷积算子Apex的实现确实更激进一些。Apex最常见的用法是from apex import amp model, optimizer amp.initialize(model, optimizer, opt_levelO1)opt_level主要有两档O1相当于“自动白名单”在FP32模型基础上按算子类型自动把合适的算子转成FP16其余保留FP32。和PyTorch原生AMP的自动混合非常接近也是我推荐默认使用的档位。O2更激进的直接模式模型权重本身先转成FP16再专门保留FP32副本用于更新。这和PyTorch的master weight机制在操作层面更接近“真BF16替代前的全半精度”方案数值风险更高收益也高一些但一旦代码里有对权重的自定义操作很容易出现精度翻车。选型上我的看法很直接新项目无脑用PyTorch原生AMP。原因不只是Apex的维护频率原生AMP和PyTorch新版融合得更好尤其torch.compile、分布式训练上原生API的兼容性少很多麻烦。老项目已经在跑Apex且稳定没有迁移必要但如果要升级PyTorch大版本或切换到DDP优先考虑替换成原生AMP减少一个依赖层面的变量。4.3 多卡训练时的AMP注意项用DistributedDataParallel时有一点比单卡更容易被忽略梯度同步。DDP在做all-reduce时通信的是FP32梯度而AMP在反向传播时产生的是FP16梯度这中间会发生一次精度转换。如果不做任何处理通信开销和精度都会出问题。PyTorch的处理是在DDP初始化时检测autocast状态自动把梯度转成FP32再同步。但如果你手动在模型外部另外做梯度的自定义通信比如一些模型并行代码就要自己确认通信的dtype。我的经验是多卡时先把单卡AMP跑通再切DDP否则很难判断问题是来自分布式还是来自精度。5. 训练过程中我踩过的坑与排查链路AMP本身不难用真正麻烦的是它把“数值异常”的排查范围扩大了。下面是几个我实际遇到过、并且有清晰根因的问题。5.1 loss突然变成NaN第一反应先查scaler我曾经在一个语音模型训练里遇到训练到第4000步loss突然NaN。关掉AMP之后就恢复正常于是第一反应是AMP导致了数值不稳定。排查后发现在我的数据管线里有个操作会把inputs里的padding部分置成极大的负数用于mask attention。这些负数在FP16里其实没问题但经过一个自己实现的指数运算后中间结果超过了65504直接溢出为inf反向时inf传回梯度loss就NaN了。所以不要一见到NaN就甩锅给AMP。排查链路应该是看scaler的scale值变化如果scale在NaN前已经下降过说明前面出现过梯度异常问题在更早的step定位第一个出现inf/NaN的内部张量用detect_anomaly虽然慢但有效或者手动在几个关键层后打印梯度范数看是不是自己实现的损失函数里出现了非白名单算子自定义损失部分经常在FP32和FP16之间来回转这恰恰是数值不稳定的高发区确认数据预处理没有产生超大值输入。5.2 收敛变慢但loss不爆多半是缩放策略太保守NaN是一类问题另一类是训练在合理范围内但收敛速度比FP32慢。这种情况我遇到最多的是损失本身非常小比如一些辅助损失只有1e-3量级初始缩放因子虽然很大但动态调整没有跟上。这时候可以从两个方向调整把GradScaler的init_scale调大一些比如从65536改成2的20次方适当减少scaler.update()的检查间隔让scale更快上升。还有一个隐藏问题如果自己改了损失权重把loss乘了一个常数比如loss 0.1 * main_loss aux_loss那缩放因子整体作用在混合后的loss上各个损失分量的梯度占比会变。PyTorch的GradScaler是对整体loss进行缩放这在数学上等价于对所有梯度乘同一个常数不同分量的相对比值不会变所以理论上没问题。但实践中损失跨度过大的时候比如两个loss差了100倍小loss的梯度很容易被压到FP16的不可表示区即便整体缩放也只是同时放大两者而已。我处理这类问题时会在梯度裁剪前看两个分量的梯度范数是否在一个数量级如果差太多先调整损失权重而不是死磕scale。5.3 显存没省多少检查动态图和临时张量有读者开AMP后看显存只降了10%远低于预期。常见原因有三个输入图像直接以FP32放进GPU并在autocast外做了归一化导致后面一串算子的输入都是FP32AMP的策略不强制转换白名单失效模型里有自定义算子没有注册AMP的custom_fwd装饰器在autocast上下文外执行FP32路径显存大头在优化器状态AdamW本身就是FP32状态参数多时优化器状态的显存不会因为AMP减少。对第一点把输入归一化放进autocast上下文内或者用非阻塞转换提前把输入转成合适dtype能让收益恢复不少。对第三点如果你用的是Adam系优化器它的fp32动量状态不会随AMP缩小省显存的期望要合理——通常省的是激活值缓存和一部分梯度不是优化器状态。5.4 用DDP时每个进程各存一个scaler别共享DDP下每个进程独立计算自己的梯度并同步后各自调用scaler.step()。这三种状态必须每个进程独立维护scaler、optimizer和loss。关键点是DDP会做梯度同步但scaler不会同步scale值每个进程根据自己看到的那一份loss做异常检测。当batch在数据层面被shard后各进程的loss略有差异scale的升降节奏可能不完全一致。这通常不影响训练但会让日志里的scale值看起来有细微抖动不必担心。真正的坑是我见过有人把scaler也放进DDP module里想一起广播结果scaler的scale参数被当作模型参数同步相当于给优化器的scale值加了隐含的约束非常诡异。scaler、optimizer和loss这三种状态必须每个进程独立维护。6. 哪些任务适合吃下AMP红利哪些不划算最后给一些判断标准帮你自己评估要不要上AMP。适合上AMP的任务一般有三个特征模型足够深或宽、计算密度高、batch内样本数多。典型是图像分类、目标检测、Transformer类NLP模型、语音识别中的特征序列模型。这些任务里矩阵乘和卷积占比高Tensor Core利用率直接决定迭代速度。AMP带来的收益几乎是无脑的显存省、速度提、精度波动可以通过损失缩放和主权重控制住。不太适合的任务通常是这三类模型极小比如单层线性网络计算密集度低开AMP反而增加自动类型转换的开销大量自定义数值操作在CPU和GPU之间来回搬数据速度瓶颈不在Tensor Core对精度极度敏感的研究任务比如某些科学计算式的训练需要FP32的位级可复现性。注意AMP在不同GPU型号上的数学结果不完全一样跨硬件复现实验时如果对确定性要求很高建议显式关闭AMP再做最终验证。从工程角度看我的习惯是把AMP作为默认选项跑一轮哪怕是一个很小的subnet先跑几十步用tensorboard看loss曲线是否和FP32基线基本重合再决定全量训练用哪套配置。这个AB测试成本很低省掉后面调NaN的麻烦。还有一个经验之谈混合精度的收益在较长序列、大batch的NLP任务上体现得比CV更明显。因为Transformer的计算瓶颈几乎全在FP16可以覆盖的矩阵乘法上LayerNorm、softmax这些FP32算子只占很小比重所以BERT、GPT类模型用AMP非常划算。说到底自动混合精度不是玄学也不是黑科技它是围绕“精度-速度-显存”三个维度的工程折中。理解FP16的数值边界、掌握GradScaler的缩放机制、知道每个算子在混合精度策略中的位置比背API重要得多。只要抓住这三条AMP在绝大多数深度学习项目里都能稳定地帮你省时间、省显存。