Der Vortrag präsentiert „Jaxator“, ein von AMD entwickeltes Framework, das darauf abzielt, die hochoptimierten AMD-AIT-Operator-Bibliotheken (AIT) nahtlos in JAX/XLA-Workloads auf der ROCm-Plattform für AMD-Instinct-GPUs zu integrieren. Jaxator fungiert als Brücke, die JAX-Arrays über XLA-Custom-Calls mit AIT-Kernen verbindet und dabei Sharding-Regeln, Layouts und benutzerdefinierte VJP-Unterstützung für das Training bereitstellt. Anhand von zwei Fallstudien (Multi-Head Attention und dem MXFP4-End-to-End-Rezept für Llama-Modelle) demonstriert das Projekt erhebliche Leistungssteigerungen (ca. 4,39 % bei MHA, bis zu 30 % bei MXFP4 im Vergleich zu FP8-Baselines), während gleichzeitig Herausforderungen bei der Entkopplung von PyTorch-Abhängigkeiten und der Laufzeitkoordinierung zwischen HIP und XLA adressiert werden.
Wichtigste Punkte & Zeitstempel
0:00 Einführung in Jaxator: Vorstellung von Rutaj Vidya (AMD Software Engineer) bezüglich der Integration von AMD-AIT-Kernen in JAX/XLA für ROCm-Workloads.
0:19 Präsentationsstruktur: Gliederung in AIT-Übersicht, Jaxator-Framework, zwei Fallstudien (MHA und MXFP4) sowie Ausblick.
2:01 Definition von AIT: AIT ist AMDs Hochleistungs-KI-Operator-Bibliothek auf ROCm für Training und Inferenz, die LLM-Bausteine (Attention, GEMMs, MoE, LayerNorm) über mehrere Backends (Composable Kernel, Triton, HIP, Hand-Assembly) bereitstellt.
3:21 Leistungsgewinne von AIT: AIT erzielt je nach Operator Geschwindigkeitssteigerungen von 2x bis zu 17x im Vergleich zu Standardimplementierungen.
4:26 Jaxator-Framework: Fungiert als Integrationsschicht, die Low-Level-Lowering-Regeln, Runtime-Integration, Sharding-Regeln und Custom VJP (Gradienten) für JAX bereitstellt.
6:38 Verwendung von XLA Custom Calls: Architektonische Entscheidung, XLA-Custom-Calls zu nutzen, um die Kompilierungszyklen von XLA von den AIT-Freigabezyklen zu entkoppeln.
10:02 Erste Fallstudie (Multi-Head Attention): MHA stellt einen kritischen Flaschenhals beim Llama-Training dar; die Integration über Jaxator liefert einen mittleren Geschwindigkeitszuwachs von 4,39 %.
12:03 Zweite Fallstudie (MXFP4): Stresstest mit einem 4-Bit-Low-Precision-Fließkomma-Rezept für Llama 3/3.1, das die Quantisierung von BF16 zu MXFP4, Vorwärts- und Rückwärtskerne sowie Sharding-Kollokationen umfasst.
16:33 MXFP4-Leistungsresultate: Erzielung von ca. 20 % Leistungssteigerung für MLP-only und ca. 30 % für das vollständige Rezept (MLP + QKV + O-Projektionen) im Vergleich zur FP8-Baseline, wobei die Verlustkurve nahe an der FP8-Baseline verläuft.
18:09 Lektionen & Herausforderungen: Notwendigkeit der Entkopplung tiefer PyTorch-Abhängigkeiten aus AIT, Beherrschung zweier Runtimes (HIP und XLA) sowie Behandlung von Kernen als Blackboxes für XLA.
21:01 Zukunftsaussichten: Eröffnung von Pull-Requests, Bereitstellung von Alpha-Wheels im GitHub-Repository von ROCm/Jaxator und Ausweitung der Operator-Abdeckung.
Glossar
JAX: Ein von Google entwickeltes Python-Framework für hochleistungsfähiges numerisches Rechnen und maschinelles Lernen mit automatischer Differenzierung.
XLA (Accelerated Linear Algebra): Ein domainspezifischer Compiler für lineare Algebra, der JAX- und TensorFlow-Berechnungen für maximale Ausführungsgeschwindigkeit optimiert.
ROCm (Radeon Open Compute): AMDs offene Softwareplattform für GPU-Computing und parallele Programmierung (das AMD-Gegenstück zu Nvidias CUDA).
AIT: AMDs Hochleistungs-KI-Operator-Bibliothek auf ROCm, die speziell optimierte Rechenkerne für Deep-Learning-Workloads bereitstellt.
Custom Calls: Mechanismen in XLA, mit denen extern geschriebene, plattformspezifische C++- oder GPU-Funktionen direkt aus dem kompilierten Rechengraphen aufgerufen werden können.
GEMM (General Matrix Multiply): Eine grundlegende mathematische Operation (Matrix-Multiplikation), die den Großteil der Rechenzeit beim Training und Betrieb von neuronalen Netzen ausmacht.
MoE (Mixture of Experts): Eine Architektur für neuronale Netze, bei der nur bestimmte Teilnetzwerke („Experten“) für spezifische Eingaben aktiviert werden, um die Recheneffizienz zu steigern.
MHA (Multi-Head Attention): Ein zentraler Baustein von Transformator-Modellen (wie Llama), der es dem Modell ermöglicht, gleichzeitig auf Informationen aus verschiedenen Repräsentationsunterräumen zu achten.
Sharding: Das Aufteilen von Daten oder Modellgewichten auf mehrere GPUs, um paralleles Training extrem großer Modelle zu ermöglichen.
VJP (Vector-Jacobian Product): Eine mathematische Methode zur Implementierung des Rückwärtspfads (Gradientenberechnung) in Frameworks wie JAX für das Training von KI-Modellen.
MXFP4: Ein extrem präzisionsreduziertes 4-Bit-Fließkommaformat (Microscaling Format), das den Speicherbedarf und die Rechenzeit drastisch senkt.
BF16 (Bfloat16): Ein 16-Bit-Fließkommaformat, das häufig im Deep Learning verwendet wird, um den Speicherverbrauch im Vergleich zu Standard-32-Bit-Zahlen zu halbieren.
FP8: Ein 8-Bit-Fließkommaformat für noch höhere Recheneffizienz und Durchsatz bei modernem KI-Training.
HIP: Eine C++-Laufzeitumgebung und Programmiersprache von AMD, die die Ausführung von hocheffizientem Code auf AMD-Grafikkarten ermöglicht.
Composable Kernel (CK): Eine Programmierumgebung und Bibliothek von AMD zum Schreiben hochoptimierter GPU-Rechenkerne.
Triton: Eine Open-Source-Programmiersprache und ein Compiler, mit dem Entwickler effiziente GPU-Rechenkerne in Python schreiben können, ohne Low-Level-Code schreiben zu müssen.
MLPerf: Ein internationaler Industriestandard zur objektiven und standardisierten Messung der Leistung von Hard- und Software für maschinelles Lernen.