Précision et compilation

La basse précision et la compilation sont deux façons d’accélérer l’inférence de Diffusers. Chargez les poids en bfloat16 ou float16, puis compilez le débruiteur (le UNet ou le transformer) avec torch.compile ou une compilation régionale.

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. Diffusers charge les poids du modèle en float32 lorsque vous omettez dtype, ce qui fait du changement de type de données une méthode simple pour accélérer rapidement l’inférence.

bfloat16

float16

TensorFloat-32

Attention par produit scalaire

L’attention par produit scalaire mis à l’échelle (SDPA) est l’attention par défaut sur PyTorch 2.0 et versions ultérieures, via AttnProcessor2_0 ou le backend native du dispatcheur d’attention. PyTorch sélectionne un noyau SDPA adapté à votre matériel. Pour FlashAttention, SageAttention, xFormers, les noyaux Hub et d’autres backends, consultez Backends d’attention.

torch.compile

torch.compile accélère l’inférence en compilant le code et les opérations PyTorch en noyaux optimisés. Vous compilez généralement le débruiteur, qui domine le temps d’exécution, à savoir pipeline.unet ou pipeline.transformer, et parfois le VAE également.

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 débruiteur et le VAE. Utilisez pipeline.unet sur les pipelines UNet tels que Stable Diffusion XL, ou pipeline.transformer sur les pipelines de type Flux et DiT. Plusieurs modes différents sont disponibles, mais "max-autotune" recherche les noyaux les plus rapides et utilise les graphes CUDA. Les graphes CUDA réduisent la surcharge en lançant plusieurs opérations GPU via une seule opération CPU.

Changer la disposition de la mémoire en channels_last peut également accélérer les modèles intensifs en convolution, comme les UNets et les VAEs. Effectuez d’abord un benchmark, car certains modèles s’exécutent plus lentement.

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.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.to(memory_format=torch.channels_last)
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 lors de la première exécution, en particulier avec "max-autotune", car le compilateur recherche les noyaux et réutilise ensuite le même objet de pipeline compilé (voir Compile Time Caching et Caching Configuration pour savoir comment configurer le comportement de mise en cache). Appeler le pipeline compilé avec une taille d'image différente déclenche une recompilation.

Compilation de formes dynamiques

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 - généralement 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 des accélérations d'exécution comparables à la compilation du graphe complet et peut réduire le temps de compilation jusqu'à 8 à 10 fois.

Utilisez la méthode compile_repeated_blocks(), un utilitaire qui encapsule torch.compile, sur le débruiteur. Appelez-la sur pipeline.unet ou pipeline.transformer, selon le pipeline.

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"
pipeline.unet.compile_repeated_blocks(fullgraph=True)

Pour activer la compilation régionale pour un nouveau modèle, définissez _repeated_blocks avec les noms de classes des blocs répétés (chaînes de caractères). Pour l’UNet de SDXL, utilisez BasicTransformerBlock. Les débruiteurs Transformer listent de la même manière les noms de classes de leurs blocs répétés.

_repeated_blocks = ["BasicTransformerBlock"]

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

Définissez fullgraph=True afin que torch.compile lève une erreur en cas de rupture de graphe au lieu de scinder silencieusement le graphe, ce qui réduit l’accélération.

Synchronisation GPU

Après chaque prédiction du débruiteur, le pipeline appelle la méthode step du planificateur, qui indexe sigmas. Laisser sigmas sur le GPU peut forcer une synchronisation CPU↔GPU à chaque étape. Ce coût est facile à manquer tant que le débruiteur n’est pas compilé. Il est préférable de laisser sigmas sur le CPU (les planificateurs Diffusers tels que Euler le font déjà).

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. L'article Custom kernels for all from Codex and Claude fournit plus de détails.

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 des activations en int8 est disponible dans torchao. Utilisez TorchAoConfig avec une configuration telle que Int8DynamicActivationInt8WeightConfig, ou l’API quantize_ de torchao. N’utilisez pas l’ancien utilitaire apply_dynamic_quant.

Matrices de projection fusionnées

~DiffusionPipeline.fuse_qkv_projections est expérimental.

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

Étapes suivantes

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

Update on GitHub