Dieser Vortrag eines Google-Produktmanagers für KI-Infrastruktur analysiert das Profiling, Debugging und Optimieren von Machine-Learning-Workloads auf Google Tensor Processing Units (TPUs) mithilfe der Toolsuite XPROF. Nach einer Einführung in die TPU-Architektur – insbesondere die Balance zwischen Tensor Cores für dichte Matrixmultiplikationen, SparseCores für speichergebundene Operationen, High Bandwidth Memory (HBM) mit bis zu 35 TB/s und latenzarmen Inter-Chip Interconnects (ICI) – wird die Vermeidung von Ineffizienzen („Flying Blind“) adressiert. Anhand einer Fallstudie (Mistral 8x22B auf Trillium TPU v6e) wird demonstriert, wie die Vergrößerung der Batch-Größe von 1 auf 8 das Verhältnis von Rechenleistung zu Kommunikation optimiert: Die FLOP-Auslastung steigt von 27,6 % auf 60,74 %, der AllReduce-Overhead sinkt von 31 % auf unter 9 % und die Speichernutzung wird von 12 GB auf 29 GB (von 32 GB) skaliert. Abschließend stellt der Vortrag Erweiterungen von XPROF vor, darunter Low-Level-Kernel-Profiling (MLIR, LLO-Instruktionsbündel), ereignisbasierte Telemetrie mit Zugriff auf über 16.000 Hardware-Register (TPU v7 „Ironwood“), kontinuierliche 90-Sekunden-Snapshots via 2-GB-Ringpuffer sowie native Out-of-Tree-Integration in das PyTorch-Ökosystem via torch.profiler.
Wichtigste Punkte & Zeitstempel
0:00 Problemstellung bei mangelnder Hardware-Transparenz: Blindes Ausführen von ML-Workloads ohne Profiling führt zu verzögerten Forschungszyklen, explodierenden Cloud-Kosten und unidentifizierten Engpässen durch unausgelastete Beschleuniger.
0:01 TPU-Systemarchitektur und Rechenteilung: Maximale Auslastung basiert auf Systembalance. Die Rechenlast teilt sich auf Tensor Cores (dichte Matrixmultiplikation) und SparseCores (speicherintensive Einbettungs-Lookups) auf, gespeist durch HBM mit Bandbreiten von rund 35 TB/s und direkte ICI-Chip-Kopplung unter Umgehung der Host-CPU.
0:03 Datenfluss und Phasenunterschiede: Der Host-CPU-Orchestrator steuert Vorverarbeitung sowie XLA-Kompilierung via PCIe zur TPU. Training fokussiert Vorwärts-/Rückwärtspfade und Gradienten-Updates; Serving priorisiert Kompilierung, Quantisierung und dynamisches Batching.
0:04 Ursachen von Memory-Bound-Zuständen: Wenn Matrixmultiplikationseinheiten (MXUs) Daten schneller verarbeiten, als HBM sie über das Vektor-Memory (VM) in Register laden kann, verhungert die Recheneinheit und gerät in ein Speicher-Engpass-Szenario.
0:05 Fallstudie Mistral 8x22B (HLO-Op-Profil): Bei Ausführung auf Trillium (4K Sequenzlänge, Batch-Größe 1) blockiert die Kommunikationsoperation AllReduce 31 % der Laufzeit bei nur 27,6 % TPU-FLOP-Auslastung, da Chips auf Gradientenabgleiche warten.
0:06 Trace-Viewer-Analyse der Kollektivkommunikation: Eine Rückwärtsschicht benötigt 89 ms, wovon AllReduce 33 ms belegt. Die Datenübertragungsgröße entspricht mit 8×16K×16K exakt den Feed-Forward-Gewichten und ist unabhängig von der Batch-Größe.
0:07 Memory-Viewer und Speichergrenzen: Das Tool visualisiert dynamische Zuweisungen und den Speicher-Höchststand (High-Water Mark) zur OOM-Diagnose. Die Baseline nutzt lediglich 12 GB des verfügbaren 32-GB-Speichers aus.
0:08 Performance-Steigerung durch Batch-Skalierung: Das Anheben der Batch-Größe auf 8 verschiebt die Auslastung von 27,6 % auf 60,74 % FLOPs, drückt den AllReduce-Zeitanteil auf 6,7–8,7 %, steigert die Speichernutzung auf 29 GB/32 GB und überführt den Workload in den angestrebten Compute-Bound-Zustand.
0:09 Kernel-Profiling für Custom-Code: Zur Analyse spezialisierter Kernel (Pallas, Mosaic, Triton), die in Standard-Profilern als Blackbox erscheinen, bietet XPROF Tiefeninspektion von MLIR-Code über XLA-Debug-Flags (xla_enable_custom_call).
0:11 LLO-Instruktionsbündel und statische Analyse: XPROF visualisiert maschinennahe Instruktionsbündel pro Taktzyklus für MXU-Auslastung, Skalar-/Vektor-ALUs, Vektorfüllungen (HBM-zu-Register) sowie Cross-Core-Datenbewegungen (XLU).
0:12 Dynamische Hardware-Zähler (TPU v7 Ironwood): Statische Analysen übersehen Pipeline-Stalls; dynamisches Sampling von Hardware-Zählern auf Sub-Mikrosekunden-Ebene deckt Synchronisationspausen und die reale Auslastungslücke auf.
0:13 Hierarchie der Metriken und Rohdaten-Register: Hardware-Registermetriken von HBM und TPU bilden die absolute Entscheidungsgrundlage („Ground Truth“). XPROF stellt dafür über 16.000 unkorrumpierte Rohdaten-Hardware-Counter tabellarisch bereit.
0:14 Kontinuierliche Profiling-Snapshots: Ein stets aktiver 2-GB-Ringpuffer auf dem Host zeichnet die letzten 90 Sekunden Telemetrie bei minimalem Overhead auf und ermöglicht programmatische Snapshots direkt beim Auftreten sporadischer Regressionen.
0:15 Native PyTorch-Integration: Über Kineto bzw. Torch-TPU bindet sich XPROF nativ und out-of-tree ein, wodurch Entwickler ohne modifizierten Code direkt die Standard-API torch.profiler auf TPUs nutzen können.
0:17 Ressourcen und Profiling Summit: Bereitstellung von Repositories und Dokumentation über OpenXLA, JAX-Guides und den TPU Developer Hub sowie Ankündigung des jährlichen Profiling Summits (20.–21. August).
Glossar
AllReduce: Eine kollektive Kommunikationsoperation im verteilten Rechnen, die Daten (z. B. Gradienten) von allen beteiligten Chips sammelt, mathematisch aggregiert (z. B. summiert) und das Gesamtergebnis an alle Knoten zurückverteilt.
Compute-Bound: Zustand, in dem die Ausführungsgeschwindigkeit eines Programms primär durch die Rechenleistung der Prozessorkerne (z. B. FLOP-Kapazität) und nicht durch Speicherbandbreite oder Netzwerktransfer limitiert wird.
FLOPs / FLOP-Auslastung (Floating Point Operations per Second): Maß für die Anzahl der Gleitkommaoperationen pro Sekunde; die Auslastung beziffert den prozentualen Anteil der real erreichten Rechenleistung gegenüber dem theoretischen Hardware-Maximum.
HBM (High Bandwidth Memory): Hochleistungsspeicherarchitektur mit gestapelten Speicherchips, die extrem breite Datenbusse und Durchsatzraten (im Vortrag ca. 35 TB/s) für Beschleuniger bereitstellt.
High-Water Mark: Der exakte Spitzenwert des dynamisch allozierten Arbeitsspeichers während der gesamten Laufzeit eines Programms; kritisch für die Vorbeugung von Speicherüberläufen.
HLO (High-Level Optimizer): Eine Zwischenrepräsentation (Intermediate Representation) im XLA-Compiler, die Rechenoperationen auf Graphebene abstrahiert und optimiert.
ICI (Inter-Chip Interconnect): Ein proprietäres Hochgeschwindigkeitsnetzwerk von Google, das TPUs direkt chip-zu-chip ohne Umweg über Host-CPU oder standardmäßige Netzwerkkarten miteinander verbindet.
Ironwood (TPU v7): Codename für die siebte Generation der Tensor Processing Units von Google.
Kineto: Eine von Meta entwickelte Profiling-Bibliothek für PyTorch, die Hardware-Aktivitäten von Beschleunigern erfasst.
LLO (Low-Level Optimizer): Compiler-Stufe unterhalb von HLO, die maschinennahe Instruktionsbündel direkt für die Funktionseinheiten der TPU generiert.
Memory-Bound: Zustand, in dem Recheneinheiten auf Daten warten müssen, weil die Speicherbandbreite oder die Latenz des Speichertransfers die Verarbeitungsgeschwindigkeit begrenzt.
Mistral 8x22B: Ein hochdimensionales Open-Weight-Sprachmodell, basierend auf einer Mixture-of-Experts-Architektur (MoE).
MLIR (Multi-Level Intermediate Representation): Ein modulares Compiler-Framework zur Standardisierung und Optimierung von Zwischencode zwischen High-Level-Frameworks und Maschinencode.
MXU (Matrix Multiply Unit): Dedizierte Hardware-Rechenwerke innerhalb des Tensor Cores einer TPU, die auf extrem parallele Matrixmultiplikationen spezialisiert sind.
OOM (Out of Memory): Programmabbruch infolge der vollständigen Erschöpfung des physikalisch verfügbaren Arbeitsspeichers.
Pallas / Mosaic / Triton: Spezialisierte Programmier-Frameworks und Domain Specific Languages (DSLs) zur manuellen Implementierung hocheffizienter GPU-/TPU-Rechenkerne.
Performance Counter: Physikalische Zählerregister auf dem Chip, die spezifische Hardware-Ereignisse (z. B. Taktzyklen, Cache-Misses, Pipeline-Stalls) in Echtzeit zählen.
SparseCore: Spezieller Rechenkern auf TPUs zur beschleunigten Abwicklung unregelmäßiger, spärlicher Speicherzugriffe (z. B. Embedding-Lookups), die klassische Matrix-Einheiten ineffizient belasten würden.
Tensor Core: Hauptrechenkern einer TPU, spezialisiert auf dichte, hochgradig parallele Matrixberechnungen.
Trillium (TPU v6e): Die sechste Generation der Google Cloud TPU, optimiert für Energieeffizienz und Skalierbarkeit bei KI-Modellen.
XLA (Accelerated Linear Algebra): Ein domänenspezifischer Compiler für Machine Learning, der Frameworks wie JAX, PyTorch oder TensorFlow in optimierten Hardware-Binärcode überführt.
XLU: Hardware-Funktionseinheit auf der TPU für Datenumordnung (Data Shuffling) und Kommunikation zwischen Prozessorkernen.
XPROF: Die integrierte Performance-Profiling-Software-Suite von Google zur Diagnose und Optimierung von ML-Workloads auf TPUs und GPUs.