Dieser Vortrag präsentiert den aktuellen Entwicklungsstand von Torch TPU, einer kollaborativen Initiative von Google und Meta zur nativen Ausführung von PyTorch-Workloads auf Google-TPUs. Die Vortragenden (Clauddio von Google und Colin von Meta) erläutern die architektonischen Grundlagen, die Integration von Software-Stacks wie Stable HLO und dem XLA-Compiler, neue Eager-Ausführungsmodi (inklusive des hybriden „Defer and Fuse“-Verfahrens) sowie produktionsreife Inferenz- und Trainingsintegrationen mit vLLM und SGLang.
Wichtigste Punkte & Zeitstempel
0:00 Motivation & Kollaboration: Google und Meta kooperieren bei Torch TPU, um fragmentierte Hardware-Flotten zu vereinheitlichen und den Entwicklungsaufwand beim Portieren von PyTorch-Modellen auf TPUs zu minimieren.
0:25 Code-Integration & Bereitstellung: Demonstration der Ausführung von nativem PyTorch-Code auf TPUs mittels torch.device("tpu"), Feinabstimmung von Llama 3 8B über LoRA und Ausführung über XPK auf 256 Ironwood-TPU-Chips.
6:03 Architektur-Stack: Der Software-Stack leitet PyTorch-Operationen über Stable HLO an den bewährten XLA-Compiler weiter, wobei konsequent PyTorch-native Programmierparadigmen bevorzugt werden.
9:16 Ausführungsmodi (Strict, Debug & Defer-and-Fuse): Neben dem standardmäßigen Strict-Modus für numerische Stabilität und dem blockierenden Debug-Modus führt Torch TPU den neuen Defer and Fuse-Modus ein, der Operationen heuristisch akkumuliert und durch Fusion über 50 % Leistungssteigerung liefert.
11:39 Kompilierungs-Cache & torch.compile: Implementierung eines persistenten Kompilierungscaches zur Vermeidung redundanter Neukompilierungen, wobei torch.compile über Dynamo und ein "Defer-all"-Verfahren direkt in den XLA-Stack eingespeist wird.
18:00 Benutzerdefinierte Kernel & Verteilung: Unterstützung von Palace- und Jax-Kerneln, Validierung von verteilten Paradigmen inklusive SPMD und MPND, FSDP v2 sowie Optimierungen für Sparse Cores.
36:11 Inferenz-Integrationen (vLLM & SGLang): Enger Zusammenschluss mit Inferenz-Engines durch Wiederverwendung von Scheduling, kontinuierlichem Batching und KV-Cache-Managern sowie Nutzung statischer Shape-Bucketisierung.
42:53 Benchmarks & Modellabdeckung: Umfangreiche Tests mit der Qwen 3-Familie (Dense, MoE, Vision-Language-Modelle), wobei auf Qwen 3.0 eine Leistungsparität oder -überlegenheit gegenüber dem reinen vLLM-TPU-Stack erreicht wird.
48:47 Trainings-Optimierungen: Einsatz eines make_effects-basierten Dispatch-Tracers für das Training, um vollständige Graphen inklusive Optimizer-Schritten direkt zu optimieren.
55:03 Ausblick & Open-Source-Pläne: Geplante Open-Source-Freigabe der Repositories und Integrationen zum Jahresende zeitgleich mit dem offiziellen Torch-TPU-Release.
Glossar
TPU (Tensor Processing Unit): Von Google entwickelte, spezialisierte ASIC-Prozessoreinheit, die für beschleunigte Machine-Learning- und Deep-Learning-Workloads optimiert ist.
PyTorch: Ein weit verbreitetes Open-Source-Framework für maschinelles Lernen, das auf Python basiert und für dynamische Berechnungsdiagramme bekannt ist.
Stable HLO: Eine stabile High-Level-Optimierungs-Zwischenrepräsentation (IR) für Machine-Learning-Compiler, die als standardisierte Schnittstelle zwischen Frontend-Frameworks und Backend-Compilern dient.
XLA (Accelerated Linear Algebra): Ein domainspezifischer Compiler für lineare Algebra von Google, der Berechnungsdiagramme von Machine-Learning-Frameworks optimiert und in ausführbare Binärdateien übersetzt.
LoRA (Low-Rank Adaptation): Eine effiziente Fine-Tuning-Technik für große Sprachmodelle, bei der nur eine minimale Anzahl von Gewichtsparametern durch Rang-Zerlegung angepasst wird.
Vertex AI: Eine von Google Cloud bereitgestellte Managed-Machine-Learning-Plattform zum Trainieren, Verwalten und Bereitstellen von KI-Modellen.
XPK: Ein Open-Source-Tool zur Orchestrierung und Bereitstellung von Workloads auf Google Cloud TPUs im Großmaßstab.
Ironwood: Eine spezifische Generation von Google-TPU-Chips.
Eager Mode: Ein Ausführungsmodus in Frameworks wie PyTorch, bei dem Operationen sofort beim Aufruf ausgeführt werden, anstatt sie vorab zu kompilieren.
Dynamo (TorchDynamo): Eine PyTorch-Frontend-Komponente, die Python-Bytecode zur Laufzeit abfängt und in ein FX-Berechnungsdiagramm zur Kompilierung umwandelt.
Inductor: Der standardmäßige Code-Generator und PyTorch-Compiler für torch.compile, der primär für GPUs optimiert ist.
Palace Kernel: Eine domänenspezifische Programmiersprache (DSL) bzw. ein Kernel-Schreibstil zur Erstellung von benutzerdefinierten Low-Level-Tuning-Kerneln auf TPUs.
SPMD (Single Program, Multiple Data): Ein paralleles Programmiermodell, bei dem dasselbe Programm auf mehreren Prozessoren mit unterschiedlichen Daten ausgeführt wird.
MPND (Multi-Program, Multiple Data): Ein paralleles Programmiermodell, bei dem verschiedene Programme auf verschiedenen Prozessoren mit unterschiedlichen Daten ausgeführt werden.
FSDP (Fully Sharded Data Parallel): Eine Speicher- und Rechenoptimierungstechnik im verteilten Training, bei der Modellparameter, Gradienten und Optimizer-Zustände über mehrere Beschleuniger gesharded (aufgeteilt) werden.
vLLM / SGLang: Beliebte Open-Source-Inferenz-Engines für Large Language Models, die auf hohem Durchsatz, Paged Attention und effizienter Speicherverwaltung basieren.
KV-Cache (Key-Value Cache): Ein Zwischenspeicher für Schlüssel- und Wertvektoren bei der Transformer-Inferenz, der die Neuberechnung bereits generierter Token-Historien vermeidet.
MoE (Mixture of Experts): Eine Architektur für neuronale Netze, bei der nur eine kleine Teilmenge von Teilnetzwerken ("Experten") dynamisch für jeden Token aktiviert wird.
Qwen 3 / Llama 3: Bekannte Open-Weights-Familien von großen Sprachmodellen.