Dieser Vortrag von Kai Wong (Rice University) präsentiert einen neuartigen Ansatz für das gemeinsame Sharding und die Platzierung von Diffusion Transformer (DiT)-Modellen zur Videogenerierung auf Tensor Processing Units (TPUs). Da über 97 % der Rechenzeit auf das DiT-Modell entfallen und hardwarebedingte Netzwerkkonflikte (Contention) zu erheblichen Latenzschwankungen führen, stellt der Autor ein zweistufiges Suchframework vor. Dieses kombiniert die schnelle Grobabschätzung von StableHLO mit einer detaillierten Kompilierungsanalyse, um optimale parallele Strategien ohne Quellcode-Modifikationen zu ermitteln.
Wichtigste Punkte & Zeitstempel
0:00 Hohe Rechenlast von DiT: Diffusion Transformers erzeugen sehr lange Sequenzen ohne KV-Cache, wodurch über 97 % der gesamten Rechenzeit auf das DiT-Modell entfallen.
0:55 Hardware-Topologie und Latenz: TPUv5- und TPUv6e-Subslices nutzen ein 2D-Mesh, bei dem gleichzeitiger Datenverkehr zu Netzwerk-Contention und stark variierenden Ausführungslatenzen führt, selbst bei identischem Sharding.
2:20 Platzierung als Ausführungsplan: Aufgrund von Bandbreiten-Degradationen bei gleichzeitigem Datenverkehr muss die Geräteplatzierung als inhärenter Bestandteil des Compiler-Ausführungsplans behandelt werden.
3:26 Rolle des JAX-Compilers: Der JAX-Compiler nutzt StableHLO und Shardy für die Sharding-Propagierung, während erst das kompilierte HLO exakte Kommunikationsprimitive (Collectives) offenlegt.
5:33 Zweistufiges Suchframework: Die erste Stufe nutzt StableHLO für eine günstige Grobabschätzung (prunt 85 % der Konfigurationen), während die zweite Stufe die verbleibenden Top-15 % via Compiler-HLO präzise bewertet.
6:56 Integration in MaxDiffusion: Der ermittelte Ausführungsplan wird direkt in MaxDiffusion integriert, ohne Modifikationen am Modellcode, an Low-Level-Lowering-Prozessen oder an Kerneln vorzunehmen.
8:59 TPUv5- und TPUv6e-Unterstützung: Der Ansatz konzentriert sich auf TPUv5- und TPUv6e-Architekturen, deren 2D-Mesh-Topologien mathematisch effizient modellierbar sind.
Glossar
TPU (Tensor Processing Unit): Ein von Google entwickelter, anwendungsspezifischer Prozessor (ASIC), der speziell zur Beschleunigung von Workloads im Bereich des maschinellen Lernens entwickelt wurde.
Diffusion Transformer (DiT): Eine generative KI-Architektur (eingesetzt u. a. für die Videogenerierung), die Diffusionsprozesse mit der Transformer-Architektur anstelle klassischer U-Net-Strukturen kombiniert.
Sharding: Das Aufteilen großer Tensoren, Datensätze oder Modellgewichte auf mehrere Hardware-Beschleuniger, um Speicher- und Rechenlasten zu verteilen.
KV Cache: Ein Zwischenspeicher für Schlüssel- und Wert-Vektoren in Transformator-Modellen, der die autoregressive Generierung beschleunigt, bei DiT-Modellen jedoch oft nicht vorhanden ist.
Tensor-Parallelismus: Eine parallele Verarbeitungsmethode, bei der einzelne mathematische Operationen (Matrizenmultiplikationen von Tensoren) auf mehrere Prozessoren aufgeteilt werden.
Kontext-Parallelismus: Eine Technik zur Verteilung extrem langer Eingabesequenzen (Kontexte) über mehrere Recheneinheiten, um den Speicherbedarf pro Gerät zu senken.
XLA (Accelerated Linear Algebra): Ein domänenspezifischer Compiler für lineare Algebra, der Rechengraphen von Frameworks wie JAX oder TensorFlow analysiert und für Hardware optimiert.
StableHLO: Ein standardisiertes, stabiles Zwischenformat (High-Level Operations) für den portablen Austausch von Rechengraphen zwischen ML-Compilern und Frameworks.
HLO (High-Level Optimizer): Eine maschinennahe, optimierte Zwischenrepräsentation im Compiler, die explizite Kommunikationsanweisungen und Kernel-Strukturen offenlegt.
Kernel Fusion: Ein Compiler-Optimierungsverfahren, bei dem mehrere aufeinanderfolgende Operationen zu einem einzigen Ausführungsschritt zusammengefasst werden, um teure Speicherzugriffe zu vermeiden.
MaxDiffusion: Ein Open-Source-Framework von Google, das optimierte Implementierungen von Diffusionsmodellen für die Ausführung auf TPU-Clustern bereitstellt.