xDiT est une bibliothèque Open Source qui accélère l'inférence pour les modèles Diffusion Transformer (DiT) à l'aide de techniques de parallélisme et d'optimisation. Ces techniques permettent une configuration multi-GPU évolutive pour les charges de travail exigeantes. Cette page explique comment déployer des modèles DiT à l'aide de xDiT et de GPU Cloud sur Vertex AI.
Pour en savoir plus sur xDiT, consultez le projet GitHub xDiT.
Avantages
La liste suivante décrit les principaux avantages de l'utilisation de xDiT pour diffuser des modèles DiT sur Vertex AI :
- Génération jusqu'à trois fois plus rapide : générez des images et des vidéos haute résolution en un temps record par rapport aux autres solutions de diffusion.
- Prise en charge évolutive de plusieurs GPU : distribuez efficacement les charges de travail sur plusieurs GPU pour des performances optimales.
- Parallélisme hybride : xDiT est compatible avec différentes approches de traitement parallèle, telles que le parallélisme de séquence unifié, PipeFusion, le parallélisme CFG et le parallélisme de données. Ces méthodes peuvent être combinées dans une recette unique pour optimiser les performances.
- Performances optimisées avec un seul GPU : xDiT offre une inférence plus rapide, même avec un seul GPU.
- Accélération du GPU : xDiT intègre plusieurs méthodes d'accélération du noyau et utilise des techniques de DiTFastAttn pour accélérer l'inférence sur un seul GPU.
- Déploiement facile : démarrez rapidement grâce au déploiement en un clic ou aux notebooks Colab Enterprise dans Vertex AI Model Garden.
Modèles compatibles
xDiT est disponible pour certaines architectures de modèle DiT dans Vertex AI Model Garden, telles que les variantes de modèle Flux.1 Schnell, CogVideoX-2b et Wan2.1 de conversion de texte en vidéo. Pour savoir si un modèle DiT est compatible avec xDiT dans Model Garden, consultez sa fiche de modèle dans Model Garden.
Parallélisme hybride pour les performances multi-GPU :
xDiT utilise une combinaison de techniques de parallélisme pour maximiser les performances sur les configurations multi-GPU. Ces techniques fonctionnent conjointement pour répartir la charge de travail et optimiser l'utilisation des ressources :
- Parallélisme de séquence unifié : cette technique divise les données d'entrée (par exemple, en divisant une image en blocs) sur plusieurs GPU, ce qui réduit l'utilisation de la mémoire et améliore l'évolutivité.
- PipeFusion : PipeFusion divise le modèle DiT en étapes et attribue chaque étape à un GPU différent, ce qui permet le traitement parallèle de différentes parties du modèle.
- Parallélisme CFG : cette technique optimise spécifiquement les modèles en utilisant le guidage sans classificateur, une méthode courante pour contrôler le style et le contenu des images générées. Il parallélise le calcul des branches conditionnelles et inconditionnelles, ce qui accélère l'inférence.
- Parallélisme des données : cette méthode réplique l'intégralité du modèle sur chaque GPU, chaque GPU traitant un lot différent de données d'entrée, ce qui augmente le débit global du système.
Pour en savoir plus sur les améliorations des performances, consultez le rapport de xDiT sur Flux.1 Schnell ou CogVideoX-2b. Google a pu reproduire ces résultats sur Vertex AI Model Garden.
Accélération par un seul GPU
La bibliothèque xDiT offre des avantages pour le service à GPU unique en utilisant torch.compile et onediff pour améliorer la vitesse d'exécution sur les GPU. Ces techniques peuvent également être utilisées en parallèle avec le parallélisme hybride.
xDiT dispose également d'une technique de calcul d'attention efficace, appelée DiTFastAttn, pour résoudre le problème de goulot d'étranglement de calcul de DiT. Pour le moment, cette technique n'est disponible que pour les configurations à un seul GPU ou en combinaison avec le parallélisme des données.
Premiers pas dans Model Garden
Le conteneur de diffusion de GPU Cloud optimisé xDiT est fourni dans Vertex AI Model Garden. Pour les modèles compatibles, les déploiements utilisent ce conteneur lorsque vous utilisez les déploiements en un clic ou les exemples de notebooks Colab Enterprise.
Les exemples suivants utilisent le modèle Flux.1-schnell pour montrer comment déployer un modèle DiT sur un conteneur xDiT.
Utiliser le déploiement en un clic
Vous pouvez déployer un point de terminaison Vertex AI personnalisé avec le conteneur xDiT à l'aide d'une fiche de modèle.
Accédez à la page Fiche de modèle, puis cliquez sur Déployer.
Pour la variante de modèle à utiliser, sélectionnez un type de machine pour votre déploiement.
Cliquez sur Déployer pour lancer le processus de déploiement. Vous recevez deux notifications par e-mail : une lorsque le modèle est importé et une autre lorsque le point de terminaison est prêt.
Utiliser le notebook Colab Enterprise
Pour plus de flexibilité et de personnalisation, utilisez les exemples de notebooks Colab Enterprise pour déployer un point de terminaison Vertex AI avec le conteneur xDiT à l'aide du SDK Vertex AI pour Python.
Accédez à la page Fiche de modèle, puis cliquez sur Ouvrir le notebook.
Sélectionnez le notebook Vertex Serving. Le notebook s'ouvre dans Colab Enterprise.
Parcourez le notebook pour déployer un modèle à l'aide du conteneur xDiT et envoyez des requêtes de prédiction au point de terminaison. L'extrait de code pour le déploiement se présente comme suit :
import vertexai
from vertexai import model_garden
vertexai.init(project=<YOUR_PROJECT_ID>, location=<REGION>)
model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")
endpoint = model.deploy()
Arguments xDiT
xDiT propose une gamme d'arguments de serveur qui peuvent être configurés pour optimiser les performances pour des cas d'utilisation spécifiques. Ces arguments sont définis comme variables d'environnement lors du déploiement. Voici les principaux arguments que vous devrez peut-être configurer :
Configuration du modèle
MODEL_ID(chaîne) : spécifie l'identifiant du modèle à charger. Il doit correspondre au nom du modèle dans votre registre ou chemin d'accès.
Arguments d'optimisation de l'exécution
N_GPUS(entier) : spécifie le nombre de GPU à utiliser pour l'inférence. La valeur par défaut est 1.WARMUP_STEPS(entier) : nombre d'étapes de préchauffage requises avant le début de l'inférence. Cela est particulièrement important lorsque PipeFusion est activé pour garantir des performances stables. La valeur par défaut est de 1.USE_PARALLEL_VAE(booléen) : permet de traiter efficacement les images haute résolution (plus de 2 048 pixels) en parallélisant le composant VAE sur les appareils. Cela permet d'éviter les problèmes de mémoire insuffisante pour les images volumineuses. La valeur par défaut est "false".USE_TORCH_COMPILE(booléen) : active l'accélération par GPU unique via torch.compile, ce qui permet d'optimiser les performances au niveau du noyau. La valeur par défaut est "false".USE_ONEDIFF(booléen) : active la technologie d'accélération de la compilation OneDiff pour optimiser la vitesse d'exécution du kernel GPU. La valeur par défaut est "false".
Arguments de parallélisme des données
DATA_PARALLEL_DEGREE(entier) : définit le degré de parallélisme des données. Laissez ce champ vide pour désactiver le parallélisme ou définissez-le sur le degré de parallélisme sélectionné.USE_CFG_PARALLEL(booléen) : active le calcul parallèle pour le guidage sans classificateur (CFG), également appelé "fractionnement de lot". Lorsque cette option est activée, le degré de parallélisme constant est de 2. Définissez sur "true" lorsque vous utilisez CFG pour contrôler le style et le contenu de la sortie. La valeur par défaut est "false".
Arguments de parallélisme de séquence (USP – Unified Sequence Parallelism)
ULYSSES_DEGREE(entier) : définit le degré Ulysses pour l'approche parallèle de séquence unifiée, qui combine DeepSpeed-Ulysses et Ring-Attention. Cela contrôle le schéma de communication all-to-all. Laissez ce champ vide pour utiliser la valeur par défaut.RING_DEGREE(entier) : définit le degré de l'anneau pour la communication peer-to-peer dans le parallélisme de séquence. Fonctionne conjointement avec ULYSSES_DEGREE pour former le maillage de processus 2D. Laissez ce champ vide pour utiliser la valeur par défaut.
Arguments du parallélisme tensoriel
TENSOR_PARALLEL_DEGREE(entier) : définit le degré de parallélisme de tenseur, qui répartit les paramètres du modèle sur les appareils selon les dimensions des caractéristiques pour réduire les coûts de mémoire par appareil. Laissez ce champ vide pour désactiver la fonctionnalité.SPLIT_SCHEME(chaîne) : définit comment répartir les Tensors du modèle sur les appareils (par exemple, par têtes d'attention ou dimensions cachées). Laissez ce champ vide pour utiliser le schéma de fractionnement par défaut.
Arguments distribués Ray
USE_RAY(booléen) : active le framework d'exécution distribuée Ray pour mettre à l'échelle les calculs sur plusieurs nœuds. La valeur par défaut est "false".RAY_WORLD_SIZE(entier) : nombre total de processus dans le cluster Ray. La valeur par défaut est 1.VAE_PARALLEL_SIZE(entier) : nombre de processus dédiés au traitement parallèle VAE lors de l'utilisation de Ray. La valeur par défaut est 0.DIT_PARALLEL_SIZE(entier) : nombre de processus dédiés au traitement parallèle du backbone DiT lors de l'utilisation de Ray. La valeur par défaut est 0.
Arguments parallèles PipeFusion
PIPEFUSION_PARALLEL_DEGREE(entier) : définit le degré de parallélisme pour PipeFusion, un parallélisme de pipeline au niveau de la séquence qui tire parti des caractéristiques de redondance temporelle d'entrée des modèles de diffusion. Des valeurs plus élevées augmentent le parallélisme, mais nécessitent plus de mémoire. La valeur par défaut est de 1.NUM_PIPELINE_PATCH(entier) : nombre de correctifs dans lesquels la séquence doit être divisée pour le traitement du pipeline. Laissez ce champ vide pour une détermination automatique.ATTN_LAYER_NUM_FOR_PP(chaîne) : spécifie les couches d'attention à utiliser pour le parallélisme de pipeline. Peut être séparé par une virgule (par exemple, "10,9") ou séparées par un espace (par exemple, "10 9"). Laissez ce champ vide pour utiliser tous les calques.
Arguments d'optimisation de la mémoire
ENABLE_MODEL_CPU_OFFLOAD(booléen) : décharge les pondérations du modèle dans la mémoire du processeur lorsqu'elles ne sont pas utilisées, ce qui réduit l'utilisation de la mémoire du GPU au prix d'une latence accrue. La valeur par défaut est "false".ENABLE_SEQUENTIAL_CPU_OFFLOAD(booléen) : décharge séquentiellement les couches de modèle sur le CPU lors de la transmission directe, ce qui permet l'inférence de modèles plus grands que la mémoire du GPU. La valeur par défaut est "false".ENABLE_TILING(booléen) : réduit l'utilisation de la mémoire GPU en décodant le composant VAE une tuile à la fois. Cet argument est utile pour les images ou vidéos volumineuses, et pour éviter les erreurs de mémoire insuffisante. La valeur par défaut est "false".ENABLE_SLICING(booléen) : réduit l'utilisation de la mémoire GPU en divisant le Tensor d'entrée en tranches pour le décodage VAE. La valeur par défaut est "false".
Arguments DiTFastAttn (optimisation de l'attention)
USE_FAST_ATTN(booléen) : active l'accélération DiTFastAttn pour l'inférence à un seul GPU, en utilisant la réduction temporelle des entrées pour réduire la complexité de calcul. La valeur par défaut est "false".N_CALIB(entier) : nombre d'échantillons de calibration pour l'optimisation DiTFastAttn. La valeur par défaut est 8.THRESHOLD(float) : seuil de similarité pour la réduction de la similarité temporelle dans DiTFastAttn. La valeur par défaut est 0,5.WINDOW_SIZE(entier) : taille de la fenêtre pour l'attention de la fenêtre avec mise en cache résiduelle afin de réduire la redondance spatiale. La valeur par défaut est 64.COCO_PATH(chaîne) : chemin d'accès à l'ensemble de données COCO pour la calibration DiTFastAttn. Obligatoire lorsque USE_FAST_ATTN est défini sur "true". Laissez ce champ vide si vous ne l'utilisez pas.
Arguments d'optimisation du cache
USE_CACHE(booléen) : active les mécanismes de mise en cache généraux pour réduire les calculs redondants. La valeur par défaut est "false".USE_TEACACHE(booléen) : active la méthode d'optimisation TeaCache pour la mise en cache des résultats intermédiaires. La valeur par défaut est "false".USE_FBCACHE(booléen) : active la méthode d'optimisation du cache du premier bloc. La valeur par défaut est "false".
Arguments d'optimisation de la précision
USE_FP8_T5_ENCODER(booléen) : active la précision FP8 (virgule flottante de 8 bits) pour l'encodeur de texte T5, ce qui réduit l'utilisation de la mémoire et peut améliorer le débit avec un impact minimal sur la qualité. La valeur par défaut est "false".
Personnalisation de la diffusion
Model Garden fournit des configurations de parallélisation xDiT par défaut pour les modèles compatibles. Vous pouvez inspecter ces paramètres par défaut à l'aide du SDK Vertex AI pour Python.
Pour afficher la configuration de déploiement par défaut d'un modèle, tel que "black-forest-labs/FLUX.1-schnell", vous pouvez exécuter l'extrait de code suivant :
import vertexai
from vertexai import model_garden
vertexai.init(project=<YOUR_PROJECT_ID>, location=<REGION>)
model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")
deploy_options = model.list_deploy_options()
# Example Response
# ['black-forest-labs/[email protected]']
# [model_display_name: "Flux1-schnell"
# container_spec {
# image_uri: "us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/xdit-serve.cu125.0-2.ubuntu2204.py310"
# env {
# name: "DEPLOY_SOURCE"
# value: "UI_NATIVE_MODEL"
# }
# env {
# name: "MODEL_ID"
# value: "gs://vertex-model-garden-restricted-us/black-forest-labs/FLUX.1-schnell"
# }
# env {
# name: "TASK"
# value: "text-to-image"
# }
# env {
# name: "N_GPUS"
# value: "2"
# }
# env {
# name: "USE_TORCH_COMPILE"
# value: "true"
# }
# env {
# name: "RING_DEGREE"
# value: "2"
# }
# ..........]
La méthode list_deploy_options() renvoie les spécifications du conteneur, y compris les variables d'environnement (env) qui définissent la configuration xDiT.
Pour personnaliser la stratégie de parallélisme, vous pouvez remplacer ces variables d'environnement lors du déploiement du modèle. L'exemple suivant montre comment modifier RING_DEGREE et ULYSSES_DEGREE pour une configuration à deux GPU, en modifiant l'approche du parallélisme :
import vertexai
from vertexai import model_garden
# Replace with your project ID and region
vertexai.init(project="<YOUR_PROJECT_ID>", location="<REGION>")
model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")
# Custom environment variables to override default settings
# This example sets N_GPUS as 2, so RING_DEGREE * ULYSSES_DEGREE must equal 2
container_env_vars = {
"N_GPUS": "2",
"RING_DEGREE": "1",
"ULYSSES_DEGREE": "2"
# Add other environment variables to customize here
}
machine_type = "a3-highgpu-2g"
accelerator_type = "NVIDIA_H100_80GB"
accelerator_count = 2
# Deploy the model with the custom environment variables
endpoint = model.deploy(
machine_type=machine_type,
accelerator_type=accelerator_type,
accelerator_count=accelerator_count,
container_env_vars=container_env_vars
)
N'oubliez pas de consulter la section "Comprendre les arguments spécifiques à xDiT" pour en savoir plus sur chaque variable d'environnement. Assurez-vous que le produit des degrés de parallélisme (par exemple, PIPFUSION_PARALLEL_DEGREE, ULYSSES_DEGREE, RING_DEGREE et USE_CFG_PARALLEL) est égal au nombre total de GPU (N_GPUS).
Pour obtenir d'autres exemples de recettes et de configurations de diffusion pour différents modèles, consultez la documentation officielle de xDiT. Pour en savoir plus sur le SDK Model Garden, consultez la documentation.