当前位置：首页 > news >正文

PyTorch 混合精度训练：省显存不等于省时间

news 2026/7/6 23:23:47

PyTorch 混合精度训练：省显存不等于省时间

一、AMP 开启后训练反而变慢了：你被"混合精度即加速"的迷思骗了

在 PyTorch 训练脚本里加一行torch.cuda.amp.autocast()就能加速训练——这是大多数教程传递的信息。

你在 ResNet-50 上启用了 AMP，batch size 从 64 翻到 128，显存占用确实降了 30%。但 wall-clock 时间反而多了 12 秒。profiler 输出显示 GPU 利用率在 60% 附近振荡，而全精度模式能稳定在 90%。

问题根源在于：AMP 减少了 Python 层的显存带宽占用，但如果你的瓶颈是 CPU 数据加载或 kernel launch 开销，省显存并不能转化为省时间。

二、FP32 到 FP16 的数值链路与损失缩放机制

混合精度训练的核心是维护三套数值表示：

FP32 主权重（master weights）：参数的权威副本
FP16 前向权重 + 梯度：实际计算使用
FP32 梯度累加器：用于更新主权重

flowchart TB W32["FP32 主权重<br/>Master Weights"] -->|Cast| W16["FP16 前向权重"] W16 --> FP["前向传播<br/>Forward Pass (FP16)"] FP --> L["Loss 计算 (FP32)"] L -->|Scale| SL["缩放 Loss<br/>(防止梯度下溢)"] SL --> BP["反向传播<br/>Backward Pass (FP16)"] BP --> G16["FP16 梯度"] G16 -->|Unscale| G32["FP32 梯度"] G32 --> U["优化器更新<br/>Optimizer Step"] U --> W32 style W32 fill:#e8f5e9 style SL fill:#fff3e0 style G16 fill:#fce4ec style U fill:#e8f5e9

损失缩放（Loss Scaling）的必要性：

FP16 的最小正规范数约为 6.0×10⁻⁸。小于此值的梯度将在反向传播中被截断为 0，导致该参数无法更新。损失缩放将 loss 乘以一个大的因子（如 1024），使小梯度在 FP16 范围内可表示，反向传播后再除以该因子。

这个"先乘再除"的过程是混合精度训练正确性的关键，但也是引入额外 kernel launch 开销的来源——scaling/unscaling 不是免费的。

三、混合精度训练的完整最佳实践

import torch import torch.nn as nn from torch.cuda.amp import GradScaler, autocast from torch.utils.data import DataLoader from typing import Optional import time class MixedPrecisionTrainer: """混合精度训练器 设计哲学：配置项显式化。所有 AMP 相关参数集中在一个类中， 确保训练循环代码不散落 casts/enabled 等分散配置 """ def __init__( self, model: nn.Module, optimizer: torch.optim.Optimizer, scaler: GradScaler, device: str = "cuda", # 以下参数的设计意图： # - 动态增长/回退间隔：损失缩放因子不能始终为固定值。 # 过于频繁的回退（NAN 发生即改 scale）会导致多一轮完整的 # forward+backward，这在长训练中累计开销可观 # - init_scale：从保守值开始，让 scaler 自行向上调整 growth_interval: int = 2000, init_scale: float = 2.0 ** 16, backoff_factor: float = 0.5, growth_factor: float = 2.0, ): self.model = model self.optimizer = optimizer self.scaler = GradScaler( init_scale=init_scale, growth_interval=growth_interval, backoff_factor=backoff_factor, growth_factor=growth_factor, ) self.device = device self.model.to(device) def train_step(self, batch: tuple) -> dict: """ 单个训练 step 的混合精度实现 关键设计点： 1. autocast 只包裹前向传播——反向传播需要在 autocast 外以 FP32 执行 2. scaler.step() 而非 optimizer.step()——scaler 内部处理 unscale + step 3. scaler.update() 调整缩放因子——在 optimizer step 之后必须调用 """ inputs, targets = batch inputs, targets = inputs.to(self.device), targets.to(self.device) self.optimizer.zero_grad() # autocast 的作用范围精确控制：只包裹前向计算。 # 将 loss.backward() 放在 autocast 内会导致梯度在 FP16 下累加， # 而 FP16 累加器只有 10 位尾数精度，多步累加后精度损失不可接受 with autocast(): outputs = self.model(inputs) loss = nn.functional.cross_entropy(outputs, targets) # backward 在 autocast 外部，使用 FP32 的主拷贝 self.scaler.scale(loss).backward() # scaler.step 内部会先对梯度 unscale，再调用 optimizer.step() # 如果 unscale 后的梯度出现 NaN/Inf，这步会被跳过（scale 回退机制） self.scaler.step(self.optimizer) # 根据本轮梯度状态动态调整 scale 因子 self.scaler.update() return {"loss": loss.item(), "scale": self.scaler.get_scale()} @classmethod def benchmark( cls, model: nn.Module, dataloader: DataLoader, use_amp: bool = True, warmup_steps: int = 10, bench_steps: int = 50 ) -> dict: """ AMP 效果基准测试 为什么需要 warmup_steps？ 混合精度训练的 scale 因子在初期会频繁调整（adaptive scaling）， 这会导致前几个 step 的耗时偏高且不稳定。 跳过 warmup 阶段的测量才能得到真实性能数字 """ optimizer = torch.optim.SGD(model.parameters(), lr=0.1) scaler = GradScaler() if use_amp else None trainer = cls(model, optimizer, scaler) data_iter = iter(dataloader) # Warmup for _ in range(warmup_steps): try: batch = next(data_iter) except StopIteration: data_iter = iter(dataloader) batch = next(data_iter) trainer.train_step(batch) if use_amp else trainer._full_precision_step(batch) # Benchmark torch.cuda.synchronize() start = time.perf_counter() for _ in range(bench_steps): try: batch = next(data_iter) except StopIteration: data_iter = iter(dataloader) batch = next(data_iter) trainer.train_step(batch) if use_amp else trainer._full_precision_step(batch) torch.cuda.synchronize() elapsed = time.perf_counter() - start # 报告显存使用量 mem_allocated = torch.cuda.max_memory_allocated() / 1024**3 return { "total_time": round(elapsed, 3), "avg_step_ms": round(elapsed / bench_steps * 1000, 2), "use_amp": use_amp, "peak_gpu_memory_gb": round(mem_allocated, 2), "final_scale": scaler.get_scale() if scaler else None, } def _full_precision_step(self, batch: tuple) -> dict: """全精度训练对照实现，用于基准对比""" inputs, targets = batch inputs, targets = inputs.to(self.device), targets.to(self.device) self.optimizer.zero_grad() outputs = self.model(inputs) loss = nn.functional.cross_entropy(outputs, targets) loss.backward() self.optimizer.step() return {"loss": loss.item()} ## 四、混合精度训练的边界场景与架构权衡 **数值稳定性 vs 训练速度**。 FP16 的动态范围是 [5.96×10⁻⁸, 65504]。当梯度值落在此范围之外时，混合精度训练会出现问题。上溢出（>65504）触发 NAN，训练中断；下溢出（<5.96×10⁻⁸）导致梯度被截断为 0，参数停止更新。 损失缩放（Loss Scaling）解决了下溢出问题，但不能解决上溢出。当 scale 因子调整跟不上 loss 的剧烈波动时，scaler 会频繁回退，浪费了额外的 forward/backward pass。见证奇迹的时刻是当你发现 NAN 不是发生在 loss 最大的步骤，而是发生在一个看似平稳的训练阶段——因为累积的 scale 因子突然碰到了某个不稳定 batch。 **哪些层不适合 AMP**。 Softmax、LayerNorm 和某些自定义 activation 函数在 FP16 下数值不稳定。PyTorch 的 autocast 内置了这些层的 FP32 回退策略（如 `torch.nn.Softmax` 自动回退到 FP32），但自定义的 CUDA kernel 不会享受这种保护。如果训练包含自定义 kernel，必须手动在 autocast context 外以 FP32 执行。 **AMP 与数据并行的交互**。 DDP + AMP 的组合需要在 AllReduce 之前将 FP16 梯度 unscale 为 FP32。PyTorch 的 GradScaler 在 `scaler.step(optimizer)` 内部处理了这件事，但如果你需要手动操作梯度（如梯度裁剪），则必须先调用 `scaler.unscale_(optimizer)`，否则裁剪的是缩放后的梯度而非真实梯度。 ## 五、总结 混合精度训练的收益高度依赖场景：显存紧张 + 计算密集（如 ViT、大型 LM）的收益最大；显存充足 + 小模型的收益最小，甚至可能因 kernel launch 开销而变慢。 三个核心实践： 1. **始终跑 benchmark**：`benchmark(use_amp=True) vs benchmark(use_amp=False)`，用硬数据而非直觉决策。 2. **监控 scale 因子**：如果 scale 在 128 附近频繁回退，考虑降低初始学习率或使用 gradient clipping。 3. **理解哪些层不被 AMP 保护**：自定义 kernel 和自定义 activation 需要手动 FP32 回退，否则可能出现静默精度损失。

查看全文

http://www.jsqmd.com/news/1137687/