Stable Diffusion XL (SDXL) の推論速度とメモリ使用量の最適化

Hugging Faceは、Stable Diffusion XL (SDXL) を、特にコンシューマー向けGPUでより実用的にするための、一連のシンプルな最適化手法を詳細に説明しています。低精度化、メモリ効率の高いAttention、および様々なオフロード技術を実装することで、メモリ使用量を28GBから11.47GBに削減し、推論レイテンシを72.2秒から約10.3秒まで短縮できます。

推論速度の最適化

拡散モデルにおいて、速度の最適化は極めて重要です。なぜなら、画像生成の反復的な性質上、望ましい結果を得るために何度も実行が必要になる場合が多いからです。以下の手法は、レイテンシの削減に焦点を当てています：

低精度化 (fp16)

標準的な float32 (fp32) の代わりに float16 (fp16) を使用すると、メモリ使用量が半分になり、計算速度が向上します。これは、fp16がより狭い範囲の浮動小数点を扱い、最新のGPUハードウェアによってより良くサポートされているためです。テストされたSDXLパイプラインでは、fp16への切り替えにより、メモリ使用量が21.7GBに減少し、推論時間が72.2秒から14.8秒に短縮されました。

メモリ効率の高いAttention (SDPA)

TransformerのAttentionブロックは、入力シーケンスの長さに応じてメモリ要件が二次関数的に増加するため、メモリのボトルネックになる可能性があります。PyTorch 2.0では、Scaled Dot Product Attention (SDPA) が導入されました。これは、Flash Attentionとメモリ効率の高いAttention (xFormers) の融合実装を提供します。🤗Diffusersでは、PyTorch ≥ 2.0においてSDPAがデフォルトで有効になっており、21.7GBのメモリ使用量を維持したまま、推論時間をさらに11.4秒まで短縮します。

torch.compile による JIT コンパイル

PyTorch 2.0の torch.compile APIを使用すると、PyTorchコードを最適化されたカーネルにJust-In-Time (JIT) コンパイルできます。SDXLのUNetを torch.compile (mode="reduce-overhead" を使用) でラップすることで、推論時間はさらに10.2秒まで改善されます。なお、最初のコンパイル実行は低速ですが、その後の呼び出しは大幅に高速化されます。

モデルのメモリフットプリントの削減

SDXLは従来のStable Diffusionモデルよりも約3倍大きく（3.5BパラメータのUNet）、これをVRAMに収めることが主要な課題となります。いくつかの手法でメモリフットプリントを削減できます：

CPU Offloading

モデルのオフロードは、GPU上でアクティブに必要とされていないパイプラインのコンポーネントをCPUに移動させます：

Model CPU Offloading: テキストエンコーダとVAEをCPUに保持したまま、UNetをGPUメモリにロードします。これにより、メモリ使用量は20.2GBに減少しました。
Sequential CPU Offloading: 個々のUNetサブモジュールの重みをCPUにオフロードし、フォワードパスの直前にのみGPUにロードします。これによりメモリは19.9GBまで減少しますが、レイテンシは67秒へと大幅に増加します。

VAE Slicing

Variational Autoencoder (VAE) は潜在変数を画像にデコードしますが、このプロセスはバッチサイズに応じてメモリ使用量が増大します。VAE slicingは、入力テンソルを小さなスライスに分割し、複数のステップにわたってデコードします。この最適化により、レイテンシに大きな影響を与えることなく、メモリ使用量を15.4GBに削減できます。

計算のキャッシュ

SDXLは、プロンプトから埋め込み（embeddings）を計算するために2つのテキストエンコーダを使用します。これらの埋め込みは逆拡散プロセスを通じて一定であるため、事前に計算してキャッシュすることができます。埋め込みが生成された後は、テキストエンコーダとトークナイザをGPUメモリから削除できるため、SDPAとfp16を組み合わせた場合のメモリフットプリントは21.9GBになります。

Tiny Autoencoder (TAESD)

標準的なVAEを、madebyollin による蒸留版（約10MB）などの軽量なものに置き換えると、メモリ使用量は15.6GBに減少し、推論レイテンシも低下します。ただし、Tiny Autoencoderは細かいディテールを省略する可能性があるため、主にプレビュー用として推奨されます。

パフォーマンスの要約

A100 GPU (40 GB) でプロンプトあたり4枚の画像を生成するテスト結果から、メモリとレイテンシのトレードオフは以下の通りです：

手法	メモリ (GB)	推論レイテンシ (ms)
最適化なしのパイプライン	28.09	72200.5
fp16	21.72	14800.9
fp16 + SDPA (デフォルト)	21.72	11413.0
デフォルト + `torch.compile`	21.73	10296.7
デフォルト + model CPU offload	20.21	16082.2
デフォルト + sequential CPU offload	19.91	67034.0
デフォルト + VAE slicing	15.40	11232.2
デフォルト + VAE slicing + sequential CPU offload	11.47	66869.2
デフォルト + 事前計算済みテキスト埋め込み	21.85	11909.0
デフォルト + Tiny Autoencoder	15.48	10449.7