Dieses Transkript dokumentiert eine technische Präsentation über Leistungsoptimierung und Debugging für JAX- und OpenXLA-Workloads, vorgetragen von Ingenieuren von Google und NVIDIA. Im Fokus steht das Tool XPro, ein plattformübergreifender ML-Performance-Profiler mit nativer Unterstützung für TPUs und GPUs. Es werden detaillierte Einblicke in fortgeschrittene Funktionen wie Trace-Ansichten, Speicher-Viewer, Roofline-Analysen und die Überwachung von Custom Calls (z. B. via Palace- oder Triton-Kernel) gegeben. Zudem stellen die NVIDIA-Kollaborateure neue PM-Sampling-Funktionen und Range-Profiling-Methoden vor, die eine tiefe Hardware-Metrikenerfassung für Hopper-GPUs und neuere Architekturen ermöglichen.
Wichtigste Punkte & Zeitstempel
0:00 Einführung und Agenda: Präsentation von XPro (OpenXLA Performance Profiler) durch Prasana Reangasami (Google) gemeinsam mit NVIDIA-Partnern zur Leistungs- und Energieanalyse von ML-Workloads.
1:47 Überblick über XPro: Der ML-Performance-Profiler unterstützt nativ TPUs und GPUs sowie andere Beschleuniger über den PJRT-Standard und bietet Ansichten für Trace, Graph, Speicher und Roofline.
4:00 Profile-Erfassung & Konfiguration: Erläuterung der UI-basierten und programmatischen Erfassung via zax.profiler und torch.profiler inklusive erweiterter Konfigurationsmöglichkeiten.
6:01 Visualisierungs-Tools: Detaillierte Darstellung von Gantt-Chart-Zeitlinien, Speicherallokationen (HBM/VM) zur Erkennung von Fragmentierungen und Roofline-Modellen zur Identifizierung von Memory- vs. Compute-Bound-Zuständen.
9:17 Custom Calls & Analyse-Herausforderungen: Diskussion von Custom Calls (Palace- und Triton-Kernel) und den Schwierigkeiten bei der Verfolgung asynchroner DMAs und dynamischer Tensor-Formen.
13:30 TPU Performance-Zähler (PM Sampling): Nutzung von bis zu 60.000 Perf-Stat-Countern auf V6E/V7-TPUs (Ironwood/Trillium) mit Sampling-Raten von bis zu 833 MHz zur instruktionsgenauen Analyse.
24:34 NVIDIA GPU-Integration via CUPTI: Chris Ashton (NVIDIA) präsentiert die Integration des CUDA Profiler Tools Interface (CUPTI) in XPro, einschließlich PM-Sampling für Hopper++ Architekturen.
30:08 Range-Profiling & Multi-Pass-Ausführung: Einführung des Range-Profilings für NVIDIA-GPUs, das durch Kernel-Replay im Multi-Host-HLO-Runner den Zugriff auf über 6.000 Hardware-Metriken ermöglicht.
38:04 Cupy-Optimierungen & Multi-Subscriber: Neue schema-basierte Aktivitätsdatensätze zur Reduzierung von Profiling-Overhead und Unterstützung paralleler Abonnenten (z. B. XPro und NSight Systems gleichzeitig).
40:40 Frage-und-Antwort-Runde: Diskussionen über Zähler-Overhead (keiner bei TPUs, CPU-limitiert bei GPUs) und die Abgrenzung zwischen XPro und NVIDIA-eigenen Tools wie Insight Compute.
Glossar
XPro (OpenXLA Performance Profiler): Ein plattformübergreifendes Analyse- und Profiling-Werkzeug zur Untersuchung von Leistung, Energieverbrauch und thermischem Verhalten von Machine-Learning-Workloads.
OpenXLA: Ein offener Compiler- und Laufzeit-Stack für maschinelles Lernen, der optimierten Maschinencode für verschiedene Hardware-Beschleuniger generiert.
JAX: Ein Python-Framework für Hochleistungs-Numerik und automatisches Differenzieren, das eng mit XLA integriert ist.
PJRT: Eine standardisierte Plugin-Schnittstelle, die Machine-Learning-Frameworks mit unterschiedlichen Hardware-Beschleunigern verbindet.
HLO (High-Level Optimizer): Eine plattformunabhängige Zwischenrepräsentation (Intermediate Representation) im XLA-Compiler zur Durchführung von Optimierungs-Pässen.
HBM (High Bandwidth Memory): Ein extrem schneller, vertikal gestapelter Hauptarbeitsspeicher, der direkt auf modernen GPUs und TPUs verbaut ist.
VM (Vektor- oder Virtueller Speicher): Ein lokaler, schneller Speicherbereich auf TPU-Architekturen.
Roofline-Analyse: Ein visuelles Leistungsmodell zur Bestimmung, ob ein Rechenkernel durch die Rechenleistung (Compute-Bound) oder den Speicherdurchsatz (Memory-Bound) limitiert ist.
Custom Calls: Direkt in den Compiler-Graph eingebettete, handoptimierte Kernel (oft in Triton oder Palace geschrieben), die spezifische Hardware-Funktionen nutzen.
Triton / Palace: Programmiermodelle und Compiler-Dialekte zur Erstellung von Hochleistungs-Beschleunigerkerneln (für GPUs bzw. TPUs).
PM Sampling (Performance Metric Sampling): Die regelmäßige, zeitbasierte Erfassung von Hardware-Zählern (Performance Countern) während der Ausführung eines Programms.
CUPTI (CUDA Profiler Tools Interface): Die offizielle Programmierschnittstelle von NVIDIA zur Erstellung von Profiling- und Trace-Werkzeugen für CUDA-Anwendungen.
Hopper++ / Blackwell (B100): Moderne NVIDIA-GPU-Architekturen mit spezialisierten Tensor-Kernen und erweiterten Leistungsüberwachungsfunktionen.
MXU (Matrix Multiply Unit): Eine spezialisierte Hardware-Recheneinheit auf TPUs zur hocheffizienten Durchführung von Matrixmultiplikationen.