diffusers/main/es/optimization/para_attn.md

ParaAttention

Los modelos de generación de imágenes y videos de gran tamaño, como FLUX.1-dev y HunyuanVideo, pueden suponer un reto de inferencia para aplicaciones en tiempo real y despliegue debido a su tamaño.

ParaAttention es una biblioteca que implementa paralelismo de contexto y cache del primer bloque, y puede combinarse con otras técnicas (torch.compile, cuantización dinámica fp8) para acelerar la inferencia.

Esta guía le mostrará cómo aplicar ParaAttention a FLUX.1-dev y HunyuanVideo en GPUs NVIDIA L20. No se aplican optimizaciones para nuestra línea base de referencia, excepto para HunyuanVideo, para evitar errores de memoria insuficiente.

Nuestro benchmark de referencia muestra que FLUX.1-dev es capaz de generar una imagen de resolución 1024x1024 en 28 pasos en 26.36 segundos, y HunyuanVideo es capaz de generar 129 fotogramas a resolución 720p en 30 pasos en 3675.71 segundos.

Para una inferencia aún más rápida con paralelismo de contexto, intente usar GPUs NVIDIA A100 o H100 (si están disponibles) con soporte para NVLink, especialmente cuando hay un gran número de GPUs.

Caché del primer bloque

Almacenar en caché la salida de los bloques de transformers del modelo y reutilizarla en los pasos de inferencia posteriores reduce el costo de cómputo y acelera la inferencia.

Sin embargo, es difícil decidir cuándo reutilizar la caché para garantizar la calidad de las imágenes o videos generados. ParaAttention utiliza directamente la diferencia residual de la salida del primer bloque de transformer para aproximar la diferencia entre las salidas del modelo. Cuando la diferencia es lo suficientemente pequeña, se reutiliza la diferencia residual de los pasos de inferencia anteriores. En otras palabras, se omite el paso de desruido.

Esto logra una aceleración de 2x en la inferencia de FLUX.1-dev y HunyuanVideo con una calidad muy buena.

Cómo funciona AdaCache, First Block Cache es una variante de este

Para aplicar la caché del primer bloque en FLUX.1-dev, llame a apply_cache_on_pipe como se muestra a continuación. 0.08 es el valor de diferencia residual predeterminado para los 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")

Optimizaciones	Original	FBCache rdt=0.06	FBCache rdt=0.08	FBCache rdt=0.10	FBCache rdt=0.12
Vista previa
Tiempo real (s)	26.36	21.83	17.01	16.00	13.78

First Block Cache redujo la velocidad de inferencia a 17,01 segundos en comparación con la línea base, es decir, 1,55 veces más rápido, mientras mantenía una pérdida de calidad casi nula.

Para aplicar First Block Cache en HunyuanVideo, use apply_cache_on_pipe como se muestra a continuación. 0.06 es el valor de diferencia residual predeterminado para los modelos de 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)

Tu navegador no admite la etiqueta de video.

HunyuanVideo sin FBCache

Tu navegador no admite la etiqueta de video.

HunyuanVideo con FBCache

First Block Cache redujo la velocidad de inferencia a 2271.06 segundos en comparación con la línea base, es decir, 1.62 veces más rápido, mientras mantenía una pérdida de calidad casi nula.

cuantización fp8

fp8 con cuantización dinámica acelera aún más la inferencia y reduce el uso de memoria. Tanto las activaciones como los pesos deben estar cuantizados para poder utilizar los NVIDIA Tensor Cores de 8 bits.

Use float8_weight_only y float8_dynamic_activation_float8_weight para cuantizar el codificador de texto y el modelo transformer.

El método de cuantización predeterminado es la cuantización por tensor, pero si tu GPU admite la cuantización por fila, también puedes probarla para obtener una mayor precisión.

Instala torchao con el siguiente comando.

pip3 install -U torch torchao

torch.compile con mode="max-autotune-no-cudagraphs" o mode="max-autotune" selecciona el mejor kernel para el rendimiento. La compilación puede tardar mucho si es la primera vez que se llama al modelo, pero vale la pena una vez que el modelo ha sido compilado.

Este ejemplo solo cuantiza el modelo transformer, pero también puedes cuantizar el codificador de texto para reducir aún más el uso de memoria.

La cuantización dinámica puede cambiar significativamente la distribución de la salida del modelo, por lo que es necesario cambiar residual_diff_threshold a un valor mayor para que surta efecto.

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

La cuantización dinámica fp8 y torch.compile redujeron la velocidad de inferencia a 7,56 segundos en comparación con la línea base, es decir, 3,48 veces más 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)

Una GPU NVIDIA L20 solo tiene 48 GB de memoria y podría presentar errores de memoria insuficiente (OOM) después de la compilación si no se llama a enable_model_cpu_offload, ya que HunyuanVideo tiene tensores de activación muy grandes al ejecutarse con alta resolución y un gran número de fotogramas. Para GPUs con menos de 80 GB de memoria, puede intentar reducir la resolución y el número de fotogramas para evitar errores de OOM.

Los modelos de generación de video a gran escala suelen estar limitados por los cálculos de atención en lugar de por las capas totalmente conectadas. Estos modelos no se benefician significativamente de la cuantización ni de torch.compile.

Paralelismo de contexto

El paralelismo de contexto (Context Parallelism) paraleliza la inferencia y escala con múltiples GPUs. El diseño composicional de ParaAttention le permite combinar el paralelismo de contexto con la caché del primer bloque (First Block Cache) y la cuantización dinámica.

Consulte el repositorio ParaAttention para obtener instrucciones detalladas y ejemplos sobre cómo escalar la inferencia con varias GPU.

Si el proceso de inferencia debe ser persistente y servible, se sugiere usar torch.multiprocessing para escribir su propio procesador de inferencia. Esto puede eliminar la sobrecarga de iniciar el proceso y cargar y recompilar el modelo.

La muestra de código siguiente combina First Block Cache, cuantización dinámica fp8, torch.compile y Context Parallelism para la velocidad de inferencia más 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()

Guárdalo en run_flux.py y ejecútalo con torchrun.

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

La velocidad de inferencia se reduce a 8,20 segundos en comparación con la línea base, es decir, es 3,21 veces más rápida, con 2 GPU NVIDIA L20. En 4 L20, la velocidad de inferencia es de 3,90 segundos, o 6,75 veces más rápida.

La muestra de código siguiente combina First Block Cache y Context Parallelism para la velocidad de inferencia más 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()

Guárdalo en run_hunyuan_video.py y ejecútalo con torchrun.

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

La velocidad de inferencia se reduce a 649,23 segundos en comparación con la línea base, o sea, 5,66 veces más rápido, con 8 GPU NVIDIA L20.

Benchmarks

Tipo de GPU	Número de GPUs	Optimizaciones	Tiempo real (s)	Aceleración
NVIDIA L20	1	Línea 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	Optimizaciones	Tiempo real (s)	Aceleración
NVIDIA L20	1	Línea 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