logo
RunComfy
AI Toolkit LoRA 训练指南

为什么在 AI Toolkit 中关闭 Gradient Checkpointing 会导致 OOM

AI Toolkit 中 Gradient Checkpointing 的实用指南：了解它如何影响显存、何时保持开启更安全，以及在关闭 GC 之前应调整哪些设置。

使用 Ostris AI Toolkit 训练扩散模型

为什么在 AI Toolkit 中关闭 Gradient Checkpointing 会导致 OOM

如果你不确定为什么 Gradient Checkpointing 是关闭状态，最安全的做法很简单：

把它打开。

对于 AI Toolkit 的大型图像模型，关闭 gradient_checkpointing 是把一次稳定训练变成 OOM 崩溃的最快方式之一。

本指南将说明：

  • Gradient Checkpointing 做了什么
  • 关闭后为什么 VRAM 会突然不够用
  • 哪些模型系列最容易受影响
  • 什么时候可以尝试关闭 GC
  • 如果真正目的是提速，应该改什么

简要回答

Gradient Checkpointing 本质上是一个显存 vs 速度的权衡。

  • 开启 = 更低的 VRAM 峰值，通常更安全
  • 关闭 = 更高的 VRAM 峰值，有时更快，但 OOM 风险大幅增加

对于大多数用户，尤其是大型图像模型，GC 不是一个质量开关，而是一个稳定性开关。

快速修复清单

  • ✅ 点击 Show Advanced
  • ✅ 在 train: 下确认 gradient_checkpointing: true
  • ✅ 使用相同的 Batch Size 和分辨率先重试一次
  • ✅ 如果仍然 OOM，降低 Batch Size 并移除最高分辨率桶
  • ✅ 如果问题出在预览生成上，应降低 Sample Width / Height / Sample Steps，而非关闭 GC

1) Gradient Checkpointing 到底在做什么

不需要深入理解 PyTorch 理论也能正确使用它。

一个实用的理解方式：

  • GC 开启时，AI Toolkit 在内存中保留较少的中间激活数据，需要时重新计算
  • GC 关闭时，更多数据常驻 VRAM，如果内存充足，训练可能更快

听起来没问题，直到你叠加以下因素：

  • 高分辨率
  • 多个分辨率桶
  • 较大的 batch size
  • 高开销的采样
  • Qwen Edit、Z-Image、FLUX / Flex 系列等重型架构

此时额外的显存空间会迅速耗尽。

2) 为什么关闭后会"突然"OOM

用户经常这样想：

"我只改了一个开关。"

但这一个开关决定了训练过程中有多少激活数据留在显存里。

GC 开启时"勉强能跑"的配置，关闭后就变成了"刚好超限或完全不可能"。

所以故障感觉很突然：

  • 同一个数据集
  • 同一个模型
  • 同样的预览 prompt
  • 同样的 batch size
  • 只改了 GC
  • 结果在第 2 步、第 3 步或第一次预览时就 OOM

3) GC 关闭时的高风险模型系列

根据 AI Toolkit 反复出现的 OOM 模式，以下组合应默认视为危险：

模型系列 GC 关闭时的高风险条件 更安全的首次尝试
Z-Image 较大 batch + 1024–1536 分辨率桶 GC 开启，保守起步，再逐步提升
Qwen-Edit 较大 batch / 多个重条件的 1024 工作流 GC 开启，Batch Size 1，降低预览负载
FLUX-dev / Flex 类大型图像模型 较大 batch 或高分辨率多桶训练 GC 开启，根据余量选择 Batch Size 1–4

一个有用的思维模型：

  • Z-Image 在 高分辨率 + 大 batch + GC 关闭 组合下迅速变得危险
  • Qwen Edit 在 1024 + 重条件 + GC 关闭 组合下迅速变得危险
  • FLUX / Flex 类在 大 batch + GC 关闭 组合下迅速变得危险

4) 在 RunComfy AI Toolkit 中如何修改

在当前的 RunComfy 界面中，你可以直接看到：

  • Batch Size
  • Gradient Accumulation
  • Steps
  • Unload TE
  • Cache Text Embeddings

但 gradient_checkpointing 最容易通过高级配置来确认。

操作步骤

  1. 打开你的任务
  2. 点击 Show Advanced
  3. 找到 train: 块
  4. 设置： 
train:
  gradient_checkpointing: true
  1. 保存任务
  2. 不改其他任何设置，直接重试

如果同样的配置现在能跑了，就确认了 GC 是问题所在。

5) 如果真正目的是提速，应该怎么做

很多用户关闭 GC 是因为想要更快的训练。

可以理解——但这通常是错误的第一个提速尝试。

在关闭 GC 之前，先试试这些：

A. 降低预览采样开销

在 Sample 面板中：

  • 降低 Width / Height
  • 降低 Sample Steps
  • 增大 Sample Every
  • 暂时开启 Disable Sampling 来验证训练稳定性

这往往比追求高风险的 GC 关闭配置带来更多实际速度提升。

B. 保持小 batch，必要时用 Accumulation 扩展

在 Training 面板中：

  • 保持 Batch Size 保守
  • 仅在需要稍大有效 batch 而不大幅提升 VRAM 峰值时增加 Gradient Accumulation

C. 先移除最高桶

在 Datasets 中：

  • 保留 512 / 768
  • 仅在稳定运行后再加回 1024 / 1536

D. 使用模型自带的低显存路径

对于某些架构，正确做法不是"关 GC"，而是：

  • low_vram: true
  • 模型专属量化
  • 模型专属文本编码器优化

按照模型指南操作，不要凭猜测。

6) 什么时候可以尝试关闭 GC？

把关闭 GC 视为高级实验，而非默认选项。

合理的前提条件：

  • 你已经有了一次稳定的运行
  • Batch Size 仍然保守
  • 没有在最大桶上已经吃紧
  • 预览很轻或已禁用
  • 有足够的 VRAM 余量
  • 只改变一个变量，而非多个

一个好的测试流程：

  1. 先用 GC 开启稳定运行
  2. 保持其余配置不变
  3. 关闭 GC
  4. 观察是否出现：
    • 初始步骤 OOM
    • 采样 OOM
    • 间歇性桶峰值

如果出现任何一种情况，重新开启 GC 并就此停止。

7) GC 关闭不是什么

Gradient Checkpointing 不是：

  • 魔法般的质量提升开关
  • 相似度改善按钮
  • 样本效果差时的正确首选方案
  • 已经 OOM 时的好的首次实验

如果你的样本效果不好，应该检查：

  • 数据集质量
  • 标注（captions）
  • Rank
  • Steps
  • 预览一致性

不要假设关闭 GC 就是解决方案。

8) 常见问题

开启 GC 会影响输出质量吗？

实际使用中，大多数用户不会注意到质量差异。

真正的权衡主要是显存 vs 速度。

我开启了 GC 但仍然 OOM，怎么办？

接下来可以调整的是：

  1. Batch Size
  2. 最高分辨率桶
  3. 预览采样开销
  4. 对于视频，Num Frames

第一次训练应该用 GC 关闭开始吗？

对于大多数用户：不要。

先证明稳定性，再做实验。

一句话总结

如果你在 AI Toolkit 中遇到 OOM 且 gradient_checkpointing 是关闭的，先把它打开。

GC 开启是安全的默认设置。GC 关闭是高级速度实验。

相关指南

准备好开始训练了吗？