コンシューマーハードウェアでの QLoRA を用いた FLUX.1-dev のファインチューニング
Hugging Face は、QLoRA を用いた FLUX.1-dev の効率的なファインチューニング手法を詳述しており、ピークメモリ使用量を 10 GB 未満の VRAM に抑えてコンシューマーグレードのハードウェア上で実行できるようにしています。4 ビット量子化と Low-Rank Adaptation(LoRA)を組み合わせることで、NVIDIA RTX 4090 のような単一 GPU 上で高性能拡散モデルをカスタマイズでき、エンタープライズ向けハードウェアは不要です。
FLUX.1-dev のアーキテクチャとファインチューニングの焦点
FLUX.1-dev は主に 3 つのコンポーネントから構成されます:テキストエンコーダ(CLIP と T5)、Flux Transformer(本体モデル)、および変分オートエンコーダ(VAE)。効率を最大化するために、QLoRA アプローチは Transformer コンポーネント のみのファインチューニングに限定し、テキストエンコーダと VAE はトレーニング中は凍結されたままです。
QLoRA のメモリ最適化手法
10 GB 未満のメモリフットプリントを実現するために、diffusers のトレーニングパイプラインに以下の最適化技術が組み込まれています:
量子化と適応
- QLoRA(Quantized LoRA):ベースモデルは
bitsandbytesを介して 4 ビット量子化形式(NF4)でロードされ、ベースモデルのメモリ使用量が大幅に削減されます。その上で LoRA アダプタが FP16 または BF16 精度で学習されます。 - LoRA(Low-Rank Adaptation):全重み行列を更新する代わりに、LoRA は 2 つの小さな低ランク行列($A$ と $B$)を学習し、学習可能パラメータ数を大幅に減らします。デモ設定では、総パラメータ 11,906,077,760 のうち 4,669,440 が学習可能です。
計算効率化
- 8-bit AdamW Optimizer:ブロック単位の量子化でオプティマイザ状態を 8 ビット精度で保存し、標準の FP32 AdamW に比べてオプティマイザメモリ使用量を約 75% 削減します。
- Gradient Checkpointing:特定のチェックポイント活性化のみを保存し、逆伝播時に他を再計算することで計算とメモリのトレードオフを実現します。
- Cache Latents:トレーニング画像は VAE エンコーダを通して事前に処理され、トレーニング開始前にキャッシュされます。これにより冗長な VAE 計算が排除され、トレーニングフェーズ中は VAE を GPU メモリから完全に除外できます。
- Pre-computing Text Embeddings:CLIP と T5 テキストエンコーダの出力はトレーニング前に一度だけキャッシュされ、ファインチューニング中にテキストエンコーダが GPU メモリを占有することを防ぎます。
NVIDIA RTX 4090 におけるパフォーマンスベンチマーク
NVIDIA RTX 4090(24 GB VRAM)を使用し、Alphonse Mucha スタイルのデータセット(512×768 解像度、バッチサイズ 1、rank 4)で FLUX.1-dev をファインチューニングした際のメモリと時間の指標は以下の通りです:
| 手法 | ピーク VRAM 使用量 | トレーニング時間(700 ステップ) |
|---|---|---|
| QLoRA | 約 9 GB | 約 41 分 |
| BF16 LoRA | 26 GB | 記載なし |
| BF16 Full Fine-tuning | 約 120 GB(推定) | 記載なし |
計算能力 8.9 以上の NVIDIA GPU(例:H100、RTX 4090)を搭載しているユーザーは、torchao を用いた FP8 トレーニング によりさらに速度が最適化されます。H100 SXM GPU での実行では、ピークメモリ使用量が 36.57 GB、700 ステップのトレーニングが約 20 分で完了しました。
学習済みアダプタの推論戦略
LoRA アダプタが学習されたら、以下の 2 つの方法で利用できます:
1. LoRA アダプタをロード
アダプタはベースモデルの上にロードされます。これにより最大の柔軟性が得られ、set_adapters() を使用して異なるスタイル間を切り替えたり、複数のアダプタを組み合わせたりでき、ベースモデルを再ロードする必要がありません。
2. LoRA をベースモデルにマージ
LoRA 重みをベースモデルに統合します。この方法は推論効率が最も高く、アダプタ重みのメモリオーバーヘッドがなくなり、マージ後のモデルをさらに再量子化してメモリ節約が可能です。
ハードウェアのアクセシビリティ
RTX 4090 向けに最適化されていますが、このワークフローはより手軽なハードウェアでも動作します。Google Colab の T4 GPU でもファインチューニングは可能ですが、同じ 700 ステップを実行するのに約 4 時間かかり、RTX 4090 の 41 分に比べて大幅に時間がかかります。