推論を高速化

Diffusionモデルは、生成が「ステップ」と呼ばれる一定の回数にわたり、ノイズが徐々に洗練されて画像や動画になる反復プロセスであるため、推論が遅くなります。このプロセスを高速化するには、異なるschedulersの実験、モデル重みの精度を下げた高速な計算、よりメモリ効率の高いアテンション機構の利用などが考えられます。

これらの技術を組み合わせて使用することで、単一の技術のみを使用する場合よりも推論を高速化できます。

このガイドでは、推論を高速化する方法について説明します。

モデルのデータ型

モデルの重みの精度とデータ型は推論速度に影響します。精度が高いほど、読み込みに必要なメモリ量が増え、計算に要する時間も長くなるためです。PyTorch はデフォルトで float32（フル精度）でモデルの重みをロードするため、データ型を変更することは、推論を素早く高速化する簡単な方法です。

bfloat16

float16

TensorFloat-32

スケーリングドットプロダクトアテンション

Memory-efficient attention は、推論速度とメモリ使用量の両方を最適化します!

Scaled dot product attention (SDPA) は、FlashAttention、xFormers、およびネイティブの C++ 実装という複数のアテンションバックエンドを実装しています。ハードウェアに応じて最適なバックエンドを自動的に選択します。

PyTorch >= 2.0 を使用している場合、SDPA はデフォルトで有効化され、コードに追加の変更は不要です。ただし、独自のバックエンドを選択したい場合は、他のアテンションバックエンドを試すこともできます。以下の例では、torch.nn.attention.sdpa_kernel コンテキストマネージャーを使用して、効率的なアテンションを有効にしています。

from torch.nn.attention import SDPBackend, sdpa_kernel
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
  image = pipeline(prompt, num_inference_steps=30).images[0]

torch.compile

torch.compile は、PyTorch のコードと演算を最適化されたカーネルにコンパイルすることで推論を高速化します。Diffusers は通常、UNet、transformer、VAE など、計算負荷の高いモデルをコンパイルします。

最大速度を実現するために、以下のコンパイラ設定を有効にしてください（その他のオプションについては、完全なリストを参照してください）。

import torch
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

UNetとVAEを読み込み、コンパイルします。選択できるモードはいくつかありますが、"max-autotune"はCUDAグラフへコンパイルすることで最速の速度を最適化します。CUDAグラフは、単一のCPU操作で複数のGPU操作を発行することにより、オーバーヘッドを実質的に削減します。

PyTorch 2.3.1 では、torch.compile のキャッシュ動作を制御できます。これは、複数のコンパイルフラグをグリッドサーチして最適な設定を見つける "max-autotune" などのコンパイルモードにおいて特に有益です。詳細は、torch.compile におけるコンパイルタイムキャッシュチュートリアルを参照してください。

メモリレイアウトをchannels_lastに変更すると、メモリと推論速度が最適化されます。

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.decode = torch.compile(
    pipeline.vae.decode,
    mode="max-autotune",
    fullgraph=True
)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

コンパイルは初回に時間がかかりますが、一度コンパイルが完了すると、その後の処理は大幅に高速化されます。コンパイル済みパイプラインは、同じ種類の推論操作にのみ使用することをお勧めします。異なる画像サイズに対してコンパイル済みパイプラインを呼び出すと、コンパイルが再実行され、これが遅く非効率になるためです。

動的形状コンパイル

より良いサポートを得るために、常に PyTorch の nightly バージョンを使用してください。

torch.compile は入力形状と条件をトラッキングし、これらが異なる場合はモデルを再コンパイルします。例えば、1024x1024 の解像度の画像でモデルをコンパイルし、異なる解像度の画像で使用した場合、再コンパイルがトリガーされます。

再コンパイルを回避するために、dynamic=True を追加し、条件が変更された際に再コンパイルを回避するためのより動的なカーネルの生成を試みてください。

+ torch.fx.experimental._config.use_duck_shape = False
+ pipeline.unet = torch.compile(
    pipeline.unet, fullgraph=True, dynamic=True
)

use_duck_shape=False を指定すると、コンパイラに、同じである入力サイズを表すのに同じシンボリック変数を使用すべきかどうかを指示します。詳細については、このコメントを参照してください。

すべてのモデルが、そのままの状態で動的コンパイルの恩恵を受けられるとは限りません。変更が必要な場合があります。AuraFlowPipelineの実装を改善し、動的コンパイルの恩恵を受けられるようにしたこのPRを参照してください。

Diffusersモデルで動的コンパイルが期待通りに動作しない場合は、お気軽にissueを作成してください。

地域別コンパイル

Regional compilationは、モデルの小さく頻繁に繰り返されるブロック（通常はtransformer層）のみをコンパイルすることで、コールドスタートのレイテンシを削減します。これにより、以降のすべての出現箇所に対してコンパイル済み成果物を再利用できます。多くの拡散アーキテクチャにおいて、これはグラフ全体のコンパイルと同等のランタイム高速化を実現し、コンパイル時間を8〜10分の1に短縮します。

以下に示すように、compile_repeated_blocks() メソッドは torch.compile をラップするヘルパーであり、transformer モデルなどのコンポーネントに適用できます。

# pip install -U diffusers
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"

# compile only the repeated transformer layers inside the UNet
pipeline.unet.compile_repeated_blocks(fullgraph=True)

新しいモデルでリージョナルコンパイルを有効にするには、コンパイル対象のブロックのクラス名（文字列）を含む_repeated_blocks属性をモデルクラスに追加します。

class MyUNet(ModelMixin):
    _repeated_blocks = ("Transformer2DModel",)  # ← compiled by default

地域別のコンパイルの例については、PRを参照してください。

Accelerate には、モデル内でコンパイルする候補ブロックを自動的に選択する compile_regions メソッドもあります。残りのグラフは個別にコンパイルされます。コンパイル対象のブロックを指定したり、コンパイルフラグを調整したりするオプションが少ないため、素早い実験に便利です。

# pip install -U accelerate
import torch
from diffusers import StableDiffusionXLPipeline
from accelerate.utils import compile_regions

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)

compile_repeated_blocks() は意図的に明示的な設計です。_repeated_blocks に反復するブロックを列挙し、ヘルパーはそれらのブロックのみをコンパイルします。これにより、1 行のコードでキャッシュ再利用について予測可能な動作と容易な推論が可能になります。

グラフのブレイク

torch.compile で fullgraph=True を指定することは、基盤となるモデルにグラフの中断が発生しないことを保証するために重要です。これにより、パフォーマンスの低下なしで torch.compile の恩恵を受けることができます。UNet と VAE については、戻り値の変数へのアクセス方法が変更されます。

- latents = unet(
-   latents, timestep=timestep, encoder_hidden_states=prompt_embeds
-).sample

+ latents = unet(
+   latents, timestep=timestep, encoder_hidden_states=prompt_embeds, return_dict=False
+)[0]

GPU同期

step() 関数は、デノイザーが予測を行うたびにスケジューラーに対して呼び出され、sigmas 変数がインデックス付けられます。GPU に配置すると、CPU と GPU 間の通信同期によりレイテンシが発生します。デノイザーがすでにコンパイルされている場合、この影響はより顕著になります。

一般的に、sigmasは通信の同期とレイテンシを避けるために、CPU上に保持しておく必要があります。

拡散モデルでtorch.compileによるパフォーマンスの最大化については、torch.compile と Diffusers: ピークパフォーマンスを実現する実践ガイドというブログ記事をご参照ください。

ベンチマーク

コンパイル済みパイプラインの推論レイテンシとメモリ使用量のデータを確認するには、diffusers/benchmarks データセットを参照してください。

diffusers-torchao リポジトリには、Flux と CogVideoX のコンパイル済みバージョンに関するベンチマーク結果も含まれています。

カーネル

Kernels は、Hub 上で最適化された計算カーネルを構築・配布・読み込むためのライブラリです。attention カーネルや、RMSNorm、GEGLU、RoPE、AdaLN などの演算用のカスタム CUDA カーネルに対応しています。

Diffusers Pipeline Integrationガイドでは、add cuda-kernelsスキルを使用してカーネルを統合する方法を示しています。このスキルにより、ClaudeやCodexなどのエージェントが、特定のモデルやお使いのハードウェア向けのカスタムカーネルを作成できるようになります。

Install the add cuda-kernels skill to teach an agent how to write a kernel. The Custom kernels for all from Codex and Claude blog post covers this in more detail.

例えば、add cuda-kernels スキルによって生成されたカスタム RMSNorm カーネルに torch.compile を組み合わせることで、H100 上で LTX-Video の生成を 1.43 倍高速化できます。

動的量子化

Dynamic quantizationは、精度を低減してより高速な数学演算を可能にすることで、推論速度を向上させます。このタイプの量子化は、固定のスケーリング係数を使用するのではなく、実行時のデータに基づいてアクティベーションのスケーリング方法を決定します。その結果、スケーリング係数がデータにより正確に整合するようになります。

以下の例では、torchao ライブラリを使用して、UNet と VAE に動的 int8 量子化を適用します。

Diffusersのtorchao統合の使用方法については、torchaoのドキュメントを参照してください。

最大速度を達成するためにコンパイラタグを設定します。

import torch
from torchao import apply_dynamic_quant
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True

dynamic_quant_filter_fn を使用して、動的量子化の恩恵を受けない UNet と VAE の一部の線形層を除外します。

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

apply_dynamic_quant(pipeline.unet, dynamic_quant_filter_fn)
apply_dynamic_quant(pipeline.vae, dynamic_quant_filter_fn)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

融合された射影行列

fuse_qkv_projections メソッドは実験的な機能であり、サポートは主に Stable Diffusion パイプラインに限定されています。他のパイプラインでこれを有効にする方法については、この PR をご覧ください。

Q、K、Vの射影行列によって表される3つの部分空間に、アテンションブロック内で入力が射影されます。これらの射影は通常、個別に計算されますが、これらを水平方向に結合して単一の行列にし、1ステップで射影を実行することもできます。これにより、入力射影の行列積のサイズが増大し、量子化の影響も向上します。

pipeline.fuse_qkv_projections()

リソース

Presenting Flux Fast: Making Flux go brrr on H100s のブログ記事を読み、TorchInductor と AOTInductor をこれらの最適化すべてと組み合わせることで、flux-fast のレシピを使用して約2.5倍の高速化を実現する方法について詳しく学びましょう。

これらのレシピは、AMD ハードウェアおよび Flux.1 Kontext Dev をサポートしています。
torch.compile を使用してパフォーマンスを最大化するには、torch.compile と Diffusers：ピークパフォーマンスに向けた実践ガイドというブログ記事をお読みください。

Update on GitHub