diffusers/main/id/optimization/para_attn.md

ParaAttention

Model generasi gambar dan video berukuran besar, seperti FLUX.1-dev dan HunyuanVideo, dapat menjadi tantangan inferensi untuk aplikasi real-time dan deployment karena ukurannya.

ParaAttention adalah pustaka yang mengimplementasikan context parallelism dan first block cache, serta dapat dikombinasikan dengan teknik lain (torch.compile, fp8 dynamic quantization) untuk mempercepat inferensi.

Panduan ini akan menunjukkan cara menerapkan ParaAttention pada FLUX.1-dev dan HunyuanVideo di GPU NVIDIA L20. Tidak ada optimasi yang diterapkan pada benchmark baseline kami, kecuali untuk HunyuanVideo guna menghindari error out-of-memory.

Benchmarks baseline kami menunjukkan bahwa FLUX.1-dev mampu menghasilkan gambar dengan resolusi 1024x1024 dalam 28 langkah selama 26,36 detik, dan HunyuanVideo mampu menghasilkan 129 frame pada resolusi 720p dalam 30 langkah selama 3675,71 detik.

Untuk inferensi yang lebih cepat dengan paralelisme konteks, coba gunakan GPU NVIDIA A100 atau H100 (jika tersedia) dengan dukungan NVLink, terutama ketika ada sejumlah besar GPU.

Cache Blok Pertama

Mencache output dari blok transformers dalam model dan menggunakannya kembali pada langkah inferensi berikutnya mengurangi biaya komputasi dan membuat inferensi lebih cepat.

Namun, sulit untuk menentukan kapan harus menggunakan ulang cache guna memastikan kualitas gambar atau video yang dihasilkan. ParaAttention secara langsung menggunakan selisih residual dari output blok transformer pertama untuk mengaproksimasi perbedaan di antara output model. Ketika perbedaannya cukup kecil, selisih residual dari langkah inferensi sebelumnya digunakan ulang. Dengan kata lain, langkah denoising dilewati.

Hal ini mencapai percepatan 2x pada inferensi FLUX.1-dev dan HunyuanVideo dengan kualitas yang sangat baik.

Cara kerja AdaCache, First Block Cache adalah variasinya

Untuk menerapkan first block cache pada FLUX.1-dev, panggil apply_cache_on_pipe seperti yang ditunjukkan di bawah. 0.08 adalah nilai residual difference default untuk model FLUX.

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

Optimasi	Asli	FBCache rdt=0.06	FBCache rdt=0.08	FBCache rdt=0.10	FBCache rdt=0.12
Pratinjau
Waktu Nyata (s)	26.36	21.83	17.01	16.00	13.78

First Block Cache mengurangi kecepatan inferensi menjadi 17,01 detik dibandingkan baseline, atau 1,55x lebih cepat, sambil mempertahankan penurunan kualitas yang nyaris nol.

Untuk menerapkan First Block Cache pada HunyuanVideo, gunakan apply_cache_on_pipe seperti yang ditunjukkan di bawah. 0.06 adalah nilai residual difference default untuk model HunyuanVideo.

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)

Browser Anda tidak mendukung tag video.

HunyuanVideo tanpa FBCache

Browser Anda tidak mendukung tag video.

HunyuanVideo dengan FBCache

First Block Cache mengurangi kecepatan inferensi menjadi 2271,06 detik dibandingkan baseline, atau 1,62x lebih cepat, sambil mempertahankan penurunan kualitas yang nyaris nol.

kuantisasi fp8

fp8 dengan kuantisasi dinamis mempercepat inferensi lebih lanjut dan mengurangi penggunaan memori. Baik aktivasi maupun bobot harus dikuantisasi untuk menggunakan NVIDIA Tensor Cores 8-bit.

Gunakan float8_weight_only dan float8_dynamic_activation_float8_weight untuk mengkuantisasi model text encoder dan transformer.

Metode kuantisasi default adalah kuantisasi per tensor, tetapi jika GPU Anda mendukung kuantisasi per baris, Anda juga dapat mencobanya untuk akurasi yang lebih baik.

Instal torchao dengan perintah di bawah ini.

pip3 install -U torch torchao

torch.compile dengan mode="max-autotune-no-cudagraphs" atau mode="max-autotune" memilih kernel terbaik untuk performa. Kompilasi dapat memakan waktu lama jika ini adalah pertama kalinya model dipanggil, tetapi hasilnya sepadan setelah model dikompilasi.

Contoh ini hanya mengkuantisasi model transformer, tetapi Anda juga dapat mengkuantisasi text encoder untuk mengurangi penggunaan memori lebih lanjut.

Kuantisasi dinamis dapat mengubah distribusi output model secara signifikan, sehingga Anda perlu mengubah residual_diff_threshold menjadi nilai yang lebih besar agar perubahan tersebut berlaku.

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

Kuantisasi dinamis fp8 dan torch.compile mengurangi kecepatan inferensi menjadi 7,56 detik dibandingkan baseline, atau 3,48x lebih cepat.

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)

GPU NVIDIA L20 hanya memiliki memori 48GB dan dapat mengalami error out-of-memory (OOM) setelah kompilasi jika enable_model_cpu_offload tidak dipanggil, karena HunyuanVideo memiliki tensor aktivasi yang sangat besar saat dijalankan dengan resolusi tinggi dan jumlah frame yang banyak. Untuk GPU dengan memori kurang dari 80GB, Anda dapat mencoba mengurangi resolusi dan jumlah frame untuk menghindari error OOM.

Model generasi video besar biasanya dibatasi oleh komputasi perhatian, bukan oleh lapisan terhubung penuh. Model-model ini tidak mendapatkan manfaat signifikan dari kuantisasi dan torch.compile.

Paralelisme Konteks

Context Parallelism memparalelkan inferensi dan berskala dengan beberapa GPU. Desain komposisional ParaAttention memungkinkan Anda untuk menggabungkan Context Parallelism dengan First Block Cache dan kuantisasi dinamis.

Lihat repositori ParaAttention untuk instruksi dan contoh terperinci tentang cara menskalakan inferensi dengan beberapa GPU.

Jika proses inferensi perlu bersifat persisten dan dapat dilayani, disarankan untuk menggunakan torch.multiprocessing untuk menulis prosesor inferensi Anda sendiri. Hal ini dapat menghilangkan overhead peluncuran proses serta pemuatan dan kompilasi ulang model.

Contoh kode di bawah ini menggabungkan First Block Cache, kuantisasi dinamis fp8, torch.compile, dan Context Parallelism untuk kecepatan inferensi tercepat.

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

Simpan ke run_flux.py dan jalankan dengan torchrun.

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

Kecepatan inferensi berkurang menjadi 8,20 detik dibandingkan baseline, atau 3,21x lebih cepat, dengan 2 GPU NVIDIA L20. Pada 4 L20, kecepatan inferensi adalah 3,90 detik, atau 6,75x lebih cepat.

Contoh kode di bawah ini menggabungkan First Block Cache dan Context Parallelism untuk kecepatan inferensi tercepat.

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

Simpan ke run_hunyuan_video.py dan jalankan dengan torchrun.

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

Kecepatan inferensi berkurang menjadi 649,23 detik dibandingkan dengan baseline, atau 5,66x lebih cepat, dengan 8 GPU NVIDIA L20.

Benchmarks

Jenis GPU	Jumlah GPU	Optimasi	Waktu Nyata (s)	Percepatan
NVIDIA L20	1	Baseline	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

Jenis GPU	Jumlah GPU	Optimasi	Waktu Nyata (s)	Percepatan
NVIDIA L20	1	Baseline	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