Accélérer l'inférence

Les modèles de diffusion sont lents à l’inférence car la génération est un processus itératif au cours duquel le bruit est progressivement affiné en une image ou une vidéo sur un certain nombre d’« étapes ». Pour accélérer ce processus, vous pouvez expérimenter avec différents schedulers, réduire la précision des poids du modèle pour des calculs plus rapides, utiliser des mécanismes d’attention plus économes en mémoire, et plus encore.

Combinez et utilisez ces techniques ensemble pour accélérer l’inférence par rapport à l’utilisation d’une seule technique à la fois.

Ce guide détaille comment accélérer l'inférence.

Type de données du modèle

La précision et le type de données des poids du modèle influencent la vitesse d’inférence, car une précision plus élevée nécessite plus de mémoire pour le chargement et plus de temps pour effectuer les calculs. PyTorch charge les poids du modèle en float32 ou en précision complète par défaut, ce qui en fait un moyen simple d’accélérer rapidement l’inférence en changeant le type de données.

bfloat16

float16

TensorFloat-32

Attention par produit scalaire

L'attention à mémoire efficace optimise la vitesse d'inférence et l'utilisation de la mémoire !

L'attention par produit scalaire mis à l'échelle (SDPA) implémente plusieurs backends d'attention, FlashAttention, xFormers et une implémentation C++ native. Il sélectionne automatiquement le backend le plus optimal pour votre matériel.

SDPA est activé par défaut si vous utilisez PyTorch >= 2.0 et aucune modification supplémentaire de votre code n’est requise. Vous pouvez toutefois expérimenter avec d’autres backends d’attention si vous souhaitez en choisir un vous-même. L’exemple ci-dessous utilise le gestionnaire de contexte torch.nn.attention.sdpa_kernel pour activer une attention efficace.

from torch.nn.attention import SDPBackend, sdpa_kernel
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
  image = pipeline(prompt, num_inference_steps=30).images[0]

torch.compile

torch.compile accélère l’inférence en compilant le code et les opérations PyTorch en noyaux optimisés. Diffusers compile généralement les modèles les plus gourmands en calcul, comme le UNet, le transformer ou le VAE.

Activez les paramètres de compilateur suivants pour une vitesse maximale (reportez-vous à la liste complète pour plus d'options).

import torch
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

Chargez et compilez le UNet et le VAE. Plusieurs modes sont disponibles, mais "max-autotune" optimise la vitesse maximale en compilant vers un graphe CUDA. Les graphes CUDA réduisent efficacement le surcoût en lançant plusieurs opérations GPU via une seule opération CPU.

Avec PyTorch 2.3.1, vous pouvez contrôler le comportement de mise en cache de torch.compile. Cela est particulièrement bénéfique pour les modes de compilation tels que "max-autotune", qui effectue une recherche sur grille sur plusieurs drapeaux de compilation afin de trouver la configuration optimale. En savoir plus dans le tutoriel Compile Time Caching in torch.compile.

Changer la disposition de la mémoire en channels_last optimise également la mémoire et la vitesse d'inférence.

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.decode = torch.compile(
    pipeline.vae.decode,
    mode="max-autotune",
    fullgraph=True
)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

La compilation est lente la première fois, mais une fois compilé, le pipeline est nettement plus rapide. Essayez d'utiliser le pipeline compilé uniquement pour le même type d'opérations d'inférence. Appeler le pipeline compilé avec une taille d'image différente relance la compilation, ce qui est lent et inefficace.

Compilation de formes dynamiques

Assurez-vous d'utiliser toujours la version nocturne de PyTorch pour un meilleur support.

torch.compile suit les formes des entrées et les conditions, et si celles-ci diffèrent, il recompile le modèle. Par exemple, si un modèle est compilé sur une image de résolution 1024x1024 puis utilisé sur une image de résolution différente, cela déclenche une recompilation.

Pour éviter la recompilation, ajoutez dynamic=True afin de générer un noyau plus dynamique et d'éviter ainsi la recompilation lorsque les conditions changent.

+ torch.fx.experimental._config.use_duck_shape = False
+ pipeline.unet = torch.compile(
    pipeline.unet, fullgraph=True, dynamic=True
)

Spécifier use_duck_shape=False indique au compilateur s'il doit utiliser la même variable symbolique pour représenter les tailles d'entrée identiques. Pour plus de détails, consultez ce commentaire.

Tous les modèles ne bénéficient pas nécessairement de la compilation dynamique dès l’origine et peuvent nécessiter des modifications. Consultez cette PR qui a amélioré l’implémentation de AuraFlowPipeline afin de tirer parti de la compilation dynamique.

N’hésitez pas à ouvrir un ticket si la compilation dynamique ne fonctionne pas comme prévu pour un modèle Diffusers.

Compilation régionale

La compilation régionale réduit la latence de démarrage à froid en ne compilant que les blocs petits et fréquemment répétés d'un modèle - typiquement une couche de transformer - et permet de réutiliser les artefacts compilés pour chaque occurrence ultérieure. Pour de nombreuses architectures de diffusion, cela offre les mêmes gains de vitesse d'exécution que la compilation du graphe complet et réduit le temps de compilation de 8 à 10 fois.

Utilisez la méthode compile_repeated_blocks(), un utilitaire qui encapsule torch.compile, sur n'importe quel composant, tel que le modèle transformer, comme indiqué ci-dessous.

# pip install -U diffusers
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"

# compile only the repeated transformer layers inside the UNet
pipeline.unet.compile_repeated_blocks(fullgraph=True)

Pour activer la compilation régionale pour un nouveau modèle, ajoutez un attribut _repeated_blocks à une classe de modèle contenant les noms de classes (sous forme de chaînes de caractères) des blocs que vous souhaitez compiler.

class MyUNet(ModelMixin):
    _repeated_blocks = ("Transformer2DModel",)  # ← compiled by default

Pour plus d’exemples de compilation régionale, consultez la PR de référence.

Il existe également une méthode compile_regions dans Accelerate qui sélectionne automatiquement les blocs candidats d’un modèle à compiler. Le reste du graphe est compilé séparément. C’est utile pour des expériences rapides, car il y a moins d’options pour définir quels blocs compiler ou ajuster les drapeaux de compilation.

# pip install -U accelerate
import torch
from diffusers import StableDiffusionXLPipeline
from accelerate.utils import compile_regions

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)

compile_repeated_blocks() est volontairement explicite. Énumérez les blocs à répéter dans _repeated_blocks et l’assistant ne compile que ces blocs. Il offre un comportement prévisible et un raisonnement simple sur la réutilisation du cache en une seule ligne de code.

Interruptions de graphe

Il est important de spécifier fullgraph=True dans torch.compile pour s'assurer qu'il n'y a pas de coupures de graphe dans le modèle sous-jacent. Cela vous permet de tirer parti de torch.compile sans aucune dégradation des performances. Pour le UNet et le VAE, cela modifie la manière dont vous accédez aux variables de retour.

- latents = unet(
-   latents, timestep=timestep, encoder_hidden_states=prompt_embeds
-).sample

+ latents = unet(
+   latents, timestep=timestep, encoder_hidden_states=prompt_embeds, return_dict=False
+)[0]

Synchronisation GPU

La fonction step() est appelée sur le planificateur à chaque fois que le débruiteur effectue une prédiction, et la variable sigmas est indexée. Lorsqu'elle est placée sur le GPU, elle introduit une latence en raison de la synchronisation de communication entre le CPU et le GPU. Cela devient plus évident lorsque le débruiteur a déjà été compilé.

En général, les sigmas doivent rester sur le CPU pour éviter la synchronisation de communication et la latence.

Consultez le billet de blog torch.compile and Diffusers: A Hands-On Guide to Peak Performance pour maximiser les performances avec torch.compile pour les modèles de diffusion.

Benchmarks

Consultez le jeu de données diffusers/benchmarks pour consulter les données de latence d'inférence et d'utilisation de la mémoire des pipelines compilés.

Le dépôt diffusers-torchao contient également des résultats de benchmark pour les versions compilées de Flux et CogVideoX.

Kernels

Kernels est une bibliothèque pour construire, distribuer et charger des noyaux de calcul optimisés sur le Hub. Elle prend en charge les noyaux attention et les noyaux CUDA personnalisés pour des opérations telles que RMSNorm, GEGLU, RoPE et AdaLN.

Le guide Diffusers Pipeline Integration montre comment intégrer un noyau avec la compétence add cuda-kernels. Cette compétence permet à un agent, comme Claude ou Codex, d'écrire des noyaux personnalisés ciblant un modèle spécifique et votre matériel.

Installez la compétence add cuda-kernels pour apprendre à un agent à écrire un noyau. L'article de blog Custom kernels for all from Codex and Claude aborde ce sujet en détail.

Par exemple, un noyau RMSNorm personnalisé (généré par la compétence add cuda-kernels) avec torch.compile accélère la génération de LTX-Video de 1,43 fois sur une H100.

Quantification dynamique

La quantification dynamique améliore la vitesse d'inférence en réduisant la précision afin de permettre des opérations mathématiques plus rapides. Ce type de quantification détermine comment mettre à l'échelle les activations en fonction des données au moment de l'exécution, plutôt que d'utiliser un facteur d'échelle fixe. En conséquence, le facteur d'échelle est plus précisément aligné sur les données.

L'exemple ci-dessous applique la quantification dynamique int8 au UNet et au VAE avec la bibliothèque torchao.

Consultez notre documentation torchao pour en savoir plus sur l'utilisation de l'intégration torchao de Diffusers.

Configurez les balises du compilateur pour une vitesse maximale.

import torch
from torchao import apply_dynamic_quant
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True

Filtrez certaines couches linéaires du UNet et du VAE qui ne bénéficient pas de la quantification dynamique à l’aide de dynamic_quant_filter_fn.

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

apply_dynamic_quant(pipeline.unet, dynamic_quant_filter_fn)
apply_dynamic_quant(pipeline.vae, dynamic_quant_filter_fn)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

Matrices de projection fusionnées

La méthode fuse_qkv_projections est expérimentale et son support est limité, principalement aux pipelines Stable Diffusion. Consultez ce PR pour en savoir plus sur la manière de l’activer pour d’autres pipelines.

Une entrée est projetée dans trois sous-espaces, représentés par les matrices de projection Q, K et V, au sein d'un bloc d'attention. Ces projections sont généralement calculées séparément, mais vous pouvez les combiner horizontalement en une seule matrice et effectuer la projection en une seule étape. Cela augmente la taille des multiplications de matrices des projections d'entrée et améliore également l'impact de la quantification.

pipeline.fuse_qkv_projections()

Ressources

Lisez le billet de blog Presenting Flux Fast: Making Flux go brrr on H100s pour en savoir plus sur la manière de combiner toutes ces optimisations avec TorchInductor et AOTInductor afin d’obtenir un gain de vitesse d’environ 2,5x en utilisant les recettes de flux-fast.

Ces recettes prennent en charge le matériel AMD et Flux.1 Kontext Dev.
Lisez le billet de blog torch.compile and Diffusers: A Hands-On Guide to Peak Performance pour maximiser les performances lors de l'utilisation de torch.compile.

Update on GitHub