你是否曾为训练时的显存不足而烦恼？是否因漫长的训练周期而焦虑？自动混合精度（AMP）技术正是解决这些痛点的利器。作为PyTorch官方推荐的标准训练方案，AMP能够在保持模型精度的同时，显著提升训练速度并降低显存占用。本文从核心原理到实战技巧，带你彻底掌握这一深度学习训练的必备技能。无论你是正在应对大规模模型训练，还是希望在有限资源下提升效率，相信本文都能为你提供有力的技术支撑。

1 自动混合精度概述

自动混合精度（Automatic Mixed Precision，AMP）是一种深度学习训练加速技术，它通过在单精度（FP32）和半精度（FP16）之间智能切换计算精度，实现训练速度提升和显存占用减少，同时保持模型的精度不受影响。这一技术主要由NVIDIA在其Volta及之后的GPU架构中通过Tensor Core实现支持，并已成为现代深度学习训练的标准实践之一。

传统深度学习训练通常使用FP32格式表示所有参数和激活值，这确保了数值精度和表示范围，但存在显存占用大和计算效率低的问题。相比之下，FP16仅需FP32一半的存储空间，且在支持FP16的硬件上计算速度更快，理论上吞吐量可达FP32的2-8倍。然而，直接使用FP16训练会面临两大挑战：数值溢出问题（FP16的表示范围远小于FP32）和舍入误差（梯度值过小无法在FP16中表示）。

自动混合精度训练通过三大核心技术解决这些问题：

• 权重备份（保留FP32主参数副本）
• 损失缩放（放大损失值以保留梯度信息）
• 精度累加（使用FP16计算但用FP32累加）。这些技术组合使用使得AMP能够在保持模型精度的同时，显著提升训练效率。

表：不同浮点格式的比较

浮点格式	位宽	指数位	小数位	表示范围	精度	内存占用
FP64	64位	11位	52位	~10^(-308)到10^308	高	8字节
FP32	32位	8位	23位	~10^(-38)到10^38	中	4字节
FP16	16位	5位	10位	~6.1×10^(-5)到65504	低	2字节
BFloat16	16位	8位	7位	~10^(-38)到10^38	较低	2字节

AMP的优势不仅体现在计算速度的提升上，还能显著降低显存占用，使得训练更大模型或使用更大批次大小成为可能。根据NVIDIA官方数据，在Tensor Core支持的GPU上，AMP可带来1.5-3倍的训练加速。此外，通过减少GPU间通信量，AMP还能优化分布式训练性能。

目前主流的深度学习框架如PyTorch、TensorFlow和PaddlePaddle均已内置AMP功能，使其成为大规模模型训练的标准配置。特别是对于Transformer、ResNet等计算密集型模型，AMP带来的性能提升尤为明显。

• 官方教程：https://docs.pytorch.ac.cn/tutorials/recipes/recipes/amp_recipe.html

2 AMP核心技术原理

2.1 权重备份（Weight Backup）

权重备份是AMP的核心技术之一，其主要目的是解决舍入误差问题。在混合精度训练中，前向计算和反向传播使用FP16，但优化器更新参数时使用FP32。具体而言，训练过程中维护两套参数：一套是FP16格式的模型参数，用于计算前向和反向传播；另一套是FP32格式的主参数副本，用于优化器更新。

这种设计解决了深度模型中学习率与梯度乘积可能过小的问题。当梯度值很小（例如小于FP16能表示的最小精度2^(-24)时，FP16无法正确表示学习率与梯度的乘积，导致参数更新无效。通过使用FP32进行参数更新，可以有效避免这一舍入误差问题。

🎯 权重备份的实现机制如下：

• 前向传播：使用FP16参数进行计算，得到FP16的激活输出
• 反向传播：计算得到FP16的梯度
• 参数更新：将FP16梯度转换为FP32，与FP32的主参数副本进行更新操作

⛳️ 初学者的矛盾点

虽然维护两套参数会增加部分显存占用，但由于激活值和梯度仍然使用FP16存储，整体显存占用仍可降低约30%-50%。这是因为训练过程中的动态内存（中间激活值）通常占大部分，而参数占用的静态内存相对较小。

2.2 损失缩放（Loss Scaling）

损失缩放技术旨在解决FP16的数值下溢问题。FP16的可表示范围有限（约6.1×10^(-5)至65504），而深度学习中的梯度值往往非常小，容易低于FP16的最小表示范围，导致梯度变为零。

损失缩放通过在前向传播后对损失值乘以一个缩放因子（Scale Factor）来解决这一问题。根据链式法则，放大损失值等价于按相同比例放大梯度值，使这些原本可能下溢的梯度能够保留在FP16的表示范围内。在优化器更新参数前，再将梯度除以相同的缩放因子，确保参数更新不受影响。

🎯 动态损失缩放（Dynamic Loss Scaling）是损失缩放的高级形式，它根据训练情况动态调整缩放因子：

1. 从较大的缩放因子开始（如2^24）
2. 在训练过程中监控梯度是否出现无穷大或NaN
3. 若连续多次迭代未出现梯度异常，则适当增大缩放因子
4. 若检测到梯度异常，则跳过本次更新并减小缩放因子

这种动态调整机制确保了在训练的不同阶段使用合适的缩放因子，既防止梯度下溢，又避免梯度溢出。PyTorch的GradScaler实现了这一功能，可自动处理缩放因子的动态调整。

2.3 精度累加（Precision Accumulation）

精度累加技术主要用于解决计算过程中的舍入误差累积问题。在矩阵乘法等操作中，尽管输入和输出使用FP16，但中间累加过程使用FP32进行，最终结果再转换回FP16存储。

以矩阵乘法为例，当计算大矩阵的点积时，多个小数的乘积相加可能产生累积误差。使用FP16进行累加时，由于精度有限，这些误差可能导致结果不准确。而使用FP32进行中间累加，然后将结果舍入到FP16，可以大幅减少这种误差

现在的GPU的Tensor Core天然支持这种计算模式：它们使用FP16进行矩阵乘法，但使用FP32或更高精度进行中间累加。这使得在保持计算速度的同时，最大限度地减少了精度损失。

2.4 混合精度训练策略

混合精度训练有多种实现策略，不同的策略在精度和性能之间有不同的权衡。NVIDIA的APEX库提供了四种优化级别：

• O0：纯FP32训练，所有操作均使用FP32
• O1：混合精度，根据黑白名单自动选择操作精度（PyTorch AMP采用类似策略）
• O2：几乎FP16，保留BN等特定操作为FP32，维护FP32参数副本
• O3：纯FP16训练，所有操作使用FP16，可能影响精度

表：混合精度训练策略比较

策略	精度	速度	显存占用	稳定性	适用场景
O0（纯FP32）	高	慢	高	高	对精度要求极高的训练
O1（混合精度）	较高	较快	中等	高	大多数训练任务（推荐）
O2（几乎FP16）	中等	快	低	中等	大规模模型训练
O3（纯FP16）	较低	最快	最低	低	性能优先的实验性训练

PyTorch的AMP功能类似于O1策略，但进行了简化，更易于使用。它通过操作黑白名单自动决定每个操作的精度，用户无需手动干预。

3 PyTorch中AMP实现机制

3.1 torch.autocast上下文管理器

torch.autocast是PyTorch中实现自动精度管理的核心组件，它作为一个上下文管理器或装饰器，可以自动为区域内的操作选择合适的精度。通过torch.autocast，用户无需手动指定每个操作的精度，框架会根据预设的黑白名单自动管理精度转换。

使用torch.autocast的基本语法如下：

import torch

# 创建模型和优化器
model = Net().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

for input, target in data:
    optimizer.zero_grad()

    # 启用autocast的前向传播
    with torch.autocast(device_type='cuda', dtype=torch.float16):
        output = model(input)
        loss = loss_fn(output, target)

    # 反向传播（不在autocast上下文中）
    loss.backward()
    optimizer.step()

⛳️autocast会根据操作类型自动选择精度

• 白名单操作：如卷积、矩阵乘法等，使用FP16计算以获得性能提升
• 黑名单操作：如softmax、BatchNorm等，使用FP32计算以保持数值稳定性
• 其他操作：根据输入类型自动选择精度，保持与输入一致

autocast支持不同的设备类型（cuda、cpu）和数据类型（torch.float16、torch.bfloat16）。在CUDA设备上，通常使用torch.float16；在CPU上，则使用torch.bfloat16。

3.2 GradScaler梯度缩放

GradScaler负责实现损失缩放功能，自动管理缩放因子的动态调整。它与autocast配合使用，构成完整的AMP训练流程。

GradScaler的基本用法如下：

# 在训练开始时创建GradScaler实例
scaler = torch.cuda.amp.GradScaler()

for epoch in range(epochs):
    for input, target in data:
        optimizer.zero_grad()

        # 前向传播与损失计算
        with torch.autocast(device_type='cuda'):
            output = model(input)
            loss = loss_fn(output, target)

        # 缩放损失并反向传播
        scaler.scale(loss).backward()

        # 梯度更新
        scaler.step(optimizer)

        # 更新缩放因子
        scaler.update()

⛳️ GradScaler的主要工作流程包括：

1. 损失缩放：scaler.scale(loss)将损失乘以当前缩放因子
2. 反向传播：对缩放后的损失进行反向传播，得到缩放的梯度
3. 梯度更新：scaler.step(optimizer)先反缩放梯度，然后检查梯度是否包含无穷大或NaN
4. 缩放因子更新：scaler.update()根据梯度情况动态调整缩放因子

3.3 完整训练示例

以下是一个完整的AMP训练示例，展示了如何将标准FP32训练转换为混合精度训练：

import torch
from torch import nn

# 定义简单的网络
class SimpleNet(nn.Module):
    def __init__(self, in_size, out_size, num_layers):
        super().__init__()
        layers = []
        for _ in range(num_layers - 1):
            layers.append(nn.Linear(in_size, in_size))
            layers.append(nn.ReLU())
        layers.append(nn.Linear(in_size, out_size))
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

# 初始化模型、优化器和数据
model = SimpleNet(4096, 4096, 3).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
loss_fn = nn.MSELoss().cuda()

# 创建GradScaler
scaler = torch.cuda.amp.GradScaler()

# 训练循环
for epoch in range(epochs):
    for input, target in zip(data, targets):
        optimizer.zero_grad()

        # 使用autocast的前向传播
        with torch.autocast(device_type='cuda'):
            output = model(input)
            loss = loss_fn(output, target)

        # 使用GradScaler进行反向传播和参数更新
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

3.4 Autocast运算符行为

理解autocast中不同运算符的行为对于有效使用AMP至关重要。以下是主要运算符类别的精度选择规则

表：autocast运算符精度选择

运算符类别	自动精度选择	示例运算符	说明
白名单运算符	FP16	conv, linear, matmul, mm, bmm	这些操作在FP16下更快且数值稳定
黑名单运算符	FP32	softmax, cross_entropy, norm, pow	需要FP32的数值精度和动态范围
自动转换运算符	跟随输入类型	add, mul, concat	根据输入数据类型自动选择精度

对于不在预定义列表中的运算符，autocast不会强制改变其精度，它们将根据输入类型自动选择计算精度。这种设计确保了向前兼容性，用户可以在现有代码中安全地使用AMP。

4 AMP高级用法与最佳实践

4.1 梯度累积与梯度裁剪

🎯 在大模型训练或大批次训练场景中，梯度累积是一种常用技术。它通过多次小批次前向后传播再更新参数，模拟更大批次大小的效果。当结合AMP时，需要特别注意梯度缩放因子的处理。

以下是AMP与梯度累积的配合使用示例：

scaler = torch.cuda.amp.GradScaler()
accumulation_steps = 4

for i, (input, target) in enumerate(data):
    with torch.autocast(device_type='cuda'):
        output = model(input)
        loss = loss_fn(output, target)
        # 标准化损失以考虑累积
        loss = loss / accumulation_steps

    # 缩放损失并累积梯度
    scaler.scale(loss).backward()

    if (i + 1) % accumulation_steps == 0:
        # 可选：梯度裁剪（需要先反缩放）
        # scaler.unscale_(optimizer)
        # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

梯度裁剪在AMP训练中需要特殊处理，因为梯度被缩放因子放大。

在裁剪前，必须使用scaler.unscale_()将梯度反缩放回FP32范围：

# 不正确的梯度裁剪（直接在缩放后的梯度上裁剪）
# torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)  # 错误！

# 正确的梯度裁剪步骤
scaler.scale(loss).backward()

# 先反缩放梯度
scaler.unscale_(optimizer)

# 然后在FP32上裁剪梯度
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.1)

# 最后更新参数
scaler.step(optimizer)
scaler.update()

⚠️ 注意：每个优化器每个训练步骤只能调用一次unscale_，且必须在所有梯度计算完成后调用。

4.2 多模型、多损失与多优化器

对于生成对抗网络（GANs）、多任务学习等复杂训练场景，通常涉及多个模型、损失函数或优化器。AMP在这些场景下需要特别注意各组件间的协调。

以下是一个多优化器示例（如GAN中的生成器和判别器）：

# 初始化模型和优化器
generator = Generator().cuda()
discriminator = Discriminator().cuda()
g_optimizer = torch.optim.Adam(generator.parameters())
d_optimizer = torch.optim.Adam(discriminator.parameters())
scaler = torch.cuda.amp.GradScaler()

for epoch in epochs:
    for real_data, _ in dataloader:
        # 训练判别器
        with torch.autocast(device_type='cuda'):
            fake_data = generator(real_data)
            d_real_output = discriminator(real_data)
            d_fake_output = discriminator(fake_data.detach())
            d_loss = d_loss_fn(d_real_output, d_fake_output)

        scaler.scale(d_loss).backward()
        scaler.step(d_optimizer)

        # 训练生成器
        with torch.autocast(device_type='cuda'):
            g_output = discriminator(fake_data)
            g_loss = g_loss_fn(g_output)

        scaler.scale(g_loss).backward()
        scaler.step(g_optimizer)

        # 更新缩放因子（只需一次）
        scaler.update()

        # 清零梯度
        g_optimizer.zero_grad()
        d_optimizer.zero_grad()

当使用多个损失函数时，需要对每个损失分别应用scaler.scale()：

model = MyModel().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
scaler = torch.cuda.amp.GradScaler()

for input, target in data:
    optimizer.zero_grad()

    with torch.autocast(device_type='cuda'):
        output1, output2 = model(input)
        loss1 = loss_fn1(output1, target)
        loss2 = loss_fn2(output2, target)
        total_loss = loss1 + loss2

    # 对总损失进行缩放和反向传播
    scaler.scale(total_loss).backward()
    scaler.step(optimizer)
    scaler.update()

4.3 自定义函数与Autocast兼容性

当使用自定义Autograd函数时（通过继承torch.autograd.Function），需要确保这些函数与AMP兼容。PyTorch提供了custom_fwd和custom_bwd装饰器来简化这一过程。

以下是一个自定义矩阵乘法的示例，演示如何使其兼容AMP：

class MyMatmul(torch.autograd.Function):
    @staticmethod
    @torch.cuda.amp.custom_fwd  # 自动处理输入转换
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        return a.mm(b)

    @staticmethod
    @torch.cuda.amp.custom_bwd  # 保持与forward相同的autocast状态
    def backward(ctx, grad_output):
        a, b = ctx.saved_tensors
        return grad_output.mm(b.t()), a.t().mm(grad_output)

# 使用自定义函数
mymm = MyMatmul.apply

with torch.autocast(device_type='cuda'):
    # 现在MyMatmul会自动处理类型转换
    output = mymm(input1, input2)

对于需要特定精度（如必须使用FP32）的自定义函数，可以指定输入转换类型：

class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)  # 强制转换为FP32
    def forward(ctx, input):
        # 此函数始终在FP32精度下运行
        return some_operation_requiring_fp32(input)

    @staticmethod
    @torch.cuda.amp.custom_bwd
    def backward(ctx, grad_output):
        # 自动匹配forward的精度
        return grad_output

# 使用示例
func = MyFloat32Func.apply
with torch.autocast(device_type='cuda'):
    # 输入会自动转换为FP32
    output = func(input)

4.4 模型保存与恢复

AMP训练模型的保存与恢复需要同时考虑模型参数、优化器状态和GradScaler状态，以确保训练连续性。以下是最佳实践：

# 保存检查点
def save_checkpoint(model, optimizer, scaler, epoch, path):
    checkpoint = {
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scaler_state_dict': scaler.state_dict(),
        'epoch': epoch
    }
    torch.save(checkpoint, path)

# 加载检查点
def load_checkpoint(model, optimizer, scaler, path):
    checkpoint = torch.load(path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    scaler.load_state_dict(checkpoint['scaler_state_dict'])
    epoch = checkpoint['epoch']
    return epoch

# 训练中的保存与恢复示例
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters())
scaler = torch.cuda.amp.GradScaler()

start_epoch = 0
if resume_from_checkpoint:
    start_epoch = load_checkpoint(model, optimizer, scaler, 'checkpoint.pth')

for epoch in range(start_epoch, total_epochs):
    for input, target in dataloader:
        # ... 训练步骤 ...

        if (epoch + 1) % save_interval == 0:
            save_checkpoint(model, optimizer, scaler, epoch + 1, f'checkpoint_epoch_{epoch+1}.pth')

特殊情况下，当从非AMP训练恢复为AMP训练时，需要特殊处理：

# 从非AMP检查点恢复
checkpoint = torch.load('non_amp_checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

# 创建新的GradScaler（检查点中没有scaler状态）
scaler = torch.cuda.amp.GradScaler()

5 AMP性能优化与调试

5.1 性能优化策略

要充分发挥AMP的性能优势，需要综合考虑模型结构、数据流水线和硬件特性。以下是关键优化策略：

• 饱和GPU计算单元：确保GPU计算单元充分饱和，避免CPU成为瓶颈。这通常需要通过增大批次大小或模型复杂度来实现。当GPU利用率不足时，AMP的性能优势可能不明显。
• 优化张量形状：为了最大化Tensor Core利用率，确保矩阵乘法中的参与维度是8的倍数（如批量大小、输入/输出特征数）。这允许Tensor Core以最高效率运行

# 优化张量形状以提高Tensor Core利用率
# 好的设置：维度为8的倍数
batch_size = 512  # 8的倍数
in_features = 4096  # 8的倍数
out_features = 4096  # 8的倍数

# 可能次优的设置
batch_size = 513  # 不是8的倍数
in_features = 4097  # 不是8的倍数

• 减少CPU-GPU同步：避免在训练循环中频繁进行CPU-GPU同步操作，如：

• • 减少.item()调用次数
  • 避免在循环中打印CUDA张量值
  • 使用非阻塞数据传输

• 操作融合：将多个小操作融合为一个大操作，减少内核启动开销。PyTorch的自动优化功能会在可能时自动融合操作，但手动优化可能带来额外收益.

5.2 常见问题与解决方案

在AMP训练过程中可能会遇到各种问题，以下是常见问题及其解决方案：

梯度溢出/NaN损失

# 诊断梯度溢出问题
scaler = torch.cuda.amp.GradScaler(init_scale=2.**12)  # 可调整初始缩放因子

for input, target in data:
    optimizer.zero_grad()

    with torch.autocast(device_type='cuda'):
        output = model(input)
        loss = loss_fn(output, target)

    scaler.scale(loss).backward()

    # 检查梯度是否包含NaN/Inf
    if not torch.isfinite(loss):
        print("检测到NaN/Inf损失，考虑减小缩放因子或检查模型")
        continue

    scaler.step(optimizer)
    scaler.update()

    # 动态调整缩放因子
    if scaler.get_scale() < 1:
        print("缩放因子已减小，可能存在数值稳定性问题")

类型不匹配错误

当遇到类型不匹配错误时，可以临时禁用autocast来诊断问题：

# 方法1：全局禁用autocast调试
# with torch.autocast(device_type='cuda', enabled=False):
#     output = model(input)

# 方法2：在特定区域强制使用FP32
with torch.autocast(device_type='cuda'):
    # 大部分区域使用自动精度
    output1 = model_part1(input)

    # 特定子区域强制使用FP32
    with torch.autocast(device_type='cuda', enabled=False):
        output2 = model_part2(output1.float())  # 明确转换为FP32

内存不足问题

即使使用AMP，大型模型仍可能遇到内存不足问题。以下优化策略可以帮助减少内存占用：

1. 梯度检查点：平衡计算和内存使用
2. 激活检查点：只保存部分激活值，需要时重新计算
3. 分批处理：将大操作分解为较小批次

5.3 分布式训练中的AMP

在分布式数据并行（DDP）训练中，AMP的使用方式与单机训练类似，但需注意以下要点：

import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

def ddp_training(rank, world_size):
    # 初始化进程组
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

    # 创建模型并移至GPU
    model = MyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    optimizer = torch.optim.Adam(ddp_model.parameters())
    scaler = torch.cuda.amp.GradScaler()

    for epoch in range(epochs):
        for input, target in dataloader:
            input, target = input.to(rank), target.to(rank)
            optimizer.zero_grad()

            with torch.autocast(device_type='cuda'):
                output = ddp_model(input)
                loss = loss_fn(output, target)

            # 使用缩放损失进行反向传播
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    mp.spawn(ddp_training, args=(world_size,), nprocs=world_size)

5.4 性能分析与调试工具

PyTorch提供了一系列工具来帮助分析和调试AMP训练：

自动精度检查

# 检查当前是否处于autocast上下文中
def check_autocast_enabled():
    return torch.is_autocast_enabled()

# 检查GPU是否支持AMP
def check_amp_support():
    return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 7

# 检查当前GPU的Tensor Core支持
print(f"GPU支持AMP: {check_amp_support()}")

使用TorchScript优化

# 将模型转换为TorchScript以获得更好性能
model = MyModel().eval()

# 禁用JIT的autocast通道（如遇到问题）
torch._C._jit_set_autocast_mode(False)

with torch.autocast(device_type='cuda'):
    # 跟踪模型
    example_input = torch.randn(1, 3, 224, 224).cuda()
    traced_model = torch.jit.trace(model, example_input)
    optimized_model = torch.jit.freeze(traced_model)

  更多推荐

• 工··信··部 · A·I·G·C· 证书

  https://www.yuque.com/lhyyh/ai/ins6gx3o7hck7shb

• AI 工具集导航：

  https://tools.lhagi.com/

• AI 大模型全栈 50 万字知识库（🌍 ➕ LHYYH0001）