Diffusion Transformer-Modelle mit xDiT-Containern auf Cloud-GPUs bereitstellen

xDiT ist eine Open-Source-Bibliothek, die die Inferenz für Diffusion Transformer (DiT)-Modelle mithilfe von Parallelisierungs- und Optimierungstechniken beschleunigt. Diese Verfahren ermöglichen eine skalierbare Multi-GPU-Einrichtung für anspruchsvolle Arbeitslasten. Auf dieser Seite wird gezeigt, wie Sie DiT-Modelle mit xDiT und Cloud-GPUs in Vertex AI bereitstellen.

Weitere Informationen zu xDiT finden Sie im GitHub-Projekt für xDiT.

Vorteile

In der folgenden Liste sind die wichtigsten Vorteile der Verwendung von xDiT zum Bereitstellen von DiT-Modellen in Vertex AI aufgeführt:

  • Bis zu dreimal schnellere Generierung: Bilder und Videos in hoher Auflösung werden im Vergleich zu anderen Lösungen für die Bereitstellung in einem Bruchteil der Zeit generiert.
  • Skalierbare Unterstützung mehrerer GPUs: Arbeitslasten können effizient auf mehrere GPUs verteilt werden, um eine optimale Leistung zu erzielen.
    • Hybride Parallelität: xDiT unterstützt verschiedene Ansätze für die parallele Verarbeitung, z. B. Unified Sequence Parallelism, PipeFusion, CFG Parallelism und Data Parallelism. Diese Methoden können kombiniert werden, um die Leistung zu optimieren.
  • Optimierte Leistung einer einzelnen GPU: Mit xDiT können Sie auch mit einer einzelnen GPU schnellere Inferenzen erzielen.
    • GPU-Beschleunigung: xDiT umfasst mehrere Methoden zur Kernelbeschleunigung und verwendet Techniken aus DiTFastAttn, um die Inference auf einer einzelnen GPU zu beschleunigen.
  • Einfache Bereitstellung: Mit der Bereitstellung mit nur einem Klick 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 und CogVideoX-2b. Ob ein DiT-Modell xDiT in Model Garden unterstützt, sehen Sie auf der Modellkarte im Model Garden.

Hybride Parallelität für die Leistung mehrerer GPUs:

xDiT verwendet eine Kombination aus Parallelisierungstechniken, um die Leistung bei Multi-GPU-Konfigurationen zu maximieren. Diese Techniken werden kombiniert, um die Arbeitslast zu verteilen und die Ressourcennutzung zu optimieren:

  • Unified Sequence Parallelism: Bei dieser Methode werden die Eingabedaten (z. B. ein Bild in Patches) auf mehrere GPUs aufgeteilt, wodurch die Speichernutzung reduziert und die Skalierbarkeit verbessert wird.
  • PipeFusion: PipeFusion teilt das DiT-Modell in Phasen auf und weist jeder Phase eine andere GPU zu. So ist die parallele Verarbeitung verschiedener Teile des Modells möglich.
  • CFG-Parallelität: Mit dieser Methode werden Modelle speziell durch die Verwendung einer klassifikatorfreien Anleitung optimiert. Dies ist eine gängige Methode zur Steuerung des Stils und Inhalts generierter Bilder. Die Berechnung von bedingten und unbedingten Verzweigungen 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 dabei einen anderen Batch von Eingabedaten, wodurch der Gesamtdurchsatz des Systems erhöht wird.

Weitere Informationen zu Leistungsverbesserungen findest du im xDiT-Bericht zu Flux.1 Schnell oder CogVideoX-2b. Google konnte diese Ergebnisse in Vertex AI Model Garden reproduzieren.

Beschleunigung durch eine einzelne GPU

Die xDiT-Bibliothek bietet Vorteile für die Bereitstellung mit einer einzelnen GPU, da torch.compile und onediff verwendet werden, um die Laufzeitgeschwindigkeit auf GPUs zu verbessern. Diese Techniken können auch in Kombination mit hybridem Parallelismus verwendet werden.

xDiT hat außerdem eine effiziente Methode zur Berechnung der Aufmerksamkeit, die DiTFastAttn, um den Rechenengpass von DiT zu beheben. Derzeit ist diese Methode nur für einzelne GPU-Konfigurationen 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 Ein-Klick-Bereitstellungen oder die Colab Enterprise-Notebook-Beispiele verwenden.

In den folgenden Beispielen wird anhand des Modells „Flux.1-schnell“ veranschaulicht, 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.

  1. Rufen Sie die Seite „Modellkarte“ auf und klicken Sie auf Bereitstellen.

  2. Wählen Sie für die zu verwendende Modellvariante einen Maschinentyp für die Bereitstellung aus.

  3. Klicken Sie auf Bereitstellen, um den Bereitstellungsprozess zu starten. Sie erhalten zwei E-Mail-Benachrichtigungen: eine, wenn das Modell hochgeladen wurde, 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.

  1. Rufen Sie die Seite „Modellkarte“ auf und klicken Sie auf Notebook öffnen.

  2. Wählen Sie das Vertex Serving-Notebook aus. Das Notebook wird in Colab Enterprise geladen.

  3. Führen Sie das Notebook aus, 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:

XDIT_DOCKER_URI=us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/xdit-serve.cu125.0-1.ubuntu2204.py310

serving_env = {
    "MODEL_ID": "black-forest-labs/FLUX.1-schnell",
    "TASK": "text-to-image",
    "DEPLOY_SOURCE": "notebook",
    "N_GPUS": "2",
    "ULYSSES_DEGREE": "1",
    "RING_DEGREE": "2",
    "PIPEFUSION_PARALLEL_DEGREE": "1",
    "USE_TORCH_COMPILE": "true",
}

model = aiplatform.Model.upload(
    display_name=model_name,
    serving_container_image_uri=XDIT_DOCKER_URI,
    serving_container_ports=[7080],
    serving_container_predict_route="/predict",
    serving_container_health_route="/health",
    serving_container_environment_variables=serving_env,
)

model.deploy(
    endpoint=endpoint,
    machine_type="a3-highgpu-2g",
    accelerator_type="NVIDIA_H100_80GB",
    accelerator_count=2,
    deploy_request_timeout=1800,
    service_account=SERVICE_ACCOUNT,
)

Umgebungsvariablen

  • MODEL_ID: Gibt die ID des bereitzustellenden DiT-Modells an, z. B. „black-forest-labs/FLUX.1-schnell“.
  • TASK: Definiert die Aufgabe, die das Modell ausführt, z. B. „text-to-image-flux-xdit“.
  • N_GPUS: Legt die Anzahl der GPUs fest, die für die Inferenz verwendet werden sollen.
  • ULYSSES_DEGREE, RING_DEGREE, PIPEFUSION_PARALLEL_DEGREE: Hiermit können Sie die Parallelisierungstechniken steuern, die von xDiT verwendet werden. Weitere Informationen zu den einzelnen Argumenten finden Sie unter xDiT-Argumente.
  • USE_TORCH_COMPILE: Ermöglicht die Beschleunigung mit einer einzelnen GPU mithilfe von torch.compile.

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. In der folgenden Liste sind die wichtigsten Argumente aufgeführt, die Sie möglicherweise konfigurieren müssen:

  • N_GPUS (ganzzahl): Gibt die Anzahl der GPUs an, die für die Inferenz verwendet werden sollen. Der Standardwert ist 1.
  • ENABLE_TILING (boolescher Wert): Verringert die GPU-Arbeitsspeichernutzung, indem die VAE-Komponente jeweils eine Kachel nach der anderen decodiert wird. Dieses Argument ist nützlich für größere Bilder oder Videos und um Speichermangel-Fehler zu vermeiden. Der Standardwert ist false.
  • ENABLE_SLICING (boolescher Wert): Verringert die GPU-Arbeitsspeichernutzung, indem der Eingabetensor für die VAE-Dekodierung in Scheiben aufgeteilt wird. Der Standardwert ist false.
  • USE_TORCH_COMPILE (boolescher Wert): Aktiviert die Beschleunigung durch eine einzelne GPU über torch.compile und erhöht so die Kompilierungsgeschwindigkeit. Der Standardwert ist false.
  • PIPEFUSION_PARALLEL_DEGREE (Ganzzahl): Legt den Parallelisierungsgrad für PipeFusion fest. Höhere Werte erhöhen die Parallelität, erfordern aber möglicherweise mehr Arbeitsspeicher. Der Standardwert ist 1.
  • WARMUP_STEPS (Ganzzahl): Wenn PipeFusion aktiviert ist, gibt dieses Argument die Anzahl der Warm-up-Schritte an, die vor Beginn der Inferenz erforderlich sind. Der Standardwert ist 0.
  • ULYSSES_DEGREE (Ganzzahl): Damit wird der Ulysses-Grad festgelegt. Der Standardwert ist 1.
  • RING_DEGREE (Ganzzahl): Legt den Ringgrad fest. Der Standardwert ist 1.
  • USE_CFG_PARALLEL (boolescher Wert): Ermöglicht die parallele Berechnung für die klassifizierungsfreie Anleitung (CFG), eine Methode, mit der die Ausgabe von DiT-Modellen gesteuert wird. Wenn diese Option aktiviert ist, ist der Grad der konstanten Parallelität 2. Wird bei Verwendung von CFG auf „true“ gesetzt. Der Standardwert ist false.
  • USE_PARALLEL_VAE (boolescher Wert): Ermöglicht die effiziente Verarbeitung von Bildern mit hoher Auflösung (über 2.048 Pixel), indem die VAE-Komponente parallelisiert wird. Der Standardwert ist false.

Eine vollständige Liste der Argumente finden Sie in der xFuserArgs-Klasse im GitHub-Projekt von xDiT.