训练大模型时Python环境下显存爆满怎么解决_实现梯度累加与混合精度

夏磊小哥_9707

夏磊小哥_9707

2026-04-20

164人浏览

原创

梯度累加需对每次loss除以accumulation_steps再backward,并仅在累积完成时调用optimizer.step()和optimizer.zero_grad(),避免梯度放大与显存泄漏。

训练大模型时python环境下显存爆满怎么解决_实现梯度累加与混合精度

梯度累加(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()  # 此时不释放激活值,但也不再保存新图
<pre class="brush:php;toolbar:false;">if (i + 1) % accumulation_steps == 0:
    optimizer.step()
    optimizer.zero_grad()  # 这里才真正清空梯度缓冲区

testing-python
testing-python

使用pytest编写和评估有效的Python测试。适用于编写测试、审查测试代码、调试测试失败或提高测试覆盖率。

下载

混合精度训练(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
<p>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)</p><pre class="brush:php;toolbar:false;">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 里。

Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!

相关专题

更多
python打包成可执行文件
python打包成可执行文件

本专题为大家带来python打包成可执行文件相关的文章,大家可以免费的下载体验。

2023.07.20

1711

4

python能做什么
python能做什么

python能做的有:可用于开发基于控制台的应用程序、多媒体部分开发、用于开发基于Web的应用程序、使用python处理数据、系统编程等等。本专题为大家提供python相关的各种文章、以及下载和课程。

2023.07.25

4324

7

format在python中的用法
format在python中的用法

Python中的format是一种字符串格式化方法,用于将变量或值插入到字符串中的占位符位置。通过format方法,我们可以动态地构建字符串,使其包含不同值。php中文网给大家带来了相关的教程以及文章,欢迎大家前来阅读学习。

2023.07.31

1689

3

python教程
python教程

Python已成为一门网红语言,即使是在非编程开发者当中,也掀起了一股学习的热潮。本专题为大家带来python教程的相关文章,大家可以免费体验学习。

2023.08.03

25237

23

python环境变量的配置
python环境变量的配置

Python是一种流行的编程语言,被广泛用于软件开发、数据分析和科学计算等领域。在安装Python之后,我们需要配置环境变量,以便在任何位置都能够访问Python的可执行文件。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

2023.08.04

3087

5

python eval
python eval

eval函数是Python中一个非常强大的函数,它可以将字符串作为Python代码进行执行,实现动态编程的效果。然而,由于其潜在的安全风险和性能问题,需要谨慎使用。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

2023.08.04

3107

5

scratch和python区别
scratch和python区别

scratch和python的区别:1、scratch是一种专为初学者设计的图形化编程语言,python是一种文本编程语言;2、scratch使用的是基于积木的编程语法,python采用更加传统的文本编程语法等等。本专题为大家提供scratch和python相关的文章、下载、课程内容,供大家免费下载体验。

2023.08.11

1163

5

python合并两个列表
python合并两个列表

Python是一种强大的编程语言,具有许多方便的功能和工具。在Python中,有多种方法可以合并两个列表。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

2023.08.10

596

4

python是前端还是后端
python是前端还是后端

Python属于前端也属于后端,其灵活性和丰富的生态系统使得开发人员能够在不同的领域中灵活运用。本专题为大家提供python相关的文章、下载、课程内容,供大家免费下载体验。

2023.08.11

2403

5

热门下载

更多
网站特效
/
网站源码
/
网站素材
/
前端模板

精品课程

更多
相关推荐
/
热门推荐
/
最新课程