diffusers/main/ko/optimization/para_attn.md

ParaAttention

Large image and video generation models, such as FLUX.1-dev and HunyuanVideo, can be an inference challenge for real-time applications and deployment because of their size.

ParaAttention는 context parallelism과 first block cache를 구현한 라이브러리로, 다른 기법(torch.compile, fp8 dynamic quantization)과 결합하여 추론을 가속화할 수 있습니다.

이 가이드에서는 NVIDIA L20 GPU에서 FLUX.1-dev와 HunyuanVideo에 ParaAttention을 적용하는 방법을 보여줍니다. 베이스라인 벤치마크에는 최적화가 적용되지 않았으며, HunyuanVideo의 경우 메모리 부족 오류를 방지하기 위해 예외적으로 적용했습니다.

Our baseline benchmark에 따르면, 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배의 속도 향상을 달성합니다.

AdaCache의 작동 방식, First Block Cache는 그 변형입니다

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

최적화	원본	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배 더 빠르게 만들었으며, 거의 품질 손실이 없도록 유지했습니다.

HunyuanVideo에 First Block Cache를 적용하려면, 아래에 표시된 대로 apply_cache_on_pipe를 사용하세요. 0.06은 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)

브라우저가 비디오 태그를 지원하지 않습니다.

FBCache가 없는 HunyuanVideo

브라우저가 비디오 태그를 지원하지 않습니다.

FBCache를 사용한 HunyuanVideo

First Block Cache는 베이스라인 대비 추론 속도를 2271.06초로 줄여 1.62배 더 빠르게 만들었으며, 거의 품질 손실이 없는 상태를 유지했습니다.

fp8 양자화

동적 양자화를 적용한 fp8은 추론 속도를 더 빠르게 하고 메모리 사용량을 줄입니다. 8비트 NVIDIA Tensor Cores를 사용하려면 활성화 값과 가중치 모두를 양자화해야 합니다.

float8_weight_only와 float8_dynamic_activation_float8_weight를 사용하여 텍스트 인코더와 트랜스포머 모델을 양자화합니다.

기본 양자화 방법은 텐서별(per tensor) 양자화이지만, GPU가 행별(row-wise) 양자화를 지원한다면 더 높은 정확도를 위해 이 방법도 시도해 볼 수 있습니다.

아래 명령어로 torchao를 설치합니다.

pip3 install -U torch torchao

mode="max-autotune-no-cudagraphs" 또는 mode="max-autotune"를 사용한 torch.compile은 성능을 위해 최적의 커널을 선택합니다. 모델이 처음 호출되는 경우 컴파일 시간이 오래 걸릴 수 있지만, 모델이 컴파일된 후에는 그 가치가 있습니다.

이 예시에서는 transformer 모델만 양자화하지만, 메모리 사용량을 더 줄이기 위해 텍스트 인코더도 양자화할 수 있습니다.

Dynamic quantization은 모델 출력의 분포를 크게 변경할 수 있으므로, 효과를 보려면 residual_diff_threshold를 더 큰 값으로 변경해야 합니다.

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배 빨라집니다.

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를 호출하지 않으면 OOM(Out-Of-Memory) 오류가 발생할 수 있습니다. 이는 HunyuanVideo가 고해상도와 많은 프레임 수로 실행될 때 매우 큰 활성화 텐서를 생성하기 때문입니다. 80GB 미만의 메모리를 가진 GPU에서는 OOM 오류를 피하기 위해 해상도와 프레임 수를 줄여 보실 수 있습니다.

대형 비디오 생성 모델은 일반적으로 완전 연결 계층보다 어텐션 계산에 병목 현상이 발생합니다. 이러한 모델은 양자화 및 torch.compile에서 상당한 이점을 얻지 못합니다.

컨텍스트 병렬

Context Parallelism은 추론을 병렬화하며 여러 GPU로 확장됩니다. ParaAttention의 구성式设计은 Context Parallelism을 First Block Cache 및 동적 양자화와 결합할 수 있도록 합니다.

여러 GPU로 추론을 확장하는 방법에 대한 자세한 지침과 예제는 ParaAttention 저장소를 참조하세요.

추론 프로세스가 영속적이고 서비스 가능해야 하는 경우, torch.multiprocessing를 사용하여 자체 추론 프로세서를 작성하는 것이 좋습니다. 이를 통해 프로세스 시작 오버헤드와 모델 로딩 및 재컴파일 비용을 제거할 수 있습니다.

아래 코드 샘플은 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배 빨라집니다.

벤치마크

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