Diffusers documentation
ParaAttention
ParaAttention
Modelos grandes de geração de imagem e vídeo, como FLUX.1-dev e HunyuanVideo, podem representar um desafio de inferência para aplicações em tempo real e implantação devido ao seu tamanho.
ParaAttention é uma biblioteca que implementa paralelismo de contexto e cache do primeiro bloco e pode ser combinada com outras técnicas (torch.compile, quantização dinâmica fp8) para acelerar a inferência.
Este guia mostrará como aplicar o ParaAttention ao FLUX.1-dev e ao HunyuanVideo em GPUs NVIDIA L20. Nenhuma otimização é aplicada ao nosso benchmark de referência, exceto para o HunyuanVideo, a fim de evitar erros de memória insuficiente.
Nosso benchmark de referência mostra que o FLUX.1-dev é capaz de gerar uma imagem com resolução de 1024x1024 em 28 passos em 26,36 segundos, e o HunyuanVideo é capaz de gerar 129 quadros em resolução 720p em 30 passos em 3675,71 segundos.
Para inferência ainda mais rápida com paralelismo de contexto, tente usar GPUs NVIDIA A100 ou H100 (se disponíveis) com suporte a NVLink, especialmente quando há um grande número de GPUs.
Cache do Primeiro Bloco
O cache da saída dos blocos do transformers no modelo e a reutilização deles nas etapas de inferência subsequentes reduzem o custo de computação e tornam a inferência mais rápida.
No entanto, é difícil decidir quando reutilizar o cache para garantir a qualidade das imagens ou vídeos gerados. O ParaAttention usa diretamente a diferença residual da saída do primeiro bloco do transformer para aproximar a diferença entre as saídas do modelo. Quando a diferença é pequena o suficiente, a diferença residual das etapas de inferência anteriores é reutilizada. Em outras palavras, a etapa de desruido é ignorada.
Isso alcança uma aceleração de 2x na inferência do FLUX.1-dev e do HunyuanVideo, com muito boa qualidade.
Para aplicar o cache do primeiro bloco no FLUX.1-dev, chame apply_cache_on_pipe conforme mostrado abaixo. 0,08 é o valor padrão de diferença residual para os modelos 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")| Otimizações | Original | FBCache rdt=0.06 | FBCache rdt=0.08 | FBCache rdt=0.10 | FBCache rdt=0.12 |
|---|---|---|---|---|---|
| Pré-visualização |  |  |  |  |  |
| Tempo Real (s) | 26.36 | 21.83 | 17.01 | 16.00 | 13.78 |
O First Block Cache reduziu a velocidade de inferência para 17,01 segundos em comparação com a linha de base, ou seja, 1,55x mais rápido, mantendo uma perda de qualidade quase nula.
Para aplicar o First Block Cache no HunyuanVideo, use apply_cache_on_pipe conforme mostrado abaixo. 0,06 é o valor padrão de diferença residual para os modelos 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)O First Block Cache reduziu a velocidade de inferência para 2271,06 segundos em comparação com a linha de base, ou seja, 1,62x mais rápido, mantendo uma perda de qualidade quase nula.
quantização fp8
O fp8 com quantização dinâmica acelera ainda mais a inferência e reduz o uso de memória. Tanto as ativações quanto os pesos devem ser quantizados para usar os NVIDIA Tensor Cores de 8 bits.
Use float8_weight_only e float8_dynamic_activation_float8_weight para quantizar o codificador de texto e o modelo transformer.
O método de quantização padrão é a quantização por tensor, mas se sua GPU suportar quantização por linha, você também pode experimentá-la para obter melhor precisão.
Instale o torchao com o comando abaixo.
pip3 install -U torch torchao
torch.compile com mode="max-autotune-no-cudagraphs" ou mode="max-autotune" seleciona o melhor kernel para desempenho. A compilação pode levar muito tempo se for a primeira vez que o modelo é chamado, mas vale a pena uma vez que o modelo tenha sido compilado.
Este exemplo apenas quantiza o modelo transformer, mas você também pode quantizar o codificador de texto para reduzir ainda mais o uso de memória.
A quantização dinâmica pode alterar significativamente a distribuição da saída do modelo, portanto, você precisa alterar o
residual_diff_thresholdpara um valor maior para que ela tenha efeito.
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")A quantização dinâmica fp8 e o torch.compile reduziram o tempo de inferência para 7,56 segundos em comparação com a linha de base, ou seja, 3,48x mais rápido.
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)A GPU NVIDIA L20 possui apenas 48GB de memória e pode enfrentar erros de falta de memória (OOM) após a compilação, caso enable_model_cpu_offload não seja chamado, pois o HunyuanVideo possui tensores de ativação muito grandes ao ser executado com alta resolução e um grande número de quadros. Para GPUs com menos de 80GB de memória, você pode tentar reduzir a resolução e o número de quadros para evitar erros de OOM.
Modelos grandes de geração de vídeo geralmente são limitados pelos cálculos de atenção, em vez das camadas totalmente conectadas. Esses modelos não se beneficiam significativamente da quantização e do torch.compile.
Paralelismo de Contexto
O Context Parallelism paraleliza a inferência e escala com múltiplas GPUs. O design composicional do ParaAttention permite combinar o Context Parallelism com o First Block Cache e a quantização dinâmica.
Consulte o repositório ParaAttention para obter instruções detalhadas e exemplos de como escalar a inferência com múltiplas GPUs.
Se o processo de inferência precisar ser persistente e acionável, é sugerido usar torch.multiprocessing para escrever seu próprio processador de inferência. Isso pode eliminar o overhead de iniciar o processo e de carregar e recompilar o modelo.
O exemplo de código abaixo combina First Block Cache, quantização dinâmica fp8, torch.compile e Context Parallelism para a velocidade de inferência mais rápida.
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()Salve em run_flux.py e execute com torchrun.
# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=2 run_flux.pyA velocidade de inferência é reduzida para 8,20 segundos em comparação com a linha de base, ou seja, 3,21x mais rápida, com 2 GPUs NVIDIA L20. Em 4 L20s, a velocidade de inferência é de 3,90 segundos, ou seja, 6,75x mais rápida.
O exemplo de código abaixo combina First Block Cache e Context Parallelism para a velocidade de inferência mais rápida.
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()Salve em run_hunyuan_video.py e execute com torchrun.
# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=8 run_hunyuan_video.pyA velocidade de inferência é reduzida para 649,23 segundos em comparação com a linha de base, ou seja, 5,66x mais rápida, com 8 GPUs NVIDIA L20.
Benchmarks
| Tipo de GPU | Número de GPUs | Otimizações | Tempo Real (s) | Aceleração |
|---|---|---|---|---|
| NVIDIA L20 | 1 | Linha de base | 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 |
| Tipo de GPU | Número de GPUs | Otimizações | Tempo Real (s) | Aceleração |
|---|---|---|---|---|
| NVIDIA L20 | 1 | Linha de base | 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 |