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 關閉是進階速度實驗。

相關指南

準備好開始訓練了嗎？