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