Diese Präsentation stellt die Werkzeuge und Methoden zur Profilierung und Fehlerbehebung von JAX- und OpenXLA-Workloads mittels XPro auf Google TPUs und NVIDIA-GPUs vor.
XPro bietet spezialisierte Ansichten zur Leistungsevaluierung (Trace Viewer, Graph Viewer, Memory Viewer und Roofline-Analysen). Während Standard-HLO-Operationen deterministisch analysiert werden können, stellen Custom Calls (z. B. Palace- oder Triton-Kernels) aufgrund ihrer dynamischen Natur eine Blackbox dar. Zur Analyse dieser Kernels nutzt XPro feingranulares Hardware-Sampling: Auf TPUs (z. B. Ironwood, TPU v6E/v7) können aus einem Fundus von bis zu 60.000 Performance-Countern mehr als 1.100 Zähler pro Durchlauf bei Abtastraten von bis zu 833 MHz erfasst werden. Dies ermöglicht die detaillierte Befehls- und Pipeline-Analyse (z. B. Identifikation von MXU-Latenzen).
In Zusammenarbeit mit NVIDIA wurde XPro um GPU-PM-Sampling (zeitbasierte Erfassung von NVLink-, PCIe- und SM-Aktivitäten mit 10–100 kHz auf Hopper- und Blackwell-Architekturen) sowie Multi-Pass Range Profiling über den Multi-Host HLO Runner erweitert. Letzteres ermöglicht das Replay von Kernels zur exakten Messung von über 6.000 Hardware-Metriken (z. B. auf dem GB100-Chip). Neue Entwicklungen in NVIDIAs CUPTI-Schnittstelle reduzieren Zudem den CPU-Profiling-Overhead durch schemabasierte Aktivitätsdatensätze und verhindern Konflikte durch Multi-Subscriber-Unterstützung.
Wichtigste Punkte & Zeitstempel
0:03 Überblick und Agenda: Prasanna Ranganathan leitet die Session zur Profilierung von JAX- und OpenXLA-Workloads ein. Die Agenda umfasst Custom-Call-Debugging, TPU/GPU-Performance-Counter und Profiling-Features in Kooperation mit NVIDIA.
1:48 Architektur von XPro: XPro dient als herstellerübergreifender Profiler für TPUs, GPUs und PJRT-kompatible Beschleuniger. Er enthält Ansichten wie den Trace Viewer (Gantt-Diagramm), Graph Viewer (HLO-Graphen), Memory Viewer (Puffer-Zuordnungen/Fragmentierung) und Roofline-Analysen.
4:02 Erfassungsmethoden: Profile können interaktiv über das XPro-UI (Lokal oder Remote-Host) oder programmatisch via jax.profiler und torch.profiler gestartet werden, um Metriken wie Step-Time, Token-Durchsatz und Flop-Auslastung zu erfassen.
6:45 Telemetrie für Leistung und Thermik: Für TPU-Architekturen (z. B. Trillium, Ironwood) integriert XPro Echtzeitmessungen von Stromverbrauch, Modultemperaturen und thermisch bedingtem Throttling in einer einheitlichen Ansicht.
9:21 Herausforderungen bei Custom Calls: Palace- und Triton-Kernels werden direkt in Maschinencode abgesenkt und umgehen HLO-Optimierungen. Durch dynamische Kontrollflüsse und Formate erfordern sie feingranulare Befehlstraces anstelle statischer Analysen.
13:40 TPU Hardware-Performance-Counter: TPU-Generationen (v6E, v7) bieten bis zu 60.000 Hardware-Zähler. Bei Abtastraten von bis zu 833 MHz (1 µs) können über 1.100 Zähler pro Durchlauf erfasst werden, um MXU-Auslastungslücken (z. B. Behebung von Stalls durch Erhöhung der Puffer-Tiefe von 1 auf 3) zu analysieren.
22:26 GPU PM-Sampling (NVIDIA): Chris Ashton demonstriert die Integration von NVIDIA CUPTI in XPro. Das PM-Sampling erfasst zeitbasierte Hardware-Metriken (SM-Aktivität, NVLink, PCIe) auf GPU-Architekturen wie Hopper und Blackwell (GB100) bei Frequenzen von 10–100 kHz.
29:52 Range Profiling via HLO Runner: Um tiefere Hardware-Metriken (über 6.000 auf GB100) ohne Zeitabtastungsfehler zu messen, führt der Multi-Host HLO Runner Multi-Pass-Replays (z. B. 6–7 Durchläufe) aus und übergibt präzise Befehls- und Cache-Zähler an XPro.
38:05 CUPTI-Optimierungen und Multi-Subscriber: Neue CUPTI-Funktionen führen schemabasierte Aktivitätsdatensätze zur Reduzierung des Host-CPU-Overheads ein und ermöglichen über Multi-Subscriber-Support den parallelen Zugriff von XPro, JAX PGLE und Nsight Systems ohne Kollisionen.
Glossar
XPro: Ein von Google entwickeltes, plattformübergreifendes Profiling-Werkzeug für Machine-Learning-Workloads zur Performance-, Speicher- und Hardwareanalyse.
JAX: Ein Python-Framework für hochleistungsfähiges numerisches Rechnen und Machine Learning, das automatische Differenzierung und XLA-Kompilierung kombiniert.
OpenXLA / HLO: OpenXLA ist ein domänenspezifischer Compiler für ML-Frameworks; HLO (High-Level Optimizer) ist die Zwischendarstellung (IR) des Compilers zur Graph-Optimierung.
Custom Call (Palace / Triton): Benutzerdefinierte Code-Kerne, die direkt in maschinennahen Code kompiliert werden und Standard-Compiler-Optimierungen umgehen.
PJRT (Plugin Joint Runtime): Eine einheitliche Laufzeitschnittstelle in OpenXLA, die es verschiedenen Hardware-Beschleunigern ermöglicht, sich transparent in ML-Frameworks einzubinden.
Roofline-Modell: Ein anschauliches Leistungsmodell, das aufzeigt, ob die Performance einer Berechnung durch die Speicherbandbreite (memory-bound) oder die Rechenleistung (compute-bound) begrenzt wird.
MXU (Matrix Multiply Unit): Die spezialisierte Hardware-Recheneinheit auf Google TPUs, die für die Ausführung von Matrixmultiplikationen verantwortlich ist.
CUPTI (CUDA Profiling Tools Interface): Eine Software-Schnittstelle von NVIDIA, die externen Profiling-Werkzeugen Zugriff auf interne Leistungsinformationen und Hardware-Zähler von NVIDIA-GPUs gewährt.
PM-Sampling (Performance Monitor Sampling): Das periodische Abtasten von Hardware-Zählern (z. B. Auslastung, Taktrate) während der Laufzeit eines Programms zur Erstellung eines Zeitdiagramms.
Range Profiling: Ein Testverfahren, bei dem ein bestimmter Codeabschnitt mehrfach durchlaufen wird, um tausende detaillierte Hardware-Zähler ohne Genauigkeitsverlust zu erfassen.
Multi-Subscriber Support: Eine Eigenschaft von Profiling-Schnittstellen, die es mehreren Diagnose-Werkzeugen gleichzeitig erlaubt, Daten vom selben Treiber/Chip abzugreifen, ohne sich gegenseitig zu blockieren.