RTX 4090でのFlux LoRAトレーニングのメモリ不足エラーを修正する

24GBのVRAMを搭載したRTX 4090を持っているのに、ローカルでのFluxトレーニングには十分なはずなのに、毎回CUDAのメモリ不足エラーでクラッシュしてしまいます。トレーニングは正常に開始し、数ステップ実行した後に停止します。バッチサイズを減らしても、まだクラッシュします。何がそのメモリを消費しているのでしょうか？

簡潔な回答： RTX 4090でFlux LoRAトレーニングがOOMになるのは、Fluxの大規模なアーキテクチャがデフォルトのトレーニング設定で30-40GBのVRAMを必要とするためです。これを修正するには、勾配チェックポイントを有効にして計算とメモリをトレードオフし、バッチサイズを1に減らし、1024x1024ではなく512x512のトレーニング解像度を使用し、FP16またはBF16での混合精度を有効にし、メモリ効率の良いアテンション実装を使用します。これらの設定により、24GBで完全なトレーニングが可能になり、高品質なLoRAが生成されます。

RTX 4090は適切に設定すればFlux LoRAトレーニングに優れたカードです。問題は、デフォルトのトレーニング設定が24GB以上のVRAMを想定していることです。適切な設定で、高品質なFlux LoRAを効率的にトレーニングできます。安定したトレーニングのためのセットアップを設定しましょう。

なぜFluxトレーニングはそれほど多くのVRAMを必要とするのか？

トレーニング中に何がメモリを消費するかを理解することで、効果的に最適化できます。

モデルサイズ

Fluxの基本モデルはSDXLやSD 1.5よりも大幅に大きいです。フル精度のモデル重みだけで約23GBを消費します。これはトレーニングが始まる前に、すでに4090のVRAM全体を使い切っています。

トレーニング中は、モデル、勾配、オプティマイザの状態、アクティベーションのためのメモリが必要です。これらのそれぞれがモデル自体のサイズに近づく可能性があります。

特にLoRAトレーニングでは、基本モデルを凍結し、小さなアダプター層のみをトレーニングします。これは大きな助けになりますが、完全なモデルを流れるアクティベーションと勾配からのメモリ圧力を排除するわけではありません。

アクティベーションメモリ

フォワードパス中、中間アクティベーションはバックワードパスで使用するために保存されます。これらのアクティベーションはバッチサイズと解像度に応じて増加します。

1024x1024の解像度では、アクティベーションメモリはモデルサイズを超える可能性があります。単一のトレーニングバッチでアクティベーションだけで15-20GBが必要になる場合があります。

これがトレーニングが数ステップ後にクラッシュする理由です。最初のステップは収まるかもしれませんが、メモリの断片化と蓄積された状態により、後続のステップが失敗します。

オプティマイザの状態

Adamのようなオプティマイザは、すべてのトレーニング可能なパラメータに対して2つのモメンタム値を保存します。これにより、トレーニングされているパラメータに必要なメモリが2倍になります。

完全なファインチューニングでは、オプティマイザの状態メモリはモデルサイズの2倍に相当します。LoRAトレーニングでは、トレーニングされるパラメータが少ないため、オプティマイザの状態は小さくなりますが、それでもかなりの量です。

勾配メモリ

各トレーニング可能なパラメータの勾配は、バックワードパス中にストレージが必要です。これにより、メモリ要件にトレーニング可能なパラメータのサイズの完全なコピーが追加されます。

モデル、アクティベーション、オプティマイザの状態と組み合わせると、デフォルトのFluxトレーニング設定での総メモリ要件は簡単に40-50GBに達します。

24GB VRAMでトレーニングをどのように設定するか？

これらの設定により、RTX 4090での安定したFlux LoRAトレーニングが可能になります。

勾配チェックポイントを有効にする

勾配チェックポイントは、最も効果的なメモリ最適化です。20-30%の計算時間の増加と引き換えに、アクティベーションメモリを60-70%削減します。

フォワードパス中にすべてのアクティベーションを保存する代わりに、チェックポイントはほとんどを破棄し、バックワードパス中に再計算します。メモリ使用量はモデルの深さに関係なくほぼ一定になります。

Kohya SSでは、トレーニング設定で勾配チェックポイントを有効にします。オプションは通常、シンプルなチェックボックスまたはブールパラメータです。

diffusersを使用したカスタムトレーニングスクリプトでは、トレーニング開始前にmodel.enable_gradient_checkpointing()を呼び出します。

チェックポイントを有効にするとトレーニングに時間がかかりますが、クラッシュする代わりに実際に完了します。時間のトレードオフは価値があります。

バッチサイズを1に設定する

バッチサイズはアクティベーションメモリを直接倍増させます。バッチサイズ4はバッチサイズ1の約4倍のアクティベーションメモリを使用します。

バッチサイズを1に設定します。メモリコストなしでより大きな有効バッチサイズをシミュレートするために勾配累積を使用します。

例えば、バッチサイズ1で勾配累積ステップ4は、メモリ内で1サンプルのアクティベーションのみを保持しながら、有効バッチサイズ4を提供します。

勾配累積は、重みを更新する前に複数のフォワードパスにわたって勾配を累積します。メモリ使用量はバッチサイズ1で一定のまま、トレーニングダイナミクスは大きなバッチに近似します。

トレーニング解像度を下げる

解像度はメモリに対して二乗の影響を持ちます。解像度を2倍にするとアクティベーションメモリが4倍になります。

1024x1024ではなく512x512でトレーニングします。これによりアクティベーションメモリが約75%削減されます。

512x512でのトレーニングがネイティブ解像度よりも悪い結果を生むのではないかと心配するかもしれません。実際には、低解像度でトレーニングされたLoRAは高解像度の推論にうまく転送されます。トレーニングしているスタイル要素やコンセプトは、1024x1024の生成でも反映されます。

特定のユースケースで高解像度トレーニングが必要な場合、他の積極的な最適化と組み合わせれば768x768が可能な場合もあります。注意深くテストし、メモリを監視してください。

混合精度トレーニングを使用する

混合精度は、重要な値をFP32に保ちながら、ほとんどの操作にFP16またはBF16を使用します。

BF16はAmpere以降のGPUに推奨されます。より大きなダイナミックレンジにより、FP16よりも勾配をより良く処理します。

トレーニング設定で混合精度を有効にします。Kohya SSでは、精度ドロップダウンからBF16を選択します。カスタムスクリプトでは、PyTorchのautocastコンテキストマネージャを使用します。

混合精度はモデル重みとアクティベーションのメモリをほぼ半減させます。他の最適化と組み合わせると、24GBトレーニングに不可欠です。

メモリ効率の良いアテンションを有効にする

標準的なアテンション実装は大きな中間テンソルを割り当てます。xFormersやFlash Attentionのようなメモリ効率の良いバリアントは、アテンションをチャンクで処理します。

xFormersはトレーニングに適しており、広くサポートされています。別途インストールし、トレーニング設定で有効にします。

Flash Attentionは最近のGPUでより良いパフォーマンスを提供します。トレーニングフレームワークがサポートしているかどうかを確認してください。

メモリ効率の良いアテンションは、アテンションメモリ使用量を80%以上削減できます。Fluxのような大規模モデルでは、これは数GB節約されることになります。

8ビットオプティマイザを使用する

標準のAdamオプティマイザは、パラメータごとに2つのFP32値を保存します。8ビットAdamはこれらをINT8に量子化し、オプティマイザメモリを半減させます。

bitsandbytesライブラリをインストールし、トレーニングセットアップで8ビットAdamを設定します。Kohya SSはこれを直接サポートしています。カスタムスクリプトでは、bitsandbytesから8ビットオプティマイザをインポートする必要があります。

品質への影響は最小限です。8ビットAdamは、ほとんどの場合、フル精度と同様に収束します。

Kohya SSで最も効果的な設定は何ですか？

Kohya SSは最も人気のあるFlux LoRAトレーニングツールです。ここでは、具体的に動作する設定を紹介します。

推奨設定

RTX 4090でのFlux LoRAトレーニングの出発点としてこれらの設定を使用してください。

解像度: 512,512 バッチサイズ: 1 勾配累積ステップ: 4 混合精度: bf16 勾配チェックポイント: 有効 xFormers: 有効 オプティマイザ: AdamW8bit ネットワークランク: 16-32 ネットワークアルファ: ランクと同じか半分 学習率: 1e-4 トレーニングステップ: キャラクターで1000-2000、スタイルで2000-4000

この設定では約20GBのVRAMを使用し、安定性のための余裕を残します。

異なるトレーニングタイプへの調整

キャラクターLoRAは、ランク16程度とステップ1000程度の少ない設定を使用できます。キャラクターの特徴は比較的簡単にキャプチャできます。

スタイルLoRAは、ランク32-64程度の高いランクと、3000-4000程度のより多くのトレーニングステップの恩恵を受けます。芸術的なスタイルには学習すべきバリエーションがより多くあります。

特定のオブジェクトやポーズのコンセプトLoRAは大きく異なります。キャラクター設定から始めて、結果に基づいて調整してください。

高いランクにはより多くのVRAMが必要です。ランクを64以上にプッシュする場合、トレーニング中のメモリ使用量を監視してください。解像度を448x448に下げる必要があるかもしれません。

キャプション設定

良いキャプションはトレーニング品質を劇的に向上させます。Fluxは自然言語の説明によく反応します。

BLIPまたは同様のツールを使用して初期キャプションを生成し、手動で精製します。不正確な説明を削除し、トリガーワードを追加します。

Fluxでは、SD 1.5よりも長いキャプションが一般的によく機能します。被写体だけでなく、画像コンテンツに関する関連する詳細を含めてください。

データセット全体で繰り返しのキャプションを避けてください。トリガーワードを一貫して保ちながら、言語を変えてください。

サンプル画像生成

進捗を監視するためにトレーニング中のサンプル生成を有効にします。サンプル頻度を100-200ステップごとに設定します。

サンプル生成はメモリオーバーヘッドを追加します。サンプリング中にOOMが発生した場合は、勾配累積を増やすか、サンプル頻度を減らしてください。

サンプルは、トレーニングがうまくいっているかどうか、オーバーフィットし始めているかどうかを教えてくれます。サンプルが良く見えるが劣化し始めていない時点でトレーニングを停止してください。

持続的なOOMエラーをどのようにトラブルシュートするか？

最適化後もクラッシュが続く場合は、これらの追加要因を調査してください。

VRAMの断片化

PyTorchのメモリアロケータは時間の経過とともにVRAMを断片化し、総空きメモリが十分に見えても障害を引き起こす可能性があります。

PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128環境変数を使用してトレーニングを実行します。これにより、断片化を減らすために割り当て戦略が変更されます。

以前のGPU操作なしで新しく開始すると役立ちます。トレーニング前にシステムを再起動するか、少なくともすべてのPythonプロセスを終了してください。

メモリ使用量の監視

トレーニング中にVRAM使用量を監視して、OOMが正確にいつ発生するかを特定します。

別のターミナルでnvidia-smi -l 1を実行して、毎秒更新されるメモリ使用量を確認します。

クラッシュ直前のピーク使用量に注目してください。これにより、どれだけ削減する必要があるかがわかります。

クラッシュがすぐに発生する場合、モデルロード設定が間違っています。数ステップ後にクラッシュが発生する場合、アクティベーションの蓄積または断片化が問題です。

データセットの問題

データセット内の非常に高解像度の画像は、トレーニング用にロードされる際にOOMを引き起こす可能性があります。

すべての画像がトレーニング解像度以下であることを確認するためにデータセットを前処理します。512x512でトレーニングする場合、4K画像をロードする利点はありません。

アスペクト比が妥当であることを確認してください。非常に幅広いまたは高い画像は、同じ総ピクセル数でも処理中により多くのメモリを必要とする場合があります。

VRAMを使用している他のプロセス

トレーニング前にGPUメモリを消費している他のアプリケーションをチェックしてください。

ウェブブラウザ、Discord、およびその他のGPUアクセラレーションアプリケーションを閉じてください。他の場所で消費される数百MBでも、限界を超えてしまう可能性があります。

複数のPythonプロセスが以前の失敗した実行からVRAMを保持している可能性があります。クリーンな状態のために、Pythonインタープリタまたはシステム全体を再起動してください。

トレーニングフレームワークのバグ

時々、Kohya SSまたは他のトレーニングツールのバグがメモリリークを引き起こします。

トレーニングツールの最新バージョンに更新してください。メモリ関連の修正は更新で一般的です。

ツールのGitHub issuesで症状に一致するレポートを確認してください。他の人が回避策や修正を見つけているかもしれません。

これらの技術的な制限を管理せずにLoRAをトレーニングしたいユーザーには、Apatero.comがプロフェッショナルグレードのハードウェアを使用したクラウドベースのトレーニングを提供しています。トレーニングジョブとデータセットを定義すると、プラットフォームがメモリ管理と最適化を自動的に処理します。

代替のトレーニングアプローチは何ですか？

4090でのトレーニングが問題のままの場合は、これらの代替手段を検討してください。

クラウドトレーニング

48GB以上のVRAMを持つクラウドインスタンスは、メモリの制約を完全に排除します。

RunPod、Vast.ai、およびLambda Labsは時間単位でGPUインスタンスを提供しています。A100 80GBインスタンスは、フル解像度でFlux LoRAを快適にトレーニングします。

典型的なLoRAトレーニング実行のコストは、プロバイダーとGPUタイプによって$5-15です。これは、時々のトレーニングニーズには妥当です。

データセットをアップロードし、トレーニングを実行し、LoRAをダウンロードします。ほとんどのプロジェクトでプロセスには合計1-2時間かかります。

低精度トレーニング

実験的なFP8トレーニングは、BF16よりもさらにメモリを削減します。一部のコミュニティツールがこれをサポートしています。

FP8トレーニングは、BF16ほど品質が実証されていません。重要なプロジェクトに採用する前に、結果を注意深くテストしてください。

追加で30-40%のメモリ節約により、以前は不可能だった設定が機能するようになります。

小さなLoRAバリアント

LoKr、LoHa、および同様の低ランク適応は、標準のLoRAよりも少ないメモリを使用します。

これらのバリアントは一部のトレーニング目標にはうまく機能しますが、複雑なスタイルやコンセプトでは性能が劣る場合があります。

最適化にもかかわらず標準のLoRA設定が問題のままの場合は、代替手段を試してみてください。

よくある質問

なぜトレーニングは毎回正確に1ステップ後にクラッシュするのですか？

これは、特定のメモリ閾値を超えていることを示しています。フォワードパスは収まりますが、バックワードパス中に勾配を追加するとVRAMを超えます。段階的にではなく、勾配チェックポイントと解像度の削減を一緒に有効にしてください。

RTX 4090で1024x1024でトレーニングできますか？

最小ランク、重いチェックポイント、8ビットすべてを含む極端な最適化で理論的には可能です。実際には、512x512での結果は十分に良いため、メモリの苦労は価値がありません。512x512でトレーニングし、1024x1024で生成してください。

バッチサイズ1は大きなバッチよりも悪いLoRAを生成しますか？

大きな違いはありません。勾配累積は同等のトレーニングダイナミクスを提供します。非常に小さなバッチはわずかにノイズが多いと主張する人もいますが、累積された勾配がこれを滑らかにします。品質の違いは他の要因と比較して最小限です。

LoRAが正しくトレーニングされているかどうかをどうやって知りますか？

トレーニング中のサンプル画像が進捗を示します。200-400ステップ後にトリガーワードが生成に影響しているのが見えるはずです。完全なスタイル転送は通常800-1000ステップで現れます。サンプルが変化しないか、すぐに劣化する場合は、学習率を調整してください。

なぜトレーニング中にVRAM使用量が徐々に増加するのですか？

メモリの断片化またはリークが徐々に増加を引き起こします。アロケータは再利用できない小さな断片を作成します。max_split_size_mb環境変数を設定し、カスタムコードにメモリリークがないことを確認してください。

xFormersを使用すべきですか、それともネイティブPyTorchアテンションを使用すべきですか？

xFormersは、ほとんどのトレーニングシナリオでより良いメモリ効率を提供します。ネイティブアテンションは、特定のアーキテクチャでより良く機能する場合があります。xFormersから始めて、問題が発生した場合にのみ切り替えてください。

Flux LoRAにはどのネットワークランクを使用すべきですか？

キャラクターや単純なコンセプトには16から始め、スタイルや複雑な被写体には32を使用してください。高いランクはより多くの詳細をキャプチャしますが、より多くのメモリとより多くのトレーニングデータが必要です。まず低いランクをテストしてください。多くの場合うまく機能します。

トレーニング画像は何枚必要ですか？

キャラクターの場合、10-20枚の良い画像がうまく機能します。スタイルの場合、50-200枚の画像がより良いカバレッジを提供します。量よりも質が重要です。適切にキャプションが付けられた多様な画像は、数百枚の似たようなショットに勝ります。

OOMクラッシュ後にトレーニングを再開できますか？

トレーニング設定でチェックポイントを有効にしていれば、はい。Kohya SSは定期的に進捗を保存します。メモリ設定を修正した後、最後のチェックポイントから再開してください。

開始前に設定がOOMになるかどうかを予測する方法はありますか？

モデルサイズ、バッチサイズ、解像度、および最適化に基づく大まかな推定。accelerate estimate-memoryのようなツールが推定を提供します。しかし、実際のメモリは実装の詳細によって異なるため、常に控えめに始めてください。

結論と推奨ワークフロー

RTX 4090でのFlux LoRAトレーニングは、慎重なメモリ管理が必要ですが、適切に設定されれば優れた結果を生み出します。重要な設定は、勾配チェックポイント、累積付きのバッチサイズ1、512x512解像度、および混合精度トレーニングです。

このガイドで提供された控えめな設定から始めてください。完全なトレーニング実行にコミットする前に、100-200ステップの短いテストを実行して安定性を確認してください。

トレーニング中のサンプルを監視してください。オーバートレーニングから品質が劣化し始める前に、LoRAがターゲットコンセプトをキャプチャしたら停止してください。

メモリ制限に常に苦労している場合や、24GBが許可するよりも高品質な設定が必要な場合は、クラウドトレーニングが実用的な代替手段を提供します。Apatero.comのようなサービスにより、クラウドインフラストラクチャを自分で管理せずにアクセスできます。

あなたのRTX 4090は、ローカルLoRAトレーニングに対応したハードウェアです。適切な設定で、自分のマシンですべてを維持しながら、特定のキャラクター、スタイル、およびコンセプト用のカスタムFlux LoRAをトレーニングできます。