显存不足？梯度累加+混合精度轻松解决

显存爆满是深度学习训练中的常见痛点，但通过正确组合梯度累加与混合精度训练（AMP），可在不牺牲模型性能的前提下显著降低显存占用——关键在于精准控制梯度缩放、时机正确的优化器步进与清零、严格限定autocast作用域，以及规避FP32隐式张量、优化器状态残留和内存碎片等“隐形杀手”；本文直击实践中高频踩坑点，从原理本质到代码细节层层拆解，帮你把显存利用率从“总在崩溃边缘”真正拉回稳定高效的轨道。

梯度累加（Gradient Accumulation）怎么写才不爆显存

梯度累加的本质是用时间换空间：把一个大 batch 拆成多个小 batch，每次只 forward + backward 一小部分，不立即更新参数，而是累积梯度，直到达到等效 batch size 再 step。关键不是“多算几次”，而是避免单次 backward 保存全部激活值。

常见错误是忘记缩放 loss 或没清零 optimizer 的 grad，导致梯度叠加错乱或显存持续增长。

• 必须对每次 loss 除以 accumulation_steps，否则梯度会按倍数放大
• optimizer.zero_grad() 只能在累积开始前或 step 后调用，不能在每次 backward 前都调
• 反向传播后不要立刻 del output/loss，PyTorch 自动管理计算图；但要确保没有其他变量（如中间特征）意外持有引用
• 示例代码中 if (i + 1) % accumulation_steps == 0: 是安全边界，避免最后一个 batch 不足步长时提前 step

accumulation_steps = 4
for i, (data, target) in enumerate(train_loader):
    data, target = data.to(device), target.to(device)
    output = model(data)
    loss = criterion(output, target) / accumulation_steps
    loss.backward()  # 此时不释放激活值，但也不再保存新图
if (i + 1) % accumulation_steps == 0:
    optimizer.step()
    optimizer.zero_grad()  # 这里才真正清空梯度缓冲区

混合精度训练（AMP）启用后反而 OOM？检查这三点

混合精度不是“开了就省一半显存”，它依赖 CUDA 计算能力 ≥ 7.0（如 V100/A100/Tesla T4），且需正确包裹 forward 和 backward。常见失效场景是部分模块未进入 autocast 上下文，或 scaler 更新逻辑出错，导致 FP32 梯度仍被全程保留。

• model.forward() 必须完全在 with autocast(): 内，包括 loss 计算；否则 loss 仍是 FP32，scaler.scale() 失效
• scaler.step(optimizer) 前必须有 scaler.scale(loss).backward()，否则梯度未被缩放，step 会跳过更新，同时显存不释放
• 若模型含自定义 OP（如 torch.compile 后的图、或非标准 CUDA kernel），autocast 可能不生效，需手动插入 .half() 并确保输入/输出 dtype 一致

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for data, target in train_loader:
optimizer.zero_grad()
data, target = data.to(device), target.to(device)
with autocast():  # ← 必须从这里开始，到 loss 结束
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()  # ← loss 是 autocast 下的 FP16 输出
scaler.step(optimizer)
scaler.update()  # ← 更新 scaler 内部状态，否则下次 scale 失效

torch.cuda.empty_cache() 该在哪调、不该在哪调

这个函数不解决根本问题，但它能缓解碎片化引发的“明明有空闲却分配失败”。关键是：它只释放 PyTorch 缓存池中“已标记为可回收但尚未归还驱动”的内存，对正在使用的张量无效。滥用反而拖慢训练——每次调用都会阻塞 CUDA 流。

• 适合位置：每个 epoch 结束后、evaluation 前、或切换模型/数据集前
• 绝对避免位置：forward/backward 循环内部、loss.backward() 后立刻调用（此时缓存池还没来得及整理）
• 配合 del 才有效：比如 del output, loss 后再 empty_cache()，否则 Python 引用仍在，张量不会被回收
• 注意 memory_reserved 和 memory_allocated 差值：若差值长期 > 1GB，说明碎片严重，应优先检查张量生命周期而非狂刷 empty_cache

为什么梯度累加 + AMP 一起用，显存还是降不下来

两者叠加理论上可将显存压到原始的 ~30%，但实际常卡在 50–60%。核心原因是：AMP 默认只降 activations 和 weights，而优化器状态（如 Adam 的 exp_avg 和 exp_avg_sq）仍是 FP32；梯度累加又延长了梯度缓冲区存活时间。更隐蔽的问题是模型中存在隐式 FP32 张量，比如：

• LayerNorm/BatchNorm 的 running_mean/var 默认是 FP32，即使模型设为 .half() 也不会自动转
• 某些 loss 函数（如 nn.CrossEntropyLoss）内部有 FP32 累加，需显式传 label_smoothing 或改用 torch.nn.functional.cross_entropy(..., dtype=torch.float16)
• torch.compile(model) 在某些版本中会绕过 autocast，需确认 torch._dynamo.config.automatic_dynamic_shapes = True

最有效的排查方式是插入 torch.cuda.memory_summary()，重点看 “allocated by [host]” 和 “allocated by [cudaMallocAsync]” 的分布，而不是只盯 total。很多“隐形大户”藏在 DataLoader pin_memory 或 profiler 的临时 buffer 里。

今天关于《显存不足？梯度累加+混合精度轻松解决》的内容介绍就到此结束，如果有什么疑问或者建议，可以在golang学习网公众号下多多回复交流；文中若有不正之处，也希望回复留言以告知！