xDiT ist eine Open-Source-Bibliothek, die die Inferenz für Diffusion Transformer-Modelle (DiT) mithilfe von Parallelitäts- und Optimierungstechniken beschleunigt. Diese Techniken ermöglichen eine skalierbare Multi-GPU-Einrichtung für anspruchsvolle Arbeitslasten. Auf dieser Seite wird gezeigt, wie DiT-Modelle mit xDiT und Cloud-GPUs in Vertex AI bereitgestellt werden.
Weitere Informationen zu xDiT finden Sie im xDiT-GitHub-Projekt.
Vorteile
In der folgenden Liste werden die wichtigsten Vorteile der Verwendung von xDiT zum Bereitstellen von DiT-Modellen in Vertex AI beschrieben:
- Bis zu dreimal schneller generieren: Generieren Sie hochauflösende Bilder und Videos in einem Bruchteil der Zeit, die für andere Bereitstellungslösungen erforderlich ist.
- Skalierbare Unterstützung für mehrere GPUs: Arbeitslasten werden effizient auf mehrere GPUs verteilt, um eine optimale Leistung zu erzielen.
- Hybride Parallelität: xDiT unterstützt verschiedene parallele Verarbeitungsansätze wie einheitliche Sequenzparallelität, PipeFusion, CFG-Parallelität und Datenparallelität. Diese Methoden können in einem einzigartigen Rezept kombiniert werden, um die Leistung zu optimieren.
- Optimierte Leistung bei einer einzelnen GPU: xDiT ermöglicht schnellere Inferenz auch bei einer einzelnen GPU.
- GPU-Beschleunigung: xDiT enthält mehrere Kernel-Beschleunigungsmethoden und verwendet Techniken von DiTFastAttn, um die Inferenz auf einer einzelnen GPU zu beschleunigen.
- Einfache Bereitstellung: Mit der Ein-Klick-Bereitstellung oder Colab Enterprise-Notebooks in Vertex AI Model Garden können Sie schnell loslegen.
Unterstützte Modelle
xDiT ist für bestimmte DiT-Modellarchitekturen im Vertex AI Model Garden verfügbar, z. B. Flux.1 Schnell, CogVideoX-2b und Wan2.1-Varianten des Text-zu-Video-Modells. Ob ein DiT-Modell xDiT in Model Garden unterstützt, sehen Sie auf der entsprechenden Modellkarte in Model Garden.
Hybride Parallelität für Multi-GPU-Leistung:
xDiT verwendet eine Kombination aus Parallelitätstechniken, um die Leistung bei Setups mit mehreren GPUs zu maximieren. Diese Techniken arbeiten zusammen, um die Arbeitslast zu verteilen und die Ressourcennutzung zu optimieren:
- Einheitliche Sequenzparallelität: Bei dieser Technik werden die Eingabedaten (z. B. das Aufteilen eines Bildes in Patches) auf mehrere GPUs aufgeteilt, wodurch der Speicherverbrauch reduziert und die Skalierbarkeit verbessert wird.
- PipeFusion: Bei PipeFusion wird das DiT-Modell in Phasen unterteilt und jede Phase einer anderen GPU zugewiesen. So können verschiedene Teile des Modells parallel verarbeitet werden.
- CFG-Parallelität: Diese Technik optimiert Modelle speziell durch die Verwendung von Classifier-Free Guidance (CFG), einer gängigen Methode zur Steuerung des Stils und Inhalts generierter Bilder. Die Berechnung von bedingten und unbedingten Zweigen wird parallelisiert, was zu einer schnelleren Inferenz führt.
- Datenparallelität: Bei dieser Methode wird das gesamte Modell auf jeder GPU repliziert. Jede GPU verarbeitet einen anderen Batch von Eingabedaten, wodurch der Gesamtdurchsatz des Systems erhöht wird.
Weitere Informationen zu Leistungsverbesserungen finden Sie im Bericht von xDiT zu Flux.1 Schnell oder CogVideoX-2b. Google konnte diese Ergebnisse in Vertex AI Model Garden reproduzieren.
Beschleunigung mit einer einzelnen GPU
Die xDiT-Bibliothek bietet Vorteile für die Bereitstellung auf einer einzelnen GPU, da sie torch.compile und onediff verwendet, um die Laufzeitgeschwindigkeit auf GPUs zu verbessern. Diese Techniken können auch in Verbindung mit hybrider Parallelität verwendet werden.
xDiT verwendet außerdem eine effiziente Methode zur Berechnung der Aufmerksamkeit namens DiTFastAttn, um den Rechenengpass von DiT zu beheben. Derzeit ist diese Technik nur für einzelne GPUs oder in Verbindung mit Datenparallelität verfügbar.
Erste Schritte mit Model Garden
Der für xDiT optimierte Cloud-GPU-Bereitstellungscontainer ist in Vertex AI Model Garden verfügbar. Bei unterstützten Modellen wird dieser Container für Bereitstellungen verwendet, wenn Sie die Ein-Klick-Bereitstellung oder die Colab Enterprise-Notebook-Beispiele nutzen.
In den folgenden Beispielen wird das Modell „Flux.1-schnell“ verwendet, um zu veranschaulichen, wie ein DiT-Modell in einem xDiT-Container bereitgestellt wird.
Bereitstellung mit nur einem Klick verwenden
Sie können einen benutzerdefinierten Vertex AI-Endpunkt mit dem xDiT-Container mithilfe einer Modellkarte bereitstellen.
Rufen Sie die Seite „Modellkarte“ auf und klicken Sie auf Bereitstellen.
Wählen Sie für die zu verwendende Modellvariante einen Maschinentyp für die Bereitstellung aus.
Klicken Sie auf Bereitstellen, um den Bereitstellungsprozess zu starten. Sie erhalten zwei E-Mail-Benachrichtigungen: eine, wenn das Modell hochgeladen wird, und eine weitere, wenn der Endpunkt bereit ist.
Colab Enterprise-Notebook verwenden
Für mehr Flexibilität und Anpassungsmöglichkeiten können Sie Colab Enterprise-Notebook-Beispiele verwenden, um einen Vertex AI-Endpunkt mit dem xDiT-Container mithilfe des Vertex AI SDK für Python bereitzustellen.
Rufen Sie die Seite „Modellkarte“ auf und klicken Sie auf Notebook öffnen.
Wählen Sie das Vertex Serving-Notebook aus. Das Notebook wird in Colab Enterprise geladen.
Führen Sie das Notebook durch, um ein Modell mit dem xDiT-Container bereitzustellen und Vorhersageanfragen an den Endpunkt zu senden. Das Code-Snippet für die Bereitstellung sieht so aus:
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()
xDiT-Argumente
xDiT bietet eine Reihe von Serverargumenten, die konfiguriert werden können, um die Leistung für bestimmte Anwendungsfälle zu optimieren. Diese Argumente werden während der Bereitstellung als Umgebungsvariablen festgelegt. Im Folgenden sind die wichtigsten Argumente aufgeführt, die Sie möglicherweise konfigurieren müssen:
Modellkonfiguration
MODEL_ID(String): Gibt die zu ladende Modellkennung an. Er sollte mit dem Modellnamen in Ihrer Registrierung oder Ihrem Pfad übereinstimmen.
Argumente für die Laufzeitoptimierung
N_GPUS(Ganzzahl): Gibt die Anzahl der GPUs an, die für die Inferenz verwendet werden sollen. Der Standardwert ist 1.WARMUP_STEPS(Ganzzahl): Anzahl der Aufwärmschritte, die vor dem Beginn der Inferenz erforderlich sind. Das ist besonders wichtig, wenn PipeFusion aktiviert ist, um eine stabile Leistung zu gewährleisten. Der Standardwert ist 1.USE_PARALLEL_VAE(boolesch): Ermöglicht die effiziente Verarbeitung von Bildern mit hoher Auflösung (über 2048 Pixel), indem die VAE-Komponente auf mehrere Geräte verteilt wird. Dadurch werden Probleme mit dem Arbeitsspeicher bei großen Bildern vermieden. Der Standardwert ist „false“.USE_TORCH_COMPILE(boolesch): Aktiviert die Beschleunigung mit einer einzelnen GPU über torch.compile und bietet Optimierungen auf Kernelebene für eine bessere Leistung. Der Standardwert ist „false“.USE_ONEDIFF(boolesch): Aktiviert die OneDiff-Kompilierungsbeschleunigungstechnologie, um die Ausführungsgeschwindigkeit von GPU-Kernels zu optimieren. Der Standardwert ist „false“.
Argumente für die Datenparallelität
DATA_PARALLEL_DEGREE(Ganzzahl): Legt den Grad der Datenparallelität fest. Lassen Sie das Feld leer, um die Option zu deaktivieren, oder legen Sie den gewünschten Grad der Parallelität fest.USE_CFG_PARALLEL(boolesch): Aktiviert die parallele Berechnung für die klassifikatorfreie Steuerung (Classifier-Free Guidance, CFG), auch bekannt als Split Batch. Wenn diese Option aktiviert ist, ist der konstante Parallelitätsgrad 2. Wird auf „true“ gesetzt, wenn CFG zum Steuern von Ausgabestil und ‑inhalt verwendet wird. Der Standardwert ist „false“.
Argumente für die sequenzielle Parallelität (USP – Unified Sequence Parallelism)
ULYSSES_DEGREE(Ganzzahl): Legt den Ulysses-Grad für den parallelen Ansatz für einheitliche Sequenzen fest, der DeepSpeed-Ulysses und Ring-Attention kombiniert. Damit wird das All-to-All-Kommunikationsmuster gesteuert. Lassen Sie das Feld leer, um den Standardwert zu verwenden.RING_DEGREE(Ganzzahl): Legt den Ringgrad für die Peer-to-Peer-Kommunikation bei sequenzieller Parallelität fest. Wird in Verbindung mit ULYSSES_DEGREE verwendet, um das 2D-Prozess-Mesh zu bilden. Lassen Sie das Feld leer, um den Standardwert zu verwenden.
Tensor-parallele Argumente
TENSOR_PARALLEL_DEGREE(Ganzzahl): Legt den Grad des Tensorparallelismus fest. Dabei werden Modellparameter auf Geräte aufgeteilt, um den Speicherbedarf pro Gerät zu reduzieren. Leer lassen, um die Funktion zu deaktivieren.SPLIT_SCHEME(String): Definiert, wie die Modell-Tensoren auf Geräte aufgeteilt werden (z.B. nach Attention-Heads oder verborgenen Dimensionen). Lassen Sie das Feld leer, wenn das Standardschema für die Aufteilung verwendet werden soll.
Verteilte Ray-Argumente
USE_RAY(boolesch): Aktiviert das verteilte Ausführungs-Framework von Ray zum Skalieren von Berechnungen auf mehreren Knoten. Der Standardwert ist „false“.RAY_WORLD_SIZE(Ganzzahl): Gesamtzahl der Prozesse im Ray-Cluster. Der Standardwert ist 1.VAE_PARALLEL_SIZE(Ganzzahl): Anzahl der Prozesse, die für die parallele VAE-Verarbeitung mit Ray verwendet werden. Der Standardwert ist 0.DIT_PARALLEL_SIZE(Ganzzahl): Anzahl der Prozesse, die für die parallele Verarbeitung des DiT-Backbones mit Ray verwendet werden. Der Standardwert ist 0.
Parallele PipeFusion-Argumente
PIPEFUSION_PARALLEL_DEGREE(Ganzzahl): Legt den Grad der Parallelität für PipeFusion fest, eine Pipeline-Parallelität auf Sequenzebene, die die temporale Redundanz der Eingabe von Diffusionsmodellen nutzt. Höhere Werte erhöhen die Parallelität, erfordern aber mehr Arbeitsspeicher. Der Standardwert ist 1.NUM_PIPELINE_PATCH(Ganzzahl): Anzahl der Patches, in die die Sequenz für die Pipelineverarbeitung aufgeteilt werden soll. Lassen Sie das Feld leer, wenn die Sprache automatisch ermittelt werden soll.ATTN_LAYER_NUM_FOR_PP(String): Gibt an, welche Attention-Ebenen für die Pipeline-Parallelität verwendet werden sollen. Kann durch Kommas getrennt werden (z.B. „10,9“) oder durch Leerzeichen getrennt (z.B. „10 9“). Lassen Sie das Feld leer, um alle Ebenen zu verwenden.
Argumente zur Arbeitsspeicheroptimierung
ENABLE_MODEL_CPU_OFFLOAD(boolesch): Lagert Modellgewichte in den CPU-Arbeitsspeicher aus, wenn sie nicht verwendet werden. Dadurch wird die GPU-Speichernutzung auf Kosten einer erhöhten Latenz reduziert. Der Standardwert ist „false“.ENABLE_SEQUENTIAL_CPU_OFFLOAD(boolesch): Lagert Modellschichten während des Forward Pass sequenziell auf die CPU aus, sodass die Inferenz von Modellen möglich ist, die größer als der GPU-Arbeitsspeicher sind. Der Standardwert ist „false“.ENABLE_TILING(boolesch): Reduziert die GPU-Arbeitsspeichernutzung, indem die VAE-Komponente (Variational Autoencoder) kachelweise dekodiert wird. Dieses Argument ist nützlich für größere Bilder oder Videos und um Fehler aufgrund von zu geringem Arbeitsspeicher zu vermeiden. Der Standardwert ist „false“.ENABLE_SLICING(boolesch): Reduziert die GPU-Arbeitsspeichernutzung, indem der Eingabetensor für die VAE-Decodierung in Slices aufgeteilt wird. Der Standardwert ist „false“.
DiTFastAttn-Argumente (Aufmerksamkeitsoptimierung)
USE_FAST_ATTN(boolesch): Aktiviert die DiTFastAttn-Beschleunigung für die Inferenz mit einer einzelnen GPU. Dabei wird die temporale Eingabereduzierung verwendet, um die Rechenkomplexität zu verringern. Der Standardwert ist „false“.N_CALIB(Ganzzahl): Anzahl der Kalibrierungs-Samples für die DiTFastAttn-Optimierung. Der Standardwert ist 8.THRESHOLD(Gleitkommazahl): Ähnlichkeitsschwellenwert für die temporale Ähnlichkeitsreduzierung in DiTFastAttn. Der Standardwert ist 0,5.WINDOW_SIZE(Ganzzahl): Fenstergröße für Window Attention mit Residual Caching zur Reduzierung räumlicher Redundanz. Der Standardwert ist 64.COCO_PATH(string): Pfad zum COCO-Dataset für die DiTFastAttn-Abstimmung. Erforderlich, wenn USE_FAST_ATTN „wahr“ ist. Lassen Sie das Feld leer, wenn Sie es nicht verwenden.
Argumente für die Cache-Optimierung
USE_CACHE(boolesch): Aktiviert allgemeine Caching-Mechanismen, um redundante Berechnungen zu reduzieren. Der Standardwert ist „false“.USE_TEACACHE(boolesch): Aktiviert die TeaCache-Optimierungsmethode zum Zwischenspeichern von Zwischenergebnissen. Der Standardwert ist „false“.USE_FBCACHE(boolesch): Aktiviert die Optimierungsmethode „First-Block-Cache“. Der Standardwert ist „false“.
Argumente für die Optimierung der Genauigkeit
USE_FP8_T5_ENCODER(boolesch): Aktiviert die FP8-Präzision (8-Bit-Gleitkomma) für den T5-Text-Encoder, wodurch der Arbeitsspeicherbedarf reduziert und der Durchsatz möglicherweise verbessert wird, ohne dass die Qualität wesentlich beeinträchtigt wird. Der Standardwert ist „false“.
Bereitstellung von Anpassungen
Model Garden bietet standardmäßige xDiT-Parallelisierungskonfigurationen für unterstützte Modelle. Sie können diese Standardeinstellungen mit dem Vertex AI SDK für Python prüfen.
Wenn Sie die Standardkonfiguration für die Bereitstellung eines Modells wie „black-forest-labs/FLUX.1-schnell“ aufrufen möchten, können Sie das folgende Code-Snippet ausführen:
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/flux1-schnell@flux.1-schnell']
# [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"
# }
# ..........]
Die Methode list_deploy_options() gibt die Containerspezifikationen zurück, einschließlich der Umgebungsvariablen (env), die die xDiT-Konfiguration definieren.
Wenn Sie die Parallelisierungsstrategie anpassen möchten, können Sie diese Umgebungsvariablen beim Bereitstellen des Modells überschreiben. Im folgenden Beispiel wird gezeigt, wie Sie RING_DEGREE und ULYSSES_DEGREE für eine Konfiguration mit zwei GPUs ändern und dabei den Parallelisierungsansatz ändern:
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
)
Weitere Informationen zu den einzelnen Umgebungsvariablen finden Sie im Abschnitt „xDiT-spezifische Argumente“. Achten Sie darauf, dass das Produkt der Parallelitätsgrade (z.B. PIPEFUSION_PARALLEL_DEGREE,ULYSSES_DEGREE, RING_DEGREE und USE_CFG_PARALLEL) entspricht der Gesamtzahl der GPUs (N_GPUS).
Weitere Beispiele für Serving-Rezepte und Konfigurationen für verschiedene Modelle finden Sie in der offiziellen xDiT-Dokumentation. Weitere Informationen zum Model Garden SDK finden Sie in der Dokumentation.