FP16混合精度训练实战：显存减半、训练翻倍的完整指南

很多人在GPU服务器上跑AI模型训练，最先撞上的墙往往是显存不够用，或者是训练速度慢到让人怀疑人生。我之前在调一个语义分割模型的时候，输入分辨率稍微调大一点，24G的显存就直接OOM，换batch size、改模型结构折腾了好几天。后来真正认真研究了FP16精度训练，才发现显存和速度的问题能一起解决一大半——这项技术用好了，内存占用几乎减半，计算速度在某些GPU上还能翻倍。这篇文章我不打算讲那种泛泛而谈的概念，而是把我从原理理解到落地实操、再到处处踩坑的完整过程梳理出来，让正在被显存和算力问题卡住的你，能少走很多弯路。

这篇内容适合的人群很明确：正在用GPU服务器做深度学习训练，发现显存不够或者训练太慢的人；对FP16、混合精度只有模糊印象但不知道怎么落地的人；以及在尝试FP16后遇到Loss不降、数值溢出等问题的调参选手。我会从为什么需要FP16、FP16的位级原理、硬件适配、PyTorch代码改造、踩坑排查到BF16和TF32的选型边界，给出一整套可以直接参考的实践方案。

1. 为什么非要用FP16：显存瓶颈与算力浪费的双重困境

1.1 显存不够用的真实场景

显存不够用这件事，几乎每个做CV或者NLP训练的人都遇到过。以YOLOv8或者nnU-Net这类常用模型为例，想要提升精度就要加大输入分辨率、扩大batch size，但这些操作直接会推高显存占用。很多人第一反应是换更大显存的卡，或者去搞模型并行、梯度累积这类复杂方案。但实际上，在模型训练过程中，显存消耗的大头不只是模型参数本身，还包括每一层的激活值、梯度、优化器状态等副本。FP32训练时，一份权重存4个字节，Adam优化器还要额外保存一阶动量、二阶动量，再算上梯度，一个参数可能要占16字节甚至更多。

举个例子，一个参数量2.5亿的BERT-base模型，FP32训练时模型权重本身只要1GB左右，但加上激活值、梯度和优化器状态，实际占用轻松超过10GB。如果把权重、激活值、梯度都改成FP16存储，模型权重直接变成500MB，激活值也接近减半。这就是FP16最直观的收益：在同样的显存限制下，batch size可以翻倍，或者输入分辨率可以提得更高，模型收敛效果自然更好。

1.2 算力不对称：很多GPU的FP16速度远超FP32

显存问题只是冰山一角，更可惜的是算力浪费。NVIDIA从Volta架构开始引入Tensor Core，专门为FP16矩阵运算设计了加速单元。以V100为例，FP32的浮点算力大概15.7 TFLOPS，但FP16 Tensor Core算力能到125 TFLOPS左右。A100上FP32是19.5 TFLOPS，FP16 Tensor Core则高达312 TFLOPS。换句话说，如果你一直用FP32训练，实际吃到的算力可能只有显卡理论峰值的一小部分。

不用Tensor Core的FP32计算，就好像开着一辆带涡轮增压的车却一直踩着半油门。PyTorch默认的训练路径是FP32走CUDA Core，完全用不上Tensor Core的加速能力。而一旦切换到混合精度，矩阵乘法和卷积这些最耗时的操作就会自动落在Tensor Core上，训练速度的提升非常明显。我实测过在A100上跑一个ResNet-50图像分类任务，同样epoch数，混合精度训练比FP32快了接近2.5倍，这个数据在batch size较大的时候尤其明显。

1.3 明确目标：FP16到底能省多少、快多少

在动手之前，建议先明确一下自己的收益预期，避免后面调参时没有方向。如果你在NVIDIA T4、V100、A100、H100或者RTX 30/40系显卡上跑常规的CV、NLP模型，正常情况下显存占用可以减少40%到50%，训练吞吐能提升1.5到3倍不等。如果你用的是比较老的GPU，比如GTX 10系这种不支持Tensor Core的卡，那么FP16虽然也能省显存，但计算加速效果非常有限，甚至还可能因为数值问题变得不稳定。

不过要注意，FP16不是免费的午餐。精度降低会带来数值范围收窄和舍入误差放大，这就需要配合混合精度技术（AMP）和损失缩放（Loss Scaling）来兜底。我不会一上来就把所有代码丢出来，先把原理讲清楚，你再去看代码就不容易踩坑了。

2. FP16不是“减半精度”那么简单：浮点数的位级真相

2.1 从FP32到FP16：每一比特都花在哪里

很多同学以为FP16就是把FP32的后16位砍掉，这其实是个误解。要理解FP16，得先看清浮点数的结构。无论是FP32还是FP16，都由三部分组成：符号位、指数位、尾数位。

FP32是1位符号位加8位指数位加23位尾数位，总共32位；FP16是1位符号位加5位指数位加10位尾数位，总共16位。指数位决定了能表示的数字范围，尾数位决定了精度。所以FP16和FP32并不是简单截断关系，而是同时压缩了指数和尾数的空间。这样带来的后果是：FP16能表示的最大值是65504，最小正常数大概是6.1e-5，而FP32能表示的最大值约为3.4e38，最小正常数约为1.18e-38。

表格对比会更直观：

2.2 动态范围与精度：一个“量程”一个“刻度”

动态范围和精度是两个容易混淆的概念。动态范围可以理解成一把尺子能测量的最大长度，精度则是这把尺子上的最小刻度。FP16的指数位只有5位，可以形象地想象成这把尺子很短，超过65504就会“爆表”；而尾数位只有10位，刻度比较粗，很多小数值会被四舍五入到同一个刻度上。

深度学习训练里最怕的两件事，一个是上溢，一个是下溢。上溢就是数值超过了FP16能表示的最大范围65504，直接变成无穷大，这在反向传播计算梯度时很常见。下溢就是数值小到低于最小正常数，直接被记成0，导致梯度消失。权重值往往在0.1这个量级，看起来离65504很远，但梯度计算过程中会有大量连乘，特别是在深层网络中，梯度的数量级可能跌到1e-4甚至更小。这也是为什么直接拿FP16训练，很容易出现梯度消失、Loss不再下降。

2.3 为什么梯度是最容易“翻车”的地方

如果只把模型权重改成FP16，前向传播一般还能跑，但反向传播的梯度就很容易出问题。训练初期权重更新幅度大，梯度值可能很大，一不小心就超过FP16的上限变成NaN；训练后期梯度越来越小，又容易低于FP16的下限直接变0。这两种情况都会让模型优化陷入停滞。

因此业界的标准做法并不是全量FP16，而是混合精度：前向传播和反向传播的矩阵运算用FP16计算，但主权重、梯度累加、优化器状态保留在FP32。另外还需要一个“损失缩放”机制，在反向传播之前把Loss乘以一个较大的常数，比如1024或65536，把梯度整体抬到FP16的可表示范围内，更新权重前再缩回来。这套机制就是AMP的核心，后面第四章我会详细讲代码实现。

3. 动手前先看硬件：GPU的FP16算力差异与选型对照

3.1 Tensor Core与老显卡的“假FP16”陷阱

FP16在NVIDIA GPU上的加速效果，跟Tensor Core强相关。Tensor Core从Volta架构（V100）开始出现，Turing（T4、RTX 20系）、Ampere（A100、RTX 30系）、Hopper（H100）等后续架构都标配了。Tensor Core跟普通CUDA Core的区别在于，它能在一个时钟周期内完成4x4矩阵的乘加运算，专门为深度学习这类矩阵密集计算设计。

如果显卡不支持Tensor Core，比如GTX 10系、GTX 16系这种，虽然也支持FP16计算，但走的是普通CUDA Core，FP16计算速度最多和FP32差不多，有时候甚至更慢。这种卡上强行开FP16训练，唯一的收益就是省显存，速度不仅不提升反而可能下降，很多人在老卡上“上了FP16的当”，原因就在这。

3.2 几款主流GPU的FP16算力对照

选型时可以参考这个表格，我整理了几款常见推理和训练GPU的关键数据：

注意，表格里A100的FP32算力反而比RTX 3090低，但实际训练速度要快很多，因为训练中大量调用的是Tensor Core路径。所以看显卡不能只看FP32 TFLOPS，还得看Tensor Core能力。

另外还有一个细节容易被忽略：Tensor Core并不是纯FP16专用，Ampere架构还支持TF32和BF16。TF32用FP32的位宽但把尾数截断到10位，这样既能用Tensor Core加速，又不需要改动模型的数据类型。后面第六章我会单独展开。

3.3 如何确认自己的GPU支持情况

动手前可以用几行代码确认环境支持情况：

python复制import torch
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))
print(torch.cuda.get_device_capability(0))
# 返回如 (7, 0) 表示V100，(8, 0) 表示A100，(8, 6) 表示RTX 3090等
# 计算能力6.0以下的GPU不支持FP16加速
major, minor = torch.cuda.get_device_capability(0)
if major >= 7:
    print("支持Tensor Core，FP16训练可用")
else:
    print("不支持Tensor Core，FP16收益有限")

我建议在项目初始阶段就做一次这个检查，省得后面跑完才发现硬件根本不支持加速。

4. 实操落地：PyTorch AMP使用与损失缩放机制

4.1 AMP的基本用法：自动混合精度

PyTorch从1.6版本开始原生支持自动混合精度（Automatic Mixed Precision，AMP），不需要手动把每个张量转成FP16。AMP的核心思路是自动判断哪些算子用FP16计算、哪些算子必须保持FP32，既享受Tensor Core加速，又避免数值稳定性问题。在PyTorch 2.x版本下，推荐直接使用torch.amp模块：

python复制import torch
from torch.amp import autocast, GradScaler

device = 'cuda'
model = MyModel().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler(device)

for epoch in range(epochs):
    for batch in dataloader:
        inputs = batch['image'].to(device)
        labels = batch['label'].to(device)

        optimizer.zero_grad()

        with autocast(device_type='cuda', dtype=torch.float16):
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

这段代码里最核心的就是三行：autocast上下文管理器负责把模型前向计算切到FP16，scaler.scale(loss)在反向传播前把Loss放大，scaler.step(optimizer)在更新权重前把梯度缩回来。第一次看可能觉得有点玄乎，我拆开解释一下。

autocast做的事情不是把整个模型变成FP16，而是让支持低精度加速的算子自动改成FP16执行，比如torch.matmul、torch.conv2d、torch.nn.Linear等；同时像softmax、layer_norm这类对数值范围敏感的算子，会保持FP32，避免误差累积。

GradScaler的损失缩放机制则是在反向传播前，把loss乘以一个默认的缩放因子（通常是65536或32768），让梯度整体进入FP16的可表示范围。调用scaler.step(optimizer)时，如果梯度出现上溢或下溢，scaler会跳过这一步的参数更新，并自动降低缩放因子；如果一切正常，它会把梯度缩小回正常量级再更新权重。不用手动干预，PyTorch会自动处理这些。

4.2 完整代码改造流程：从FP32到AMP

我习惯用一套固定的流程来做代码改造，每一步都有明确的检查点。

第一步，确认环境。第二步，import需要的模块。第三步，在训练循环里包上autocast。第四步，添加GradScaler。第五步，处理梯度裁剪和分布式训练等细节。下面给出一个更完整的模板，包含权重保存和多卡场景：

python复制import torch
from torch.amp import autocast, GradScaler

scaler = GradScaler(device)

for epoch in range(num_epochs):
    for batch in train_loader:
        opt.zero_grad()
        with autocast(device_type='cuda', dtype=torch.float16):
            out = model(batch['x'])
            loss = criterion(out, batch['y'])

        # 梯度裁剪时先unscale
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()

        # 记录loss和显存峰值
        peak_mem = torch.cuda.max_memory_allocated() / 1024**3

注意其中scaler.unscale_(optimizer)这步，如果用了梯度裁剪，一定要先反缩放再裁剪，否则裁剪阈值对FP16缩放过后的梯度是无效的。

保存和加载模型时，scaler的状态也要保存下来，方便断点续训：

python复制checkpoint = {
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scaler_state_dict': scaler.state_dict(),
    'epoch': epoch,
}
torch.save(checkpoint, 'checkpoint.pt')

# 恢复时
checkpoint = torch.load('checkpoint.pt')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scaler.load_state_dict(checkpoint['scaler_state_dict'])

4.3 实测对比：显存、速度、精度三组数据

只讲代码不讲效果是耍流氓。我在一张A100 40G上做了一个小实验，模型是UNet风格的医学图像分割网络，输入image size设为512x512，batch size为16，跑50个step，结果如下：

这个结果基本符合预期：显存省了快一半，速度提升了2.5倍，精度几乎无损。如果你发现速度提升不到1.5倍，先看看批量大小是不是太小了，混合精度在小batch size下Tensor Core利用率不够，提速不明显。

如果你发现显存省得不够多，除了模型参数和梯度，还要考虑激活值。PyTorch在前向传播时默认会保存FP16的激活值用于反向传播，但中间层的输出如果需要更多精度，可以针对特定模块关闭FP16。还有一个常见优化是把输入图像从FP32转成FP16之前先归一化，避免在低精度下做归一化导致误差传播。

4.4 分布式训练的AMP写法

现在很多项目都在用单机多卡或者多机多卡，DDP和AMP的配合很重要。DDP在反向传播时会做梯度同步，同步的梯度是从loss.backward()计算出来的缩放后梯度，DDP自身会确保梯度在进程间做AllReduce时保持一致性。不过需要注意，DDP的参数在每次step前是FP32的，并不会因为AMP变成FP16。

分布式训练里每个进程都需要维护自己的GradScaler对象，但因为DDP会对梯度做同步，最后所有进程的scaler状态会趋于一致。

python复制from torch.nn.parallel import DistributedDataParallel as DDP
from torch.amp import autocast, GradScaler

scaler = GradScaler(device)

for batch in train_loader:
    opt.zero_grad()
    with autocast(device_type='cuda', dtype=torch.float16):
        loss = model(batch['x'], batch['y'])
    scaler.scale(loss).backward()
    scaler.step(opt)
    scaler.update()

多卡场景下如果发现Loss震荡，优先检查是不是每个卡上的batch size太小导致BatchNorm统计量不稳定，这跟FP16关系不大，是分布式训练的常见问题。

5. 血泪排查：训练崩溃、Loss不降与精度损失的根因

5.1 Loss变NaN的第一现场

启用FP16后，最让人崩溃的就是Loss突然变成NaN。我排查过很多次，总结下来常见原因就是梯度上溢。怎么定位呢？先把损失缩放因子打出来看看，如果缩放因子一直在减小甚至变成1，说明梯度确实一直在溢出。另外可以把每层梯度的统计量打印出来，观察哪些层的梯度范数接近65504。

一个实用的排查脚本：

python复制for name, param in model.named_parameters():
    if param.grad is not None:
        grad_norm = param.grad.norm().item()
        if grad_norm > 1e4:
            print(f"梯度异常大: {name} {grad_norm}")

print(f"当前scale因子: {scaler.get_scale()}")

如果确认是梯度溢出，有几个方向可以调整。一是检查学习率是否过大，FP16训练下梯度会被放大，学习率要相应调小一些，我之前从1e-4降到5e-5就解决了问题。二是检查Loss缩放因子初始化，默认65536在绝大多数场景够用，但部分模型需要手动调大或调小。三是检查模型里是否有数值范围异常大的操作，比如depthwise卷积在某些实现下梯度容易爆炸。

5.2 Loss不降或收敛变慢，多半不是FP16的锅

有时候Loss没有变成NaN，但训练曲线明显比FP32慢。这种情况下先不要怀疑FP16，先检查几个容易忽略的点。第一个是数据归一化：如果输入数据还是FP32，在autocast里会自动被转成FP16，但如果你在外部提前把数据转成了FP16，归一化操作很可能因为精度不足出现偏差。正确的做法是输入保持FP32，让autocast在模型内部自动处理。

第二个是BatchNorm层。PyTorch的BatchNorm在autocast作用下仍会用FP32计算统计量，但如果你手动把BatchNorm强制转成FP16，running_mean和running_var的更新精度会下降，训练时间长了会出现明显偏差。保险做法是尽量不要手动干预BatchNorm的dtype。

第三个是学习率。AMP训练中梯度被缩放，Optimizer实际看到的梯度幅度和FP32时不太一样，有些优化器（尤其是Adam）对步长敏感，可能需要重新搜索学习率。我之前在某个模型上，FP32最优点是2e-4，AMP下最优降到1e-4左右，这也是正常的。

5.3 自定义算子和不兼容层怎么处理

如果你的模型里有自定义的CUDA算子，或者调用了某些不支持FP16的库，训练可能会直接报错或者静默输出错误结果。一个常见做法是为这些层单独关闭autocast：

python复制class MyCustomLayer(torch.nn.Module):
    def forward(self, x):
        # 强制该层用FP32计算
        return torch.ops.my_ops.custom_forward(x.float())

或者干脆在模型forward入口把fp16输入转回fp32。不过这样会带来额外的显存拷贝开销，能不用就不用。

可以用torch.autocast的enabled=False来做局部关闭：

python复制with autocast(device_type='cuda', dtype=torch.float16, enabled=False):
    output = custom_layer(input)

关于哪些算子适合FP16、哪些必须FP32，我个人的经验是：卷积、线性层、矩阵乘法这些计算密集型算子放心交给FP16；softmax、layer norm、loss计算、数值范围敏感的激活函数保持FP32；BatchNorm的统计量更新保持FP32。这也是PyTorch AMP默认策略的核心逻辑。

5.4 精度损失的临界场景

FP16并不是万能的。有些场景下，FP16导致的精度损失很难接受：

• 训练精度要求极高的模型（如某些医学影像诊断模型），对小数后4位的精度都很敏感。
• 模型规模很小、层数很浅，FP32本来就能轻松跑，FP16加速不明显但精度下降相对明显。
• 自定义损失函数中存在大量小数值相乘，这些区域在FP16下很容易下溢成0。
• 某些对比学习、度量学习场景，需要计算两两相似度，特征归一化后的数值范围很小，FP16的舍入误差可能影响排序结果。

如果遇到这些情况，我会优先尝试BF16，它在动态范围上和FP32一致，下溢问题大幅减少，只是尾数更短。在Ampere及以后架构上BF16同样走Tensor Core，速度也不差。

6. BF16与TF32：同门师兄弟的适用边界

6.1 BF16：动态范围是FP32的兄弟，精度更“粗糙”

BF16的全称是Brain Floating Point，最早由Google Brain提出。它的格式是1位符号位加8位指数位加7位尾数位，指数位和FP32完全一样，所以动态范围跟FP32一样大，最大能表示到3.4e38，最小正常数也是1.18e-38。这意味着在BF16下，梯度下溢的问题基本不存在。但代价是尾数位只有7位，有效精度只有约3位十进制数，比FP16的10位尾数还少。直观理解：BF16的“量程”和FP32一样宽，但“刻度”更粗。

实际使用中，BF16在NLP大模型训练上表现不错，因为Transformer里的激活值和梯度动态范围大，但单值精度要求并不极端。PyTorch里可以直接throughweight dtype指定：

python复制with autocast(device_type='cuda', dtype=torch.bfloat16):
    out = model(x)

不过BF16在这代N卡上只对Tensor Core友好的算子有加速，某些算子的实现不如FP16成熟。如果你的模型对下溢特别敏感，优先试试BF16；如果模型主要跑在RTX 30系以后，并且对显存占用要求苛刻，FP16仍然更合适。

6.2 TF32：A100上的“隐形势力”

TF32是Ampere架构针对FP32训练提出的一种加速方案。它本质上是把FP32的23位尾数截断到10位，但保留8位指数，所以动态范围和FP32一致，精度介于FP16和FP32之间。TF32的好处是，模型代码完全不用改动，还是FP32，但矩阵乘法在Tensor Core上会自动用TF32计算，速度大约是纯FP32的两倍。

在A100上，TF32默认是关闭的。我用过之后才知道，PyTorch中要手动打开：

python复制torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

TF32特别适合那些不想改代码但又想白嫖Tensor Core提速的场景。不过它的精度还是比FP16损失得少，因为至少保留了10位尾数。需要说明的是，TF32加速效果主要取决于矩阵乘法在整个训练中的占比，如果模型是计算密集型（CNN、Transformer），收益明显；如果模型有大量IO操作或自定义算子，效果就有限。

6.3 四兄弟的选型建议

不同精度格式各有自己的主场，我整理了一个简表：

如果让我给一个快速决策路径：卡是老卡（不支持Tensor Core），就老老实实FP32；卡是V100/T4/RTX 20系，优先FP16，用AMP；卡是A100/H100，而且模型是训练GPT类大模型，优先BF16；如果只是想在不改代码的情况下提速，打开TF32。

6.4 训练期间精度相关状态的监控习惯

用FP16/AMP训练时，养成监控几个指标的习惯能省很多事后排查的力气。我每次训练都会周期性打印三样东西：loss曲线的滚动平均值、当前scaler的scale值、梯度范数。默认情况下，PyTorch AMP会在后向传播中检测到inf或nan时自动跳过更新，所以如果loss出现偶尔跳动但整体趋势正常，可以不用管；如果梯度范数持续异常大或scale值一路下降，就要立刻停下来调参。

另外，如果用了DDP和AMP，保存checkpoint时建议把scaler状态也存下来，这样断点续训时缩放因子是连续的，不会因为重启后重新从头缩放造成训练波动。这些都是在实战中调节出来的经验，看着琐碎，关键时刻能救命。

最后的一些个人体会

FP16精度训练不是一项“开启即起飞”的功能，它是一套需要理解原理再去使用的系统工程。我自己从最开始直接改dtype导致Loss爆炸，到后来老老实实用AMP、配损失缩放、针对模型层逐一确认数值分布，这个过程踩了不少坑，但也确实把训练效率拉高了一个台阶。如果你正在被GPU服务器显存不足、训练速度慢的问题困扰，我建议按这个顺序走一遍：先确认你的GPU支持Tensor Core，然后用AMP包上训练循环，跑一次对照实验对比显存和速度，再根据Loss表现调整学习率和梯度裁剪方式。多数情况下，混合精度会给你一个惊喜。如果训练中遇到数值震荡，别急着推翻FP16，先按排查清单走一遍，大概率能定位到具体原因。希望这篇经验能帮你少绕几个圈。