ParaAttention

Les grands modèles de génération d’images et de vidéos, tels que FLUX.1-dev et HunyuanVideo, peuvent poser un défi d’inférence pour les applications en temps réel et le déploiement en raison de leur taille.

ParaAttention est une bibliothèque qui implémente le context parallelism et le first block cache, et peut être combinée avec d'autres techniques (torch.compile, quantification dynamique fp8) pour accélérer l'inférence.

Ce guide vous montre comment appliquer ParaAttention à FLUX.1-dev et HunyuanVideo sur des GPU NVIDIA L20. Aucune optimisation n’est appliquée pour notre benchmark de référence, à l’exception de HunyuanVideo afin d’éviter les erreurs de mémoire insuffisante.

Notre benchmark de référence montre que FLUX.1-dev est capable de générer une image en résolution 1024x1024 en 28 étapes en 26,36 secondes, et que HunyuanVideo est capable de générer 129 images en résolution 720p en 30 étapes en 3675,71 secondes.

Pour une inférence encore plus rapide avec le parallélisme de contexte, essayez d’utiliser des GPU NVIDIA A100 ou H100 (si disponibles) avec prise en charge de NVLink, en particulier lorsqu’il y a un grand nombre de GPU.

Cache du premier bloc

Mettre en cache la sortie des blocs de transformers du modèle et la réutiliser aux étapes d'inférence suivantes réduit le coût de calcul et accélère l'inférence.

Cependant, il est difficile de déterminer quand réutiliser le cache afin de garantir la qualité des images ou vidéos générées. ParaAttention utilise directement la différence résiduelle de la sortie du premier bloc transformer pour approximer l'écart entre les sorties du modèle. Lorsque cet écart est suffisamment faible, la différence résiduelle des étapes d'inférence précédentes est réutilisée. Autrement dit, l'étape de débruitage est ignorée.

Cela permet d'obtenir un gain de vitesse de 2x lors de l'inférence de FLUX.1-dev et HunyuanVideo, avec une très bonne qualité.

Cache in Diffusion Transformer — Fonctionnement d'AdaCache, First Block Cache en est une variante

FLUX-1.dev

HunyuanVideo

Pour appliquer le cache du premier bloc sur FLUX.1-dev, appelez apply_cache_on_pipe comme indiqué ci-dessous. 0,08 est la valeur par défaut de la différence résiduelle pour les modèles 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")

Optimisations	Original	FBCache rdt=0.06	FBCache rdt=0.08	FBCache rdt=0.10	FBCache rdt=0.12
Aperçu
Temps réel (s)	26.36	21.83	17.01	16.00	13.78

First Block Cache a réduit la vitesse d’inférence à 17,01 secondes par rapport à la référence, soit 1,55 fois plus rapide, tout en maintenant une perte de qualité quasi nulle.

quantification fp8

fp8 avec quantification dynamique accélère davantage l’inférence et réduit l’utilisation de la mémoire. Les activations et les poids doivent tous deux être quantifiés afin d’utiliser les NVIDIA Tensor Cores 8 bits.

Utilisez float8_weight_only et float8_dynamic_activation_float8_weight pour quantifier le codeur de texte et le modèle transformer.

La méthode de quantification par défaut est la quantification par tenseur, mais si votre GPU prend en charge la quantification par ligne, vous pouvez également l'essayer pour une meilleure précision.

Installez torchao avec la commande ci-dessous.

pip3 install -U torch torchao

torch.compile avec mode="max-autotune-no-cudagraphs" ou mode="max-autotune" sélectionne le meilleur noyau pour les performances. La compilation peut prendre beaucoup de temps s’il s’agit de la première fois que le modèle est appelé, mais cela en vaut la peine une fois le modèle compilé.

Cet exemple ne quantise que le modèle transformer, mais vous pouvez également quantiser l’encodeur de texte pour réduire encore davantage l’utilisation de la mémoire.

La quantification dynamique peut modifier considérablement la distribution de la sortie du modèle, il est donc nécessaire d’augmenter la valeur de residual_diff_threshold pour que cela prenne effet.

FLUX-1.dev

HunyuanVideo

Parallélisme de contexte

Le parallélisme de contexte (Context Parallelism) parallélise l'inférence et s'échelonne sur plusieurs GPU. La conception compositionnelle de ParaAttention vous permet de combiner le parallélisme de contexte avec le cache du premier bloc et la quantification dynamique.

Reportez-vous au dépôt ParaAttention pour des instructions détaillées et des exemples sur la manière de mettre à l’échelle l’inférence avec plusieurs GPU.

Si le processus d’inférence doit être persistant et accessible en tant que service, il est recommandé d’utiliser torch.multiprocessing pour écrire votre propre processeur d’inférence. Cela permet d’éliminer le surcoût lié au lancement du processus ainsi qu’au chargement et à la recompilation du modèle.

FLUX-1.dev

HunyuanVideo

L'extrait de code ci-dessous combine First Block Cache, quantification dynamique fp8, torch.compile et Context Parallelism pour la vitesse d'inférence la plus rapide.

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

Enregistrez dans run_flux.py et lancez-le avec torchrun.

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

La vitesse d’inférence est réduite à 8,20 secondes par rapport à la référence, soit 3,21 fois plus rapide, avec 2 GPU NVIDIA L20. Sur 4 L20, la vitesse d’inférence est de 3,90 secondes, soit 6,75 fois plus rapide.

L'extrait de code ci-dessous combine First Block Cache et Context Parallelism pour obtenir la vitesse d'inférence la plus rapide.

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

Enregistrez dans run_hunyuan_video.py et lancez-le avec torchrun.

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

La vitesse d'inférence est réduite à 649,23 secondes par rapport à la référence, soit 5,66 fois plus rapide, avec 8 GPU NVIDIA L20.

Benchmarks

FLUX-1.dev

HunyuanVideo

Type de GPU	Nombre de GPU	Optimisations	Temps réel (s)	Accélération
NVIDIA L20	1	Ligne 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

Type de GPU	Nombre de GPU	Optimisations	Temps réel (s)	Accélération
NVIDIA L20	1	Ligne 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

Update on GitHub

Diffusers

ParaAttention

Cache du premier bloc

quantification fp8

Parallélisme de contexte

Benchmarks