AI ToolkitでGradient Checkpointingをオフにするとなぜ OOM が発生するのか

Gradient Checkpointing をオフにした理由がはっきりしないなら、最も安全な答えはシンプルです：

オンに戻してください。

AI Toolkit の大型画像モデルでは、gradient_checkpointing をオフにすることは、安定した学習を OOM クラッシュに変える最も簡単な方法のひとつです。

このガイドでは以下を説明します：

Gradient Checkpointing が何をしているか
オフにするとなぜ VRAM が突然不足するのか
どのモデルファミリーが最も影響を受けやすいか
GC オフをテストしても良いケース
速度が目的なら代わりに何を変えるべきか

クイックアンサー

Gradient Checkpointing は基本的に メモリ vs 速度 のトレードオフです。

オン = VRAM ピークが低い、一般的に安全
オフ = VRAM ピークが高い、時に高速、OOM のリスクがかなり高い

ほとんどのユーザーにとって、特に大型画像モデルでは、品質の調整ではなく安定性の調整です。

クイック修正チェックリスト

✅ Show Advanced をクリック
✅ train: の中で gradient_checkpointing: true になっていることを確認
✅ 同じ Batch Size と Resolution でまず再試行
✅ それでも OOM なら、Batch Size を下げ、最大の Resolution バケットを外す
✅ プレビュー生成が問題なら、GC をオフにする代わりに Sample Width / Height / Sample Steps を下げる

1) Gradient Checkpointing が実際にやっていること

正しく使うために PyTorch の深い理論は不要です。

実践的な考え方：

GC オン では、AI Toolkit は中間のアクティベーションデータをメモリにあまり保持せず、必要なときに再計算する
GC オフ では、より多くのデータが VRAM に常駐するため、メモリに余裕があれば学習が速くなる可能性がある

問題なさそうに聞こえますが、以下が重なると話が変わります：

高解像度
複数の解像度バケット
大きなバッチサイズ
高コストなサンプリング
Qwen Edit、Z-Image、FLUX / Flex 系のような重いアーキテクチャ

余分なメモリの余裕は一瞬でなくなります。

2) オフにすると「突然」OOM が起きる理由

ユーザーはよくこう思います：

「トグルを1つ変えただけなのに。」

しかし、その1つのトグルが学習中に保持されるアクティベーションメモリの量を変えています。

GC オンで「重いけどギリギリ大丈夫」だった設定が、GC オフでは「ギリギリアウトか不可能」になります。

だからこそ、障害が突然に感じられます：

同じデータセット
同じモデル
同じプレビュープロンプト
同じバッチサイズ
GC だけ変更
ステップ 2、ステップ 3、または最初のプレビューで OOM

3) GC オフ時の高リスクモデルファミリー

AI Toolkit で繰り返し観察される OOM パターンに基づくと、以下の組み合わせはデフォルトで危険とみなすべきです：

モデルファミリー	GC オフ時の高リスク条件	安全な最初の試行
Z-Image	大きめバッチで 1024–1536 バケット	GC オン、控えめに開始、その後スケールアップ
Qwen-Edit	大きめバッチ / 複数の重い条件での 1024 ワークフロー	GC オン、Batch Size 1、プレビュー負荷を軽減
FLUX-dev / Flex 系大型画像モデル	大きめバッチまたは高解像度マルチバケット学習	GC オン、余裕に応じて Batch Size 1–4

役立つ考え方：

Z-Image は 高解像度 + 大きめバッチ + GC オフ の組み合わせで急速に危険になる
Qwen Edit は 1024 + 重い条件付け + GC オフ の組み合わせで急速に危険になる
FLUX / Flex 系は 大きめバッチ + GC オフ の組み合わせで急速に危険になる

4) RunComfy AI Toolkit での変更方法

現在の RunComfy UI では、以下のような設定が直接確認できます：

Batch Size
Gradient Accumulation
Steps
Unload TE
Cache Text Embeddings

ただし、gradient_checkpointing はアドバンスド設定から確認するのが最も確実です。

手順

ジョブを開く
Show Advanced をクリック
train: ブロックを見つける
以下を設定：

train:
  gradient_checkpointing: true

ジョブを保存
他に何も変えずに再実行

同じ設定で動作すれば、GC が原因だったことが確認できます。

5) 速度が本当の目標なら代わりにすべきこと

多くのユーザーは、学習を速くしたくて GC をオフにします。

その気持ちは理解できますが、通常は最初の速度実験として間違った選択です。

GC オフの前にこれらを試してください：

A. プレビューサンプリングを軽くする

Sample パネルで：

Width / Height を下げる
Sample Steps を下げる
Sample Every を増やす
一時的に Disable Sampling をオンにして学習の安定性を確認する

リスクの高い GC オフ設定を追求するよりも、実用的な速度向上が得られることが多いです。

B. バッチは小さく保ち、必要なら Accumulation でスケール

Training パネルで：

Batch Size は控えめに
VRAM ピークの大幅な増加なしに有効バッチサイズを少し大きくしたい場合のみ Gradient Accumulation を増やす

C. まず最大バケットを外す

Datasets で：

512 / 768 を維持
安定した学習の後にのみ 1024 / 1536 を再導入

D. モデルの低メモリパスを使う

一部のアーキテクチャでは、正しい方法は「GC オフ」ではなく：

low_vram: true
モデル固有の量子化
モデル固有のテキストエンコーダー最適化

汎用的な推測ではなく、モデルガイドに従ってください。

6) GC オフをテストしても良いのはいつか？

GC オフは上級実験として扱い、デフォルトにしないでください。

妥当な条件：

すでに安定した学習がある
Batch Size がまだ控えめ
最大バケットでギリギリの状態でない
プレビューが軽いか無効になっている
十分な VRAM の余裕がある
変更する変数は1つだけ

良いテスト手順：

GC オンでジョブを安定させる
残りの設定はそのまま
GC をオフにする
以下を監視：

初期ステップでの OOM
サンプリング時の OOM
間欠的なバケットスパイク

いずれかが発生したら、GC をオンに戻してそこで終了。

7) GC オフがではないもの

Gradient Checkpointing は以下のいずれでもありません：

魔法の品質向上スイッチ
類似度改善トグル
サンプルが悪いときの正しい最初の対策
すでに OOM 状態での良い最初の実験

サンプルが弱い場合は、以下を確認してください：

データセットの品質
キャプション
Rank
Steps
プレビューの一致性

GC オフが答えだと決めつけないでください。

8) FAQ

GC オンにすると出力品質は落ちる？

実際のところ、ほとんどのユーザーが気づく違いではありません。

本当のトレードオフは主にメモリ vs 速度です。

GC をオンにしてもまだ OOM が出る。どうすればいい？

次の調整ポイントは：

Batch Size
最大の Resolution バケット
プレビューサンプリングコスト
動画の場合は Num Frames

最初の学習を GC オフで始めるべきケースはある？

ほとんどのユーザーにとって：いいえ。

まず安定性を証明してから実験してください。

まとめ（一行）

AI Toolkit で OOM が発生していて gradient_checkpointing がオフなら、まずそれを直してください。

GC オンが安全なデフォルト。GC オフは上級者向けの速度実験です。

AI ToolkitでGradient Checkpointingをオフにするとなぜ OOMが発生するのか