In diesem Vortrag präsentiert Sayak Paul, Senior Research Engineer bei Hugging Face, die Entwicklung einer TPU-optimierten Implementierung des Bildgenerierungsmodells „Qwen Image“ von Alibaba. Der Beitrag erläutert die Funktionsweise von Latent Diffusion Models (LDM) als mehrstufige Pipelines im Gegensatz zu Monolithen sowie die Integration der Hugging Face diffusers-Bibliothek mit PyTorch XLA. Um das Modell effizient auf Google TPU v6e-8 Hardware auszuführen, wurde das Text-Encoder-Modell (Qwen 2.5 VL) auf der CPU belassen, während der speicherintensive Diffusion Transformer (DiT) mittels einer 2D-Mesh-SPMD-Strategie über 8 TPU-Chips hinweg geshardet wurde. Trotz einer initialen XLA-Kompilierungszeit von 95 Minuten erreichte das System im Steady-State eine Inferenzzeit von 20 Sekunden für 50 Denoising-Schritte und übertraf damit eine ungeshardete A100-Referenzkonfiguration. Abschließend wird der zukünftige Wechsel zu torch_tpu zur vereinfachten Portierbarkeit diskutiert.
Wichtigste Punkte & Zeitstempel
0:00 TPU-Portierung von Qwen Image: Vorstellung des Projekts zur Bereitstellung von Alibabas Open-Source-Bildgenerierungsmodell auf Google TPUs unter Nutzung von PyTorch XLA und Hugging Face diffusers.
1:33 Architektur von Latent Diffusion Models: Bildgenerierungsmodelle arbeiten vorwiegend in einem niederdimensionalen Latentraum und nutzen iterative Denoising-Schritte anstelle klassischer autoregressiver Verfahren.
3:16 Pipeline-Komponenten: Das Qwen-Image-System besteht aus mehreren separaten Teilmodellen: dem Text-Encoder Qwen 2.5 VL, dem iterativen Diffusion Transformer (DiT) und dem Variational Autoencoder (VAE) als Decoder.
5:44 PyTorch XLA-Integration in Diffusers: Die Einbindung in PyTorch XLA erfolgt primär über gezielte mark_step()-Aufrufe zur Steuerung der Kompilierungsgrenzen sowie den Einsatz experimenteller XLA FlashAttention.
9:19 Hybride CPU-TPU-Ausführungsstrategie: Der Text-Encoder wird auf der CPU ausgeführt, da er nur 1 % der Rechenlast ausmacht und somit Kompilierungsaufwand spart, während DiT und VAE auf dem TPU-Beschleuniger laufen.
12:03 SPMD-Sharding auf TPU v6e-8: Da der Diffusion Transformer nicht auf einen einzelnen TPU-Chip passt, wird das Modell mittels 2D-Mesh und Single-Process-SPMD entlang der größten Parameterdimensionen über 8 Chips verteilt.
15:16 XLA-Debugging & Fehlerbehebung: Behebung von Fehlern bei der Prompt-Kodierung sowie Eliminierung negativer Indexierungsfehler, die von XLA im Gegensatz zu CUDA/CPU explizit als Laufzeitverletzungen ausgewiesen werden.
16:24 Kompilierung & Inferenz-Latenz: Nach einer initialen XLA-Kompilierungszeit von 95 Minuten sinkt die Inferenzdauer auf 20 Sekunden pro Bild (50 Schritte) auf der TPU v6e-8 im Vergleich zu ~35 Sekunden auf einer NVIDIA A100.
17:48 Laufzeitanalyse und Profiling: Profiling-Ergebnisse bestätigen hohe Graphdichte, effektive Operationsfusion und minimale Host-to-Device-Transfers im kompilierten Ausführungsgraphen.
19:31 Ausblick auf Torch TPU: Künftige Implementierungen zielen auf die Nutzung von torch_tpu ab, um das manuelle Setzen von Kompilierungsgrenzen zu erübrigen und die Code-Portierbarkeit zwischen CUDA und TPU weiter zu vereinfachen.
Glossar
Latent Diffusion Model (LDM): Ein Bildgenerierungsmodell, das den iterativen Rauschentfernungsprozess (Denoising) in einem komprimierten latenten Vektorraum statt direkt auf hochauflösenden Bildpixeln ausführt.
Diffusion Transformer (DiT): Eine Modellarchitektur für Diffusionsmodelle, die auf dem Transformer-Prinzip (Selbstaufmerksamkeitsmechanismen) basiert und ältere U-Net-Netzwerke ersetzt.
PyTorch XLA: Eine Bibliothek, die das PyTorch-Framework mit dem XLA-Compiler (Accelerated Linear Algebra) verbindet, um PyTorch-Code auf Google TPUs und anderen Beschleunigern auszuführen.
Single Process Multiple Data (SPMD): Ein paralleles Programmiermodell, bei dem derselbe Ausführungscode gleichzeitig auf mehreren Prozessoren läuft, wobei jeder Prozessor unterschiedliche Datensegmente verarbeitet.
Sharding: Die Aufteilung großer Modellparameter oder Datensätze auf mehrere physische Speicherbausteine bzw. TPU-Chips zur Bewältigung von Arbeitsspeicherbeschränkungen.
FlashAttention: Ein speichereffizienter und exakter Aufmerksamkeitsalgorithmus (Attention Mechanism), der den E/A-Lese- und Schreibzugriff zwischen GPU/TPU-SRAM und HBM optimiert.
Variational Autoencoder (VAE): Ein neuronales Netzwerk, das Eingabedaten in einen latenten Raum kodiert und daraus neue Daten rekonstruieren bzw. dekodieren kann.
Tensor Processing Unit (TPU): Ein von Google entwickelter anwendungsspezifischer integrierter Schaltkreis (ASIC), der speziell für die Beschleunigung von Algorithmen des maschinellen Lernens konzipiert ist.