AI Toolkit CUDA 显存不足？修复模型加载、训练和采样时的 OOM

如果你的 AI Toolkit 任务因 CUDA out of memory、OOM during training step 3 times in a row 或 0 bytes is free 而失败，不要继续原封不动地重新运行同一个任务。

实际上，AI Toolkit 中大多数 OOM 错误来自以下四个位置之一：

1. 模型加载（在真正的训练开始之前）
2. 最初几个训练步骤
3. 预览采样 / 基线样本生成
4. 视频特有的峰值（帧数太多、分辨率桶太大，或二者兼有）

本指南是快速恢复路径：确定你遇到的是哪种 OOM，在 RunComfy AI Toolkit 中修改正确的设置，更快地完成一次成功的重试。

---

快速修复清单（从这里开始）

• ✅ 在 Training 中降低 Batch Size
• ✅ 在 Datasets 中先禁用最大的 Resolution 桶
• ✅ 在 Sample 中降低 Width / Height / Sample Steps，或临时打开 Disable Sampling
• ✅ 点击 Show Advanced 确认 gradient_checkpointing: true
• ✅ 视频模型先降低 Num Frames，再动 learning rate
• ✅ 如果非常保守的配置仍然报错，当作可能的 Worker / GPU 状态问题处理，而非仅仅是配置问题

---

1) 确认是同一个问题

如果你的日志中包含以下消息，说明你来对地方了：

CUDA out of memory
torch.OutOfMemoryError
OOM during training step 3 times in a row
Tried to allocate ...
0 bytes is free
CUBLAS_STATUS_ALLOC_FAILED

常见情况：

• 任务在步骤 1 之前就失败
• 任务跑到步骤 2–10，然后反复 OOM
• 训练本身看起来没问题，但在生成样本时崩溃
• 相同配置有时成功有时失败

---

2) 首先：是哪种 OOM？

A. 模型加载时或训练开始前的 OOM

通常意味着以下之一：

• 对于当前的显存节省设置来说模型太重
• 预览 / 基线样本生成已经太贵
• Worker / GPU 状态不佳，没有从干净内存启动

典型表现：

• 在有意义的训练步骤开始之前就失败
• 模型加载完成后或第一个样本生成时立即报错
• 日志显示几乎没有空闲 VRAM，或抛出 CUBLAS 分配错误

B. 最初几个训练步骤的 OOM

这是最常见的配置导致的情况。

典型原因：

• gradient_checkpointing 关闭了
• Batch Size 太高
• 最大的数据集桶太激进
• 视频模型中，Num Frames 才是真正的显存峰值

C. 采样 / 预览生成时的 OOM

这是一个非常常见的陷阱。

训练配置可能差不多可以，但预览太贵：

• Sample Width / Height 太大
• Sample Steps 太高
• Sample Every 太频繁
• 预览视频使用了太多帧

D. 只是偶尔 OOM

这通常是接近极限的配置，不是什么神秘问题。

例如：

• 在小桶上能运行，到了最大桶就崩溃
• 视频训练只在最重的片段上失败
• 训练核心部分能放下，但样本生成把它推过了极限

---

3) RunComfy AI Toolkit 中最快的修复方法

Fix A — 重新打开 gradient checkpointing

图像模型 OOM 首先要检查的。

在哪里修改

1. 打开失败的任务
2. 点击 Show Advanced
3. 在 train: 下确认：

gradient_checkpointing: true

如果不确定，保持开启。

---

Fix B — 降低 Batch Size，用 Gradient Accumulation 保持稳定性

在哪里修改

• 打开任务编辑器
• 在 Training 面板中：
• 降低 Batch Size
• 如果想要稍大的有效 batch 而不增加 VRAM 峰值，保持或提高 Gradient Accumulation

安全重试规则

• 图像模型：OOM 时先降到 Batch Size = 1
• 视频模型：除非已证明配置稳定，否则默认 Batch Size = 1

不要把 learning rate 当作第一个显存调节手段。通常不是。

---

Fix C — 先移除最大的数据集桶

在哪里修改

• 进入 Datasets 面板
• 在 Resolutions 下先禁用最大的桶

安全回退顺序

• 1024 / 1536 → 先移除
• 验证稳定性期间保留 512 / 768
• 稳定后再逐个加回更大的桶

这是将一次接近极限的训练变成可重复训练的最快方法之一。

---

Fix D — 让预览采样更轻或临时禁用

如果崩溃发生在训练真正开始之前，或者每次采样运行时都崩溃，先修复预览。

在哪里修改

• 打开 Sample 面板

然后执行以下一项或多项：

• 降低 Width
• 降低 Height
• 降低 Sample Steps
• 增大 Sample Every
• 为验证运行打开 Disable Sampling

好的首次重试

如果目标是"证明任务能训练"，暂时不做采样的运行完全可以。

稳定后用更小的设置重新启用预览。

---

Fix E — 视频模型：先降 Num Frames

对于视频模型，帧数通常是最大的显存杠杆。

在哪里修改

• Datasets 面板 → Num Frames
• Sample 面板 → Num Frames

如果正在训练视频且出现 OOM，先降帧数，再降 batch size，再降分辨率。

不要从改 optimizer 或 LR 开始。

---

Fix F — 使用模型的低显存路径

某些架构在 VRAM 紧张时，本来就该使用显存节省设置来训练。

在哪里修改

• 点击 Show Advanced
• 在 model: 下寻找：

low_vram: true

某些模型的正确低显存路径还包括模型特定的量化或文本编码器处理。遵循对应的模型指南，不要靠猜。

---

4) 按故障时间快速诊断

---

5) 如何区分配置 OOM 和环境 / GPU 状态问题

当以下条件全部成立时，视为不仅仅是配置问题：

• Batch Size = 1
• gradient_checkpointing: true
• 保守的分辨率 / 帧数
• 任务仍在训练有意义地开始之前就失败
• 日志显示 0 bytes is free 或 CUBLAS 分配失败

这种情况下：

1. 停止重复同样的重试
2. 尽可能在新的 worker 上创建新任务
3. 如果之前能正常运行的相同保守配置现在立即失败，升级到支持团队

这很重要，因为反复重试可能浪费时间和 GPU 预算。

---

6) 最安全的回退顺序

图像模型

1. gradient_checkpointing: true
2. Batch Size → 1
3. 移除最大的 Resolution 桶
4. 缩小或禁用 Sample
5. 打开 low_vram 或模型的低显存路径

视频模型

1. 降低 Num Frames
2. Batch Size → 1
3. 移除最大的 Resolution 桶
4. 缩小或禁用 Sample
5. 打开 low_vram 或模型特定的 offload / 量化

---

7) 首次成功重试之后

拿到稳定的训练后：

• 每次只加回一个更重的设置
• 记录修改了什么
• 不要同时重新引入多个高风险设置

好的扩展顺序：

1. 保持相同的稳定显存设置
2. 增加 Steps
3. 重新启用更大的桶
4. 重新启用更丰富的采样
5. 最后才测试更大的 batch 或更长的视频

---

一句话总结

如果 AI Toolkit 报 OOM，停止随机试错。

打开 gradient_checkpointing，降低 Batch Size，移除最大的 Resolution 桶，让预览采样更轻，然后再重试。

---

相关指南

• 如何避免 AI Toolkit 中的 OOM：首次成功训练的安全设置
• 为什么在 AI Toolkit 中关闭 Gradient Checkpointing 会导致 OOM
• Wan 2.2 / LTX-2 OOM 指南：安全的帧数、Batch Size 和分辨率
• 打开 RunComfy AI Toolkit