コンシューマーハードウェアでの 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 分に比べて大幅に時間がかかります。

Sources