Diffusers と PEFT を使用した Flux の高速 LoRA 推論
Hugging Face は Flux.1-Dev テキストツーイメージモデル向けに最適化レシピを開発し、LoRA 推論速度を約 2.3 倍向上させました。この手法は、異なる LoRA アダプタを入れ替える際に発生する再コンパイルの停止という一般的な問題を解決し、レイテンシを犠牲にせず高性能なカスタマイズを可能にします。
LoRA 推論のボトルネックを克服する
LoRA 推論の最適化は困難です。なぜなら、ランクが異なり異なる層を対象とするアダプタをホットスワップすると、通常はモデル構造が変化するからです。標準的なワークフローでは、特定の LoRA を持つモデルに torch.compile を適用すると速度向上が得られますが、別の LoRA に入れ替えるとグラフの再コンパイルがトリガーされ、推論が大幅に遅くなります。
この課題を解決するために、Hugging Face は「ホットスワッピング」メカニズムを利用します。Diffusers で hotswap=True を設定すると、モデル構造は変わらず LoRA アダプタの重みだけが交換されます。これにより再コンパイルが不要となりますが、以下の条件を満たす必要があります。
- Maximum Rank Definition:
max_rankを事前に指定し、プール内で最大のアダプタに対応できるようにします。 - Layer Consistency: 後続の LoRA は、最初にロードした LoRA が対象とした層、またはそのサブセットを対象としなければなりません。
- Text Encoder Limitation: 現時点ではテキストエンコーダの対象化はホットスワップでサポートされていません。
ハイエンド GPU 向け最適化レシピ
NVIDIA H100 GPU などのハイパフォーマンスハードウェア向けに、最適化された推論パイプラインは次の 4 つの主要コンポーネントを組み合わせます。
- Flash Attention 3 (FA3): アテンション機構の効率を向上させます。
torch.compile: 実行グラフを最適化する JIT コンパイラです。- FP8 Quantization: TorchAO が提供する高速・省メモリな量子化ですが、情報はロスします。
- Hotswapping: コンパイルを再実行せずにアダプタを切り替えられます。
パフォーマンスベンチマーク (H100)
| オプション | 時間 (s) | 速度向上 (ベースライン比) | 備考 |
|---|---|---|---|
| Baseline | 7.8910 | – | ベースライン |
| Optimized | 3.5464 | 2.23× | ホットスワップ + コンパイル + FP8 |
| No FP8 | 4.3520 | 1.81× | FP8 量子化なしの最適化 |
| No FA3 | 4.3020 | 1.84× | Flash Attention 3 なしの最適化 |
| Baseline + Compile | 5.0920 | 1.55× | コンパイルは有効だが再コンパイルの停止が発生 |
コンシューマ GPU 向け最適化 (RTX 4090)
Flux.1-Dev を Bfloat16 で実行すると約 33 GB の VRAM が必要となり、RTX 4090 の 24 GB を超えてしまいます。そこで Hugging Face は以下のメモリ削減戦略を採用しました。
- T5 Text Encoder Quantization:
bitsandbytesの NF4 量子化を用いてテキストエンコーダのメモリフットプリントを削減。 - FP8 Quantization: Flux Transformer に適用。
- Regional Compilation:
compile_repeated_blocksを使用し、コンパイル時間とメモリ使用量を削減。
RTX 4090 上で FP8 量子化、torch.compile、T5 の NF4 量子化を組み合わせた結果、ベースラインの 23.6060 秒から 11.5715 秒へと 2.04 倍 の速度向上が得られました。
ホットスワッピングの技術的実装
アダプタの入れ替え時に再コンパイルを防ぐため、実装は次の 2 つの技術的ハードルに対処しています。
- Tensor Conversion: LoRA のスケーリング係数を float から torch tensor に変換。
- Weight Padding: LoRA の重みを
max_rankで定義された最大形状にパディング。これにより属性全体を再割り当てせずに重みデータを置き換えることが可能になります。パディングはゼロで埋められ、数式結果は変わりませんが、過度のパディングは計算速度に若干の影響を与える可能性があります。
このワークフローを実装する際は、LoRA のロード順序が重要です。アダプタが互いに重ならない層を対象とする場合、まずすべての必要層の合併集合を対象としたダミー LoRA を作成してから他の LoRA をロードしてください。