# ParaAttention

    

    

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

[ParaAttention](https://github.com/chengzeyi/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 秒で生成できます。

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

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

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

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

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

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

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

```python
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 |
| - | - | - | - | - | - |
| プレビュー | ![Original](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-original.png) | ![FBCache rdt=0.06](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.06.png) | ![FBCache rdt=0.08](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.08.png) | ![FBCache rdt=0.10](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.10.png) | ![FBCache rdt=0.12](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.12.png) |
| 実時間 (秒) | 26.36 | 21.83 | 17.01 | 16.00 | 13.78 |

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

HunyuanVideo に First Block Cache を適用するには、`apply_cache_on_pipe` を以下のように使用します。0.06 は HunyuanVideo モデルのデフォルトの残差差の値です。

```python
import time
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

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.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe, residual_diff_threshold=0.6)

pipe.vae.enable_tiling()

begin = time.time()
output = pipe(
    prompt="A cat walks on the grass, realistic",
    height=720,
    width=1280,
    num_frames=129,
    num_inference_steps=30,
).frames[0]
end = time.time()
print(f"Time: {end - begin:.2f}s")

print("Saving video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)
```

  
  お使いのブラウザは video タグに対応していません。

 FBCacheなしのHunyuanVideo

  
  お使いのブラウザは video タグに対応していません。

 FBCache を使用した HunyuanVideo

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

## fp8量子化

動的量子化を伴うfp8は、推論の高速化とメモリ使用量の削減をさらに実現します。8ビットの[NVIDIA Tensor Cores](https://www.nvidia.com/en-us/data-center/tensor-cores/)を使用するには、アクティベーションと重みの両方を量子化する必要があります。

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

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

以下のコマンドで[torchao](https://github.com/pytorch/ao/tree/main)をインストールします。

```bash
pip3 install -U torch torchao
```

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

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

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

```python
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.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())
pipe.transformer = torch.compile(
   pipe.transformer, mode="max-autotune-no-cudagraphs",
)

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

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

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

fp8 動的量子化と torch.compile を使用することで、ベースラインと比較して推論速度が 7.56 秒に短縮され、3.48 倍高速化されました。

```python
import time
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

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.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

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()
# pipe.enable_sequential_cpu_offload()

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,
    ).frames[0]
    end = time.time()
    if i == 0:
        print(f"Warm up time: {end - begin:.2f}s")
    else:
        print(f"Time: {end - begin:.2f}s")

print("Saving video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)
```

NVIDIA L20 GPU はメモリが 48GB しかなく、コンパイル後、`enable_model_cpu_offload` が呼び出されていない場合、HunyuanVideo は高解像度かつ大量のフレームで実行される際に非常に大きな活性化テンソルを生成するため、メモリ不足 (OOM) エラーが発生する可能性があります。メモリが 80GB 未満の GPU では、OOM エラーを回避するために、解像度とフレーム数を減らすことを試みてください。

大規模な動画生成モデルでは、全結合層よりもアテンション計算がボトルネックになることが一般的です。これらのモデルは、量子化やtorch.compileから大きな恩恵を受けません。

## コンテキスト並列

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

> [!TIP]
> 複数のGPUで推論をスケーリングする方法の詳細な手順や例については、[ParaAttention](https://github.com/chengzeyi/ParaAttention/tree/main)リポジトリを参照してください。

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

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

```python
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](https://pytorch.org/docs/stable/elastic/run.html) で起動します。

```bash
# 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 を組み合わせることで、最速の推論速度を実現します。

```python
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](https://pytorch.org/docs/stable/elastic/run.html) で起動します。

```bash
# 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倍高速化されました。

## ベンチマーク

| 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 |