Dieses Video bietet eine detaillierte technische Analyse der Workload-Optimierung für Google TPUs (Tensor Processing Units) mittels des Profiling-Tools XPRO. Ein Google-Produktmanager erläutert die TPU-Architektur, typische Leistungsengpässe (wie das „Blindfliegen“) und demonstriert anhand einer Mistral 8x22-Fallstudie, wie gezielte Anpassungen der Batch-Größe die Rechenleistung maximieren. Zudem werden neue Funktionen wie kontinuierliche Profiling-Snapshots (ein 2-GB-Host-Zirkularpuffer für 90 Sekunden Verlauf) und die native PyTorch-Integration vorgestellt.
Wichtigste Punkte & Zeitstempel
0:00 Einführung in XPRO: Präsentation des TPU-Software-Suiten-Tools zur Vermeidung von „Blindflügen“, die zu ineffizienten Ressourcenkosten und langsamen Iterationszyklen führen.
1:16 TPU-Architektur: Systemweites Gleichgewicht aus Tensor Core (dichte Matrizenrechnungen), Sparse Core (Einbettungssuchen), HBM (High Bandwidth Memory mit 35 TB/s) und direktem ICI-Interconnect.
3:01 Arbeitsablauf der Ausführung: Die Host-CPU steuert die Datenpipeline und kompiliert Code über XLA zu Hardware-Binärdateien, die über den PCIe-Bus an den TPU-Beschleuniger übertragen werden.
5:04 Fallstudie zu Leistungsengpässen: Analyse eines Mistral 8x22-Modells auf Trillium-Hardware; das anfängliche Profil zeigt eine hohe All-Reduce-Kommunikationszeit (31 %) bei geringer FLOP-Auslastung (27 %).
7:36 Speicher-Viewer & Optimierung: Eine Erhöhung der Batch-Größe von 1 auf 8 steigert die Speicherauslastung von 12 GB auf 29 GB (von 32 GB) und hebt die FLOP-Auslastung auf 60,74 %, wodurch das System in einen optimalen, rechengebundenen Zustand wechselt.
9:51 Kernel-Profiling: Tiefenanalyse für benutzerdefinierte Kernel (Palace, Mosaic, Triton) mit Zugriff auf HLO-Graphen, LLO-Bündel und über 16.000 rohe Hardware-Performance-Zähler.
14:36 Kontinuierliche Profiling-Snapshots: Ein permanenter Flight-Recorder auf Basis eines 2-GB-Host-Zirkularpuffers speichert die letzten 90 Sekunden, um transiente Anomalien und unerwartete Leistungsabfälle abzufangen.
15:59 PyTorch-Kompatibilität: Native TPU-Profilierung über Kineto/torch_tpis ohne intrusive Code-Änderungen oder benutzerdefinierte Wrapper.
Glossar
TPU (Tensor Processing Unit): Ein von Google entwickelter, spezialisierter Mikroprozessor, der gezielt für die Beschleunigung von Workloads im Bereich des maschinellen Lernens (KI) optimiert ist.
XLA (Accelerated Linear Algebra): Ein domänenspezifischer Compiler für lineare Algebra, der maschinelle Lernmodelle optimiert, um die Ausführungsgeschwindigkeit auf Hardware zu maximieren.
Tensor Core: Ein Hardware-Kern innerhalb der TPU, der für dichte, rechenintensive Matrixmultiplikationen optimiert ist.
Sparse Core: Ein spezieller Rechenkern der TPU für unregelmäßige, speicherintensive Operationen wie Embedding-Lookups (Einbettungssuchen).
HBM (High Bandwidth Memory): Ein ultraschneller, dreidimensional gestapelter Arbeitsspeicher, der direkt auf dem Prozessorsubstrat sitzt und extrem hohe Datenübertragungsraten (hier 35 TB/s) ermöglicht.
ICI Interconnect: Ein proprietäres Hochgeschwindigkeits-Netzwerkprotokoll, das die direkte Kommunikation zwischen mehreren TPU-Chips ohne Einbindung der Host-CPU erlaubt.
All-Reduce: Eine verteilte Rechenoperation in Cluster-Systemen, bei der Daten (wie Gradienten) von allen Rechenknoten gesammelt, mathematisch zusammengefasst und synchronisiert werden.
HLO (High-Level Optimizer): Eine maschinenunabhängige Zwischenrepräsentation von Berechnungen innerhalb des XLA-Compilers zur Optimierung von Operationen auf hoher Abstraktionsebene.
Kineto / PyTorch Profiler: Ein in PyTorch integriertes Werkzeug zur Messung von Laufzeiten, Speicherverbrauch und Ressourcen-Auslastung von Operationen auf Hardware-Beschleunigern.