PyTorch混合精度训练实战：如何用AMP让你的旧显卡也能跑大模型

PyTorch混合精度训练实战：如何用AMP让你的旧显卡也能跑大模型

手里只有一张RTX 3060，甚至更老的RTX 2060，看着动辄数十GB显存需求的大模型，是不是觉得训练任务遥不可及？别急着升级硬件，也别轻易放弃。很多时候，限制我们的不是显卡的绝对性能，而是对显存利用效率的认知。混合精度训练，正是打开这扇门的钥匙。它不是什么遥不可及的实验室技术，而是PyTorch框架内置的、经过实战检验的成熟方案。今天，我们不谈空洞的理论，直接上手，看看如何用PyTorch原生的AMP（Automatic Mixed Precision）工具，让你手头的消费级显卡也能承担起大模型的训练任务。

我见过太多开发者，一遇到显存不足就本能地调小batch size，或者干脆放弃更复杂的模型结构。这其实是一种资源的浪费。现代GPU，即便是消费级产品，其Tensor Core单元对半精度浮点数（FP16）的计算吞吐量，往往是单精度（FP32）的2到8倍。这意味着，合理利用FP16，不仅能省下近一半的显存，还能大幅提升计算速度。关键在于，如何安全、稳定地使用它，避免因精度损失导致的模型不收敛或数值溢出。这就是AMP要解决的问题。

1. 混合精度训练的核心原理：不只是省显存

在深入代码之前，我们有必要花几分钟理解混合精度训练到底在做什么。这能帮助你在遇到问题时，知道该从哪里排查。

深度学习训练中的数值主要有三种：模型权重（Weights）、激活值（Activations） 和梯度（Gradients）。传统训练全程使用FP32，每个数占用4字节。FP16则只占2字节，直观上能节省一半的存储空间。但直接全部换成FP16行不通，主要因为两个问题：

1. 数值范围不足：FP16能表示的最大正值约为65504，最小正值约为5.96e-8。而深度学习中的梯度值通常非常小，可能在1e-7甚至更小。在反向传播的链式乘法中，这些微小梯度用FP16表示时会直接下溢（Underflow）为零，导致权重无法更新。
2. 舍入误差：当权重值很大（例如1.0），而梯度更新量很小（例如1e-5）时，FP16可能无法精确表示这个微小的变化，导致更新无效。

AMP的聪明之处在于“混合”。它并非简单地将所有计算转为FP16，而是采用了一个主权重（Master Weights） 的概念。具体流程如下：

• 前向传播（FP16）：使用FP16的权重副本和输入数据进行计算，得到FP16的损失。
• 损失缩放（Loss Scaling）：将FP16的损失乘以一个缩放因子（如1024或2048），将其“放大”到FP16能较好表示的数值范围内。这一步是关键，它保护了那些微小的梯度不被淹没。
• 反向传播（FP16）：计算得到放大后的FP16梯度。
• 梯度反缩放与更新（转FP32）：将放大后的FP16梯度转换回FP32，并除以相同的缩放因子，得到真实的FP32梯度。然后用这个FP32梯度去更新保存在FP32精度下的主权重。
• 同步权重：将更新后的FP32主权重转换为FP16副本，用于下一次前向传播。

这个过程就像一个精密的放大器：在前向和反向计算这些“重活”上用FP16来提速省显存，而在关键的权重更新这个“精细活”上，用FP32来保证数值稳定和精度。

注意：并非所有计算都适合用FP16。像Softmax、LayerNorm这类对数值范围敏感的操作，强制使用FP16容易导致溢出或精度灾难。AMP内部维护了一个“白名单”和“黑名单”，自动为不同类型的操作分配合适的精度。

为了更直观地理解AMP模式下显存的分布变化，我们来看一个简单的对比表格：

组件	FP32 训练显存占用	AMP (O2) 训练显存占用	说明
模型参数 (Weights)	4 × Ψ	2 × Ψ	AMP下，前向/反向时使用FP16副本。
梯度 (Gradients)	4 × Ψ	2 × Ψ	反向传播时产生FP16梯度。
优化器状态 (如Adam)	8 × Ψ	12 × Ψ	Adam需维护动量（m）和方差（v）两个状态。FP32训练时它们也是FP32。AMP下，为了稳定更新FP32主权重，优化器状态也需保持为FP32，且需额外维护一份FP32的主权重，因此总占用反而增加。
激活值 (Activations)	可变，巨大	可变，约减半	前向传播中产生的中间结果，与batch size和序列长度强相关。使用FP16可大幅减少其占用。
总计趋势	非常大	显著减少	对于大模型，激活值通常是显存瓶颈。AMP通过将其转为FP16，能释放出关键的空间来增大batch size或使用更大模型。

从上表可以看出，AMP省显存的秘诀主要在于大幅削减了激活值（Activations）的占用。对于Transformer这类激活值显存消耗巨大的模型，AMP的效果尤为明显。

2. PyTorch AMP快速上手：三行代码的魔法

理论说再多，不如一行代码。PyTorch从1.6版本开始，将AMP功能集成到了torch.cuda.amp模块中，使用起来异常简单。我们从一个最简单的训练循环开始改造。

假设你原本的训练循环长这样：

import torch
import torch.nn as nn
import torch.optim as optim

model = YourModel().cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()

for data, target in dataloader:
    data, target = data.cuda(), target.cuda()
    
    optimizer.zero_grad()
    output = model(data)          # 前向传播
    loss = criterion(output, target) # 计算损失
    loss.backward()               # 反向传播
    optimizer.step()              # 更新权重

启用AMP只需要引入两个关键对象，并修改三处：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import autocast, GradScaler  # 关键导入

model = YourModel().cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()
scaler = GradScaler()  # 关键对象1：梯度缩放器

for data, target in dataloader:
    data, target = data.cuda(), target.cuda()
    
    optimizer.zero_grad()
    
    # 关键修改1：用autocast上下文管理器包裹前向计算
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    
    # 关键修改2：用scaler.scale对loss进行缩放，再backward
    scaler.scale(loss).backward()
    
    # 关键修改3：用scaler.step和scaler.update来更新权重和缩放因子
    scaler.step(optimizer)
    scaler.update()

是的，核心就是GradScaler和autocast。autocast上下文管理器会自动将其范围内的计算转换为合适的精度（大部分转为FP16）。GradScaler则负责我们前面提到的损失缩放与反缩放流程。scaler.update()还会根据梯度是否出现溢出（NaN/Inf）来动态调整缩放因子，这是一个非常鲁棒的设计。

3. 深入AMP配置：O1与O2模式的选择与调优

PyTorch的AMP主要通过autocast的上下文来管理精度。虽然没有像NVIDIA Apex库那样显式的opt_level参数，但我们可以通过配置autocast和手动管理模型权重精度来模拟不同的优化级别。

3.1 O1模式（默认推荐）：自动黑白名单

我们上面展示的用法，其实就是O1模式的思路。autocast()默认使用PyTorch预设的操作类型黑白名单：

• 白名单操作：如torch.mm, torch.matmul, torch.conv2d等计算密集型操作，在autocast区域内会使用FP16计算以获得加速。
• 黑名单操作：如torch.softmax, torch.log, torch.exp等对数值范围敏感的操作，即使在autocast区域内也会保持FP32计算以保证稳定性。
• 灰名单操作：其余操作，其输入数据类型决定计算精度。如果所有输入都是FP16，则用FP16计算；否则用FP32。

这种模式在速度和稳定性之间取得了很好的平衡，是大多数情况下的首选。你几乎不需要额外配置。

3.2 O2模式：极致性能与手动控制

O2模式更为激进，它的目标是让几乎所有操作都使用FP16，仅将少数必须的操作（如Softmax、LayerNorm）和优化器更新步骤保留为FP32。在PyTorch中实现O2，需要手动将模型权重转换为FP16。

import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler

model = YourModel().cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-4)
scaler = GradScaler()

# 关键步骤：将模型权重转换为FP16
model.half()  # 这将把模型中所有浮点参数和缓冲区转换为FP16

for data, target in dataloader:
    data, target = data.cuda(), target.cuda()
    data = data.half()  # 输入数据也需要转换为FP16
    
    optimizer.zero_grad()
    
    # 注意：此时autocast仍然需要，因为它会保护黑名单操作（如LayerNorm）使用FP32
    with autocast():
        # 模型权重和输入已是FP16，大部分计算在此进行
        output = model(data)
        loss = criterion(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

O2模式能最大程度节省显存和提升速度，因为它连模型参数本身都存储为FP16。但它对模型和任务更挑剔。如果模型中有自定义的、对精度敏感的操作，或者任务本身对数值稳定性要求极高，O2可能导致训练发散。我的经验是，对于成熟的视觉CNN模型（如ResNet）和大多数NLP Transformer模型，O2通常工作良好。但对于一些新颖的、未经充分测试的架构，建议从O1开始。

3.3 关键配置：梯度缩放器（GradScaler）

GradScaler有几个重要的初始化参数，理解它们有助于调试：

scaler = GradScaler(
    init_scale=2.**16,      # 初始缩放因子，默认65536
    growth_factor=2.0,       # 当连续`growth_interval`步无溢出时，缩放因子乘以此值
    backoff_factor=0.5,      # 当发生溢出时，缩放因子乘以此值
    growth_interval=2000,    # 连续无溢出的步数阈值，用于增大缩放因子
    enabled=True             # 是否启用缩放，可用于动态开关AMP
)

• 动态缩放：这是AMP稳定训练的核心。如果检测到梯度溢出（出现Inf/NaN），scaler.step(optimizer)会跳过本次权重更新，并且scaler.update()会将缩放因子乘以backoff_factor（如0.5）以减少下次的缩放幅度。如果连续多步（growth_interval）没有溢出，则会乘以growth_factor（如2.0）尝试增大缩放因子，以更充分地利用FP16的动态范围。
• 跳过更新：当发生溢出时，scaler.step()不会执行optimizer.step()。这是保护机制，防止NaN污染模型权重。你可以在训练日志中监控这一点。

4. 实战技巧：针对Transformer等大模型的精细调整

对于BERT、GPT这类基于Transformer的大模型，直接应用上述基础AMP方法可能还不够。我们需要一些额外的技巧来保证稳定性和最大化收益。

4.1 保持LayerNorm和Softmax为FP32

Transformer中的LayerNorm和注意力机制中的Softmax是数值敏感的操作。在O2模式下，由于模型权重已是FP16，autocast会自动将这些黑名单操作保持在FP32。但在O1模式下，或者当你手动控制时，需要确保它们运行在FP32上下文。

方法一：依赖autocast（推荐） 最简单的方式就是相信autocast。PyTorch的默认黑名单已经包含了F.layer_norm和F.softmax。只要这些操作在with autocast():块内，它们就会被安全地执行在FP32下。

方法二：手动包装（用于自定义模块或复杂控制） 如果你需要更精细的控制，或者编写自定义模块，可以手动指定某些部分的计算精度：

import torch.nn.functional as F
from torch.cuda.amp import custom_fwd, custom_bwd

class StableTransformerLayer(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.attention = nn.MultiheadAttention(dim, heads)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(...)
    
    @custom_fwd(cast_inputs=torch.float32)  # 强制前向传播的输入转换为FP32
    def forward(self, x):
        # 即使外部是autocast(FP16)，这里的norm1计算也会使用FP32
        nx = self.norm1(x)
        attn_output, _ = self.attention(nx, nx, nx)
        x = x + attn_output
        
        nx = self.norm2(x)
        mlp_output = self.mlp(nx)
        x = x + mlp_output
        return x

使用@custom_fwd装饰器可以强制该函数内的计算使用指定的数据类型。

4.2 监控显存与检测溢出

开启AMP后，监控工具变得尤为重要。

• 显存监控：在训练循环中定期打印显存使用情况。

  import torch
  print(f"Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
  print(f"Cached: {torch.cuda.memory_reserved() / 1e9:.2f} GB")
  

  你会观察到，开启AMP后，allocated内存（当前已分配的张量）会显著下降，而cached内存（PyTorch缓存的内存池）可能变化不大。这是正常的。

• 梯度溢出检测：GradScaler提供了方便的方法来检查上一步是否发生了溢出。

  scaler.step(optimizer)
  scaler.update()
  
  if scaler.get_scale() != prev_scale:
      print(f"Loss scale changed from {prev_scale} to {scaler.get_scale()}")
      # 缩放因子变化意味着发生了溢出或增长
  

  如果频繁发生溢出（缩放因子持续下降），你可能需要：

  1. 检查模型结构，确保敏感操作运行在FP32。
  2. 尝试减小init_scale。
  3. 检查数据中是否存在异常值（如NaN）。

4.3 与梯度累积（Gradient Accumulation）结合

当单卡显存即使开了AMP也无法容纳想要的batch size时，梯度累积是另一个救命稻草。它的原理是：连续进行多次前向传播和反向传播，但不更新权重（optimizer.step()），只是累积梯度。在累积了N个微批次（micro-batch）后，用累积的总梯度更新一次权重。这相当于用时间换空间，实现了更大的有效batch size。

AMP与梯度累积能完美配合：

scaler = GradScaler()
accumulation_steps = 4  # 累积4步
optimizer.zero_grad()

for i, (data, target) in enumerate(dataloader):
    data, target = data.cuda(), target.cuda()
    
    with autocast():
        output = model(data)
        loss = criterion(output, target)
        loss = loss / accumulation_steps  # 损失按累积步数平均
    
    scaler.scale(loss).backward()
    
    # 每 accumulation_steps 步更新一次权重
    if (i + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()
        print(f"Step updated at iteration {i+1}")

注意，我们将loss除以了accumulation_steps。这是因为backward()会累加梯度。如果不除以步数，那么累积N步后的梯度相当于单个大batch的N倍，这会导致更新步长过大，破坏优化过程。除以N后，累积梯度在数值上就等价于一个大batch的梯度。

5. 常见陷阱与性能调优指南

混合精度训练并非银弹，实践中会遇到一些坑。这里总结几个最常见的问题和解决方案。

问题一：训练Loss变成NaN或突然爆炸。

• 可能原因：梯度溢出。FP16的动态范围有限，即使有Loss Scaling，某些极端大的梯度值仍可能溢出。
• 解决方案：
  1. 启用梯度裁剪（Gradient Clipping）：在scaler.step(optimizer)之前，使用scaler.unscale_(optimizer)将梯度反缩放回FP32，然后进行裁剪。

     scaler.unscale_(optimizer)
     torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 裁剪梯度范数
     scaler.step(optimizer)
     scaler.update()
     

  2. 检查数据：确保输入数据没有异常值（如inf或nan）。
  3. 降低初始缩放因子：创建GradScaler时，尝试更小的init_scale，例如2.**10。

问题二：开启AMP后速度没有提升，甚至变慢。

• 可能原因1：你的GPU不支持Tensor Core（计算能力低于7.0），或者驱动/CUDA版本太旧。AMP在Volta（V100）、Turing（RTX 20系列）、Ampere（RTX 30系列）及更新架构上才有显著的Tensor Core加速。
• 可能原因2：模型的计算瓶颈不在GPU，而在数据加载或CPU预处理（DataLoader）。使用torch.utils.data.DataLoader时，设置num_workers>0和pin_memory=True，并确保数据预处理不会成为瓶颈。
• 可能原因3：模型太小或操作类型不适合FP16加速。如果模型主要由元素级操作（如ReLU、加法）构成，而不是矩阵乘（MatMul）或卷积（Conv），那么FP16的加速收益可能被类型转换的开销抵消。

问题三：验证集精度略有下降。

• 可能原因：这是精度损失的正常体现。FP16毕竟损失了部分精度。
• 解决方案：
  1. 在验证/测试时关闭AMP：评估模型时，使用FP32精度以确保最高的准确性。

     @torch.no_grad()
     def evaluate(model, dataloader):
         model.eval()
         # 不使用autocast，确保全精度计算
         for data, target in dataloader:
             data, target = data.cuda(), target.cuda()
             output = model(data)  # 这里是FP32计算
             # ... 计算指标
     

  2. 微调超参数：有时稍微降低学习率或调整其他超参数可以弥补精度损失。
  3. 坚持使用O1模式：O2模式虽然更快，但精度损失风险也更高。如果O1能达到满意的精度和速度，就没必要追求O2。

性能调优终极建议：

1. 从O1模式开始：它最稳定，适合绝大多数模型。
2. 逐步增大batch size：开启AMP后，显存占用降低，尝试逐步增加batch_size直到占满显存。更大的batch size有时能带来更稳定的训练和更好的最终性能。
3. 监控Loss Scale：如果scaler.get_scale()稳定在一个较高的值（如几千），说明训练很健康。如果它持续下降并保持在很低的值（如1或2），说明梯度溢出频繁，需要排查问题。
4. 使用更现代的优化器：像AdamW通常比原始Adam对混合精度训练更友好。LAMB优化器在大batch训练中表现也很好。

混合精度训练已经从一项前沿技术变成了深度学习训练的标配。对于个人开发者和小型团队而言，它是在有限硬件条件下探索更大模型、更复杂任务的必备技能。我自己的RTX 3060 12GB显卡，通过AMP结合梯度累积，曾经成功微调过70亿参数的模型，而这在纯FP32模式下是不可想象的。关键在于理解其原理，耐心调试，并善用监控工具。下次当你再看到“CUDA out of memory”时，不妨先别急着抱怨硬件，试试打开AMP，或许会有惊喜。