diffusers/main/pt/optimization/para_attn.md

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

Como o AdaCache funciona, o First Block Cache é uma variante dele

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)

Seu navegador não suporta a tag de vídeo.

HunyuanVideo sem FBCache

Seu navegador não suporta a tag de vídeo.

HunyuanVideo com FBCache

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_threshold para 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.py

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

A 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