ParaAttention

FLUX.1-dev や HunyuanVideo などの大規模な画像・動画生成モデルは、そのサイズゆえに、リアルタイムアプリケーションやデプロイにおける推論の課題となり得ます。

ParaAttention は、context parallelism と first block cache を実装したライブラリであり、他の技術（torch.compile、fp8 動的量子化）と組み合わせることで、推論を高速化できます。

このガイドでは、NVIDIA L20 GPU上でFLUX.1-devとHunyuanVideoにParaAttentionを適用する方法を示します。ベースラインベンチマークには最適化は適用されていませんが、HunyuanVideoについてはメモリ不足エラーを回避するために例外として適用しています。

ベースラインのベンチマークによると、FLUX.1-dev は 28 ステップで 1024x1024 解像度の画像を 26.36 秒で生成でき、HunyuanVideo は 30 ステップで 720p 解像度の 129 フレームを 3675.71 秒で生成できます。

コンテキスト並列性によるさらに高速な推論を実現するには、特にGPUの台数が多い場合、NVLink対応のNVIDIA A100またはH100 GPU（利用可能な場合）の使用をお試しください。

最初のブロックキャッシュ

モデル内のtransformersブロックの出力をキャッシュし、次の推論ステップで再利用することで、計算コストを削減し、推論を高速化できます。

ただし、生成される画像や動画の品質を確保するために、キャッシュを再利用するタイミングを判断するのは容易ではありません。ParaAttention は、最初の Transformer ブロックの出力の残差差を直接使用して、モデル出力間の差を近似します。差が十分に小さい場合、以前の推論ステップの残差差が再利用されます。言い換えれば、デノイジングステップがスキップされるということです。

これにより、FLUX.1-dev および HunyuanVideo の推論で 2 倍の高速化を実現し、非常に高い品質を維持します。

Cache in Diffusion Transformer — AdaCacheの仕組み、First Block Cacheはそのバリアントです

FLUX-1.dev

HunyuanVideo

FLUX.1-dev に first block cache を適用するには、以下に示すように apply_cache_on_pipe を呼び出します。0.08 は FLUX モデルのデフォルトの residual difference 値です。

import time
import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe, residual_diff_threshold=0.08)

# Enable memory savings
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()

begin = time.time()
image = pipe(
    "A cat holding a sign that says hello world",
    num_inference_steps=28,
).images[0]
end = time.time()
print(f"Time: {end - begin:.2f}s")

print("Saving image to flux.png")
image.save("flux.png")

最適化	Original	FBCache rdt=0.06	FBCache rdt=0.08	FBCache rdt=0.10	FBCache rdt=0.12
プレビュー
実時間 (秒)	26.36	21.83	17.01	16.00	13.78

First Block Cache は、ベースラインと比較して推論速度を 17.01 秒に短縮し、1.55 倍の高速化を実現しながら、品質の低下をほぼゼロに抑えました。

fp8量子化

動的量子化を伴うfp8は、推論の高速化とメモリ使用量の削減をさらに実現します。8ビットのNVIDIA Tensor Coresを使用するには、アクティベーションと重みの両方を量子化する必要があります。

float8_weight_onlyとfloat8_dynamic_activation_float8_weightを使用して、テキストエンコーダーとtransformerモデルを量子化します。

デフォルトの量子化手法はテンソルごとの量子化ですが、GPUが行ごとの量子化をサポートしている場合は、より高い精度を得るためにこれを試すこともできます。

以下のコマンドでtorchaoをインストールします。

pip3 install -U torch torchao

mode="max-autotune-no-cudagraphs" または mode="max-autotune" を指定した torch.compile は、パフォーマンスを最大化する最適なカーネルを選択します。モデルが初めて呼び出される場合、コンパイルに長時間かかることがありますが、一度コンパイルが完了すればその価値は十分にあります。

この例では、transformerモデルのみを量子化していますが、テキストエンコーダーも量子化することで、メモリ使用量をさらに削減できます。

動的量子化はモデルの出力分布を大きく変える可能性があるため、residual_diff_thresholdをより大きな値に変更する必要があります。

FLUX-1.dev

HunyuanVideo

コンテキスト並列

Context Parallelism は推論を並列化し、複数の GPU によってスケーリングします。ParaAttention の構成要素としての設計により、Context Parallelism を First Block Cache および動的量子化と組み合わせることができます。

複数のGPUで推論をスケーリングする方法の詳細な手順や例については、ParaAttentionリポジトリを参照してください。

推論プロセスを永続的かつサービス可能にする必要がある場合は、torch.multiprocessing を使用して独自の推論プロセッサを作成することをお勧めします。これにより、プロセスの起動、モデルの読み込み、再コンパイルに伴うオーバーヘッドを排除できます。

FLUX-1.dev

HunyuanVideo

以下のコード例では、First Block Cache、fp8 動的量子化、torch.compile、および Context Parallelism を組み合わせることで、最速の推論速度を実現します。

import time
import torch
import torch.distributed as dist
from diffusers import FluxPipeline

dist.init_process_group()

torch.cuda.set_device(dist.get_rank())

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.context_parallel import init_context_parallel_mesh
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae

mesh = init_context_parallel_mesh(
    pipe.device.type,
    max_ring_dim_size=2,
)
parallelize_pipe(
    pipe,
    mesh=mesh,
)
parallelize_vae(pipe.vae, mesh=mesh._flatten())

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(
    pipe,
    residual_diff_threshold=0.12,  # Use a larger value to make the cache take effect
)

from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only

quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
torch._inductor.config.reorder_for_compute_comm_overlap = True
pipe.transformer = torch.compile(
   pipe.transformer, mode="max-autotune-no-cudagraphs",
)

# Enable memory savings
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())

for i in range(2):
    begin = time.time()
    image = pipe(
        "A cat holding a sign that says hello world",
        num_inference_steps=28,
        output_type="pil" if dist.get_rank() == 0 else "pt",
    ).images[0]
    end = time.time()
    if dist.get_rank() == 0:
        if i == 0:
            print(f"Warm up time: {end - begin:.2f}s")
        else:
            print(f"Time: {end - begin:.2f}s")

if dist.get_rank() == 0:
    print("Saving image to flux.png")
    image.save("flux.png")

dist.destroy_process_group()

run_flux.py に保存し、torchrun で起動します。

# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=2 run_flux.py

2基のNVIDIA L20 GPUを使用した場合、推論速度はベースラインと比較して8.20秒に短縮され、3.21倍高速化されます。4基のL20を使用した場合、推論速度は3.90秒となり、6.75倍高速化されます。

以下のコード例では、First Block Cache と Context Parallelism を組み合わせることで、最速の推論速度を実現します。

import time
import torch
import torch.distributed as dist
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

dist.init_process_group()

torch.cuda.set_device(dist.get_rank())

model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    dtype=torch.bfloat16,
    revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.float16,
    revision="refs/pr/18",
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.context_parallel import init_context_parallel_mesh
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae

mesh = init_context_parallel_mesh(
    pipe.device.type,
)
parallelize_pipe(
    pipe,
    mesh=mesh,
)
parallelize_vae(pipe.vae, mesh=mesh._flatten())

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe)

# from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only
#
# torch._inductor.config.reorder_for_compute_comm_overlap = True
#
# quantize_(pipe.text_encoder, float8_weight_only())
# quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
# pipe.transformer = torch.compile(
#    pipe.transformer, mode="max-autotune-no-cudagraphs",
# )

# Enable memory savings
pipe.vae.enable_tiling()
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())

for i in range(2):
    begin = time.time()
    output = pipe(
        prompt="A cat walks on the grass, realistic",
        height=720,
        width=1280,
        num_frames=129,
        num_inference_steps=1 if i == 0 else 30,
        output_type="pil" if dist.get_rank() == 0 else "pt",
    ).frames[0]
    end = time.time()
    if dist.get_rank() == 0:
        if i == 0:
            print(f"Warm up time: {end - begin:.2f}s")
        else:
            print(f"Time: {end - begin:.2f}s")

if dist.get_rank() == 0:
    print("Saving video to hunyuan_video.mp4")
    export_to_video(output, "hunyuan_video.mp4", fps=15)

dist.destroy_process_group()

run_hunyuan_video.py に保存し、torchrun で起動します。

# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=8 run_hunyuan_video.py

8基のNVIDIA L20 GPUを使用した場合、推論速度はベースラインと比較して649.23秒に短縮され、5.66倍高速化されました。

ベンチマーク

FLUX-1.dev

HunyuanVideo

GPU タイプ	GPUの数	最適化	実時間 (秒)	高速化
NVIDIA L20	1	ベースライン	26.36	1.00x
NVIDIA L20	1	FBCache (rdt=0.08)	17.01	1.55x
NVIDIA L20	1	FP8 DQ	13.40	1.96x
NVIDIA L20	1	FBCache (rdt=0.12) + FP8 DQ	7.56	3.48x
NVIDIA L20	2	FBCache (rdt=0.12) + FP8 DQ + CP	4.92	5.35x
NVIDIA L20	4	FBCache (rdt=0.12) + FP8 DQ + CP	3.90	6.75x

GPU タイプ	GPUの数	最適化	実時間 (秒)	高速化
NVIDIA L20	1	ベースライン	3675.71	1.00x
NVIDIA L20	1	FBCache	2271.06	1.62x
NVIDIA L20	2	FBCache + CP	1132.90	3.24x
NVIDIA L20	4	FBCache + CP	718.15	5.12x
NVIDIA L20	8	FBCache + CP	649.23	5.66x

Update on GitHub

Diffusers

ParaAttention

最初のブロックキャッシュ

fp8量子化

コンテキスト並列

ベンチマーク