Browse Summaries

Abstract

Dieser Vortrag stellt JAX Quantum vor, ein am MIT entwickeltes Open-Source-Framework zur beschleunigten Simulation von Quantenhardware. Die Software nutzt Google JAX, um eine an QTip angelehnte API mit automatischer Differentiation, Batching und Hardware-Beschleunigung auf GPUs und TPUs zu verbinden. Der Schwerpunkt liegt auf der Modellierung von gekoppelten supraleitenden Schaltkreisen, Zeitentwicklungen offener Quantensysteme und bosonischer Quantenfehlerkorrektur.

Der Referent präsentiert Performance-Analysen, die zeigen, dass TPUs bei Matrixmultiplikationen (matmul) und der Lösung von Differentialgleichungen (diffrax) durch ihre Matrix-Multiply-Unit-Architektur (MXU) hervorragend skalieren, während GPUs bei Eigenwertzerlegungen (eigh) dominieren. In der Praxis führt der Einsatz von Hardware-Beschleunigern bei komplexen Multi-Start-Pulsoptimierungen zu einer Laufzeitreduktion von 14 Stunden (CPU) auf 24 Minuten (TPU v4, 35x schneller) bzw. 3,5 Minuten (NVIDIA H200 GPU, 240x schneller).

Wichtigste Punkte & Zeitstempel

• 0:00 Historischer Kontext: Matrixmechanik als fundamentale mathematische Sprache der Quantenmechanik seit Heisenberg (1925).
• 0:15 Experimentelle Rahmenbedingungen: Vorstellung der Arbeitsumgebung am MIT mit Mischungskältemaschinen, die Quantenprozessoren (QPUs) auf Temperaturen kühlen, die 300-mal kälter als der interstellare Raum sind.
• 1:39 Mathematische Operationstypen: Klassifizierung der Hardware-Simulationen in Matrix-Vektor-Multiplikationen, Kraus-Maps (Matrix-Matrix-Multiplikation), Hermitische Eigenwertprobleme (eigh) und Zeitentwicklungen mittels Schrödinger-Gleichung via diffrax.
• 2:47 Architektur von JAX Quantum: Schichtenmodell aus QTip-kompatibler High-Level-API, Modulen für supraleitende Schaltkreise und Floquet-Theorie sowie JAX-Backend-Features (Autodiff, TPU/GPU-Parallelisierung).
• 4:02 Die QArray-Datenstruktur: Kernobjekt als Hülle um JAX-Arrays zur Kapselung von Moden-, Pegel- und Batch-Informationen. Unterstützt dichte Matrizen, BCOO-Sparse, diagonale Formate, CuQuantum-Integration und Sharding.
• 4:53 Hardware-Benchmarks: Leistungsevaluierung von TPU v4, NVIDIA H200 GPU und Laptop-CPU. TPUs zeigen Stärken bei matmul mit Stufenartefakten bei Vielfachen der 128er MXU-Zellgröße. GPUs führen bei eigh.
• 6:20 Rauschanalyse & Gerätedesign: Nutzung automatischer Differentiation zur effizienten Berechnung von Frequenzgradienten in Abhängigkeit von Fluss-Bias-Fluktuationen.
• 6:53 Pulsoptimierung für Bosonische Fehlerkorrektur: Multi-Start-Optimierung zur Vorbereitung komplexer Quantenzustände. Zeiteinsparung von 14 Stunden (Laptop-CPU) auf 24 Minuten (TPU v4) und 3,5 Minuten (H200 GPU).
• 7:44 Zusammenfassung & Roadmap: Geplante Erweiterungen umfassen verteiltes Sharding, Evaluierung von TPU v5/v6, erweiterte Backends und Krylow-Unterraum-Methoden zur Optimierung von Diagonalisierungen auf TPUs.

Diskussions-Highlights

• 8:27 Hardware-Vergleichbarkeit: Diskussion über den Vergleich von TPU v4 mit NVIDIA H200 GPUs. Der Referent verweist auf den zeitnahen Zugang zu neueren TPU-Infrastrukturen (z. B. TPU v6e via Google Colab).
• 9:13 Framework-Wahl (Python vs. Julia): Die Entscheidung gegen QuantumOptics.jl fiel trotz der Performance von Julia zugunsten des Python/JAX-Ökosystems, um die Akzeptanz in der Experimentalphysik-Community zu maximieren.
• 9:31 TPU-Laufzeitverhalten bei Differentialgleichungen: Evaluierung adaptiver versus fester Schrittweiten in diffrax. Adaptivität erschwert die Auslastung von TPUs, liefert aber dennoch signifikante Geschwindigkeitsvorteile.
• 10:27 Reinforcement Learning & Hardware-Latenz: Beim Einsatz von RL zur Quantensteuerung bildet die Übertragungs-Latenz zwischen Arbitrary Waveform Generatoren (AWGs) und Beschleunigern den Engpass, welcher durch optische Direktverbindungen minimiert wird.

Glossar

• JAX: Ein von Google entwickeltes Python-Framework für hochleistungsfähige numerische Berechnungen. Es bietet automatische Differentiation (Autodiff) und kompiliert Code mittels XLA direkt für GPUs und TPUs.
• QPU (Quantum Processing Unit): Der physikalische Quantenprozessor-Chip, auf dem Qubits oder bosonische Moden manipuliert werden.
• Mischungskältemaschine (Dilution Refrigerator): Ein kryogenes Kühlsystem, das durch das Mischen von Helium-3- und Helium-4-Isotopen Temperaturen im Millikelvin-Bereich erzielt.
• Kraus-Map: Eine mathematische Operatordarstellung zur Beschreibung der Zeitentwicklung offener Quantensysteme unter Einwirkung von Umweltrauschen.
• Diffrax: Eine auf JAX basierende Open-Source-Bibliothek zur numerischen Lösung von gewöhnlichen und stochastischen Differentialgleichungen.
• QTip: Eine etablierte Python-Bibliothek (Quantum Toolbox in Python) zur Simulation der Dynamik von Quantensystemen.
• Floquet-Theorie: Ein mathematisches Verfahren zur Behandlung von Differentialgleichungen mit periodisch zeitabhängigen Koeffizienten, häufig genutzt bei wechselfeldgetriebenen Quantensystemen.
• QArray: Das zentrale Datenstrukturelement in JAX Quantum, das JAX-Tensor-Arrays mit quantenmechanischen Metadaten (z. B. Hilbertraum-Dimensionen) verknüpft.
• BCOO (Bounded Coordinate Format): Ein Speicherformat für dünnbesetzte (sparse) Matrizen, das speziell für effiziente Ausführung auf Parallelbeschleunigern (GPUs/TPUs) strukturiert ist.
• CuQuantum: Ein von NVIDIA entwickeltes Software Development Kit (SDK) zur Beschleunigung von Quantenschaltkreis- und Zustandsvektorsimulationen auf NVIDIA-GPUs.
• MXU (Matrix Multiply Unit): Die spezialisierte Hardware-Recheneinheit in Google TPUs, die für die Ausführung dichter Matrixmultiplikationen optimiert ist.
• Hermitische Matrix (eigh): Eine komplexe quadratische Matrix, die gleich ihrer konjugiert-transponierten Matrix ist. Ihre Eigenwerte repräsentieren in der Physik reelle Messgrößen wie Energieniveaus.
• Bosonische Fehlerkorrektur: Ein Quantenfehlerkorrektur-Ansatz, der unendlichdimensionale Hilberträume kontinuierlicher Variablen (z. B. Photonen in Resonatoren) nutzt, anstatt diskrete Zwei-Niveau-Qubits zu verwenden.
• AWG (Arbitrary Waveform Generator): Ein Synthesizer zur Erzeugung präziser analoger Mikrowellenpulse zur Steuerung physikalischer Qubits.

Abstract

Dieser Vortrag bietet eine tiefgehende architektonische Analyse von XLA (Accelerated Linear Algebra), dem spezialisierten Compiler-Ökosystem für maschinelles Lernen von Google. Die Referenten Kevin Gleason und Peter Gavin erläutern die historische Entwicklung und das Schichtenmodell der Zwischendarstellungen (Intermediate Representations, IR) innerhalb von XLA: ausgehend von der ursprünglichen HLO (High Level Optimizer, ca. 2017) über MHLO (Meta-HLO, 2019) bis hin zu StableHLO (2022) als offenem Industriestandard für das OpenXLA-Ökosystem. Ergänzt wird dieses IR-Spektrum durch VHLO zur Gewährleistung mehrjähriger Bytecode-Kompatibilität sowie CHLO zur Bewahrung hochstufiger mathematischer Framework-Abstraktionen.

Darüber hinaus werden die Laufzeit- und Compilerschnittstelle PJRT sowie die interne Funktionsweise des XLA-Compilers detailliert aufgeschlüsselt. Der Compiler verarbeitet funktionale, datenflussbasierte Instruktionsgraphen („Computations“) ohne Rekursion, was eine rein statische Speicherallokation zur Kompilierzeit erzwingt. Besondere Aufmerksamkeit gilt den Speicher-Layouts (Minor-to-Major-Indizierung, mehrdimensionales Tiling zur Vektorisierung) sowie der Kernel-Fusion (insbesondere Loop Fusion), die als primärer Leistungshebel zur Eliminierung physischer Speicherzugriffe dient. Abschließend adressieren die Referenten Aspekte der Schleifenentrollung („Unrolling“), der Ausführung von Custom Calls sowie Steuerungsmechanismen für die numerische Präzision.

Wichtigste Punkte & Zeitstempel

• 00:05 Einführung & Agenda: Übersicht über das Schichtenmodell von XLA, HLO-Dialekte, Framework-Kopplung, die PJRT-Laufzeitschnittstelle sowie die interne Compiler-Architektur für TPUs und andere Beschleuniger.
• 00:54 HLO-Entstehung & Grundkonzept: HLO wurde 2017 mit etwa 120 Kernel-Ops als framework-unabhängige und hardware-portable Compiler-IR eingeführt, um TensorFlow und zukünftige Zielarchitekturen (wie TPUs) zu entkoppeln.
• 02:08 Einführung von MHLO: 2019 entstand MHLO als MLIR-Repräsentation von HLO, um das MLIR-Ökosystem für dynamische Formen, Quantisierungstypen und compilerspezifische Fusions-Passes zu nutzen.
• 03:31 StableHLO als OpenXLA-Standard: StableHLO etablierte 2022 eine reduzierte, hochspezifizierte Schnittstellen-IR (~100 Ops) für OpenXLA, die von Frameworks (JAX, PyTorch, ZML) und Hardware-Backends (AWS Neuron, Tenstorrent, Apple Metal) unterstützt wird.
• 05:22 VHLO & Kompatibilität: Versioned StableHLO (VHLO) nutzt den MLIR-Bytecode sowie automatisierte Pass-Transformationen (Up-/Downgrades) zur Aufrechterhaltung mehrjähriger Vorwärts- und Rückwärtskompatibilität.
• 06:35 CHLO für High-Level-Abstraktionen: Client-HLO (CHLO) bewahrt kompositionsrelevante Operationen höherer Ordnung (z. B. top_k, scan), sichert vorhersagbare numerische Zerlegungen und wird optional in StableHLO oder Backend-IRs abgesenkt.
• 08:21 Custom Kernels & Custom Calls: Maßgeschneiderte Kernels (z. B. JAX Palace) werden als opake MLIR-Blobs in stablehlo.custom_call-Instruktionen verpackt und von spezialisierten Backend-Compilern ausgewertet.
• 09:21 PJRT-Laufzeitschnittstelle: Standardisierte API-Schicht für Beschleuniger-Plugins zur Abstraktion von Kompilierung, Programmausführung und Host-Device-Pufferverwalter.
• 10:31 HLO-Modulstruktur & Rekursionsverbot: Das HLO-Modul verwaltet datenflussbasierte Instruktionsgraphen. Rekursion ist strikt untersagt, um obere Grenzen für den Speicherverbrauch zu garantieren.
• 12:05 Speicher-Layouts, Minor-to-Major & Tiling: Arrays unterstützen bis zu 8 Dimensionen. Die physische Speicheranordnung wird über Minor-to-Major-Layouts und mehrdimensionales Tiling (z. B. Kachelung T8 für Vektorregister) gesteuert.
• 15:12 HLO-Instruktionen & Steuerfluss: Neben Arithmetik- und Array-Ops sind auch Steuerflusselemente (conditional, while), Funktionstransfers, Parameter und Tupel als explizite HLO-Instruktionen repräsentiert.
• 17:43 Fusion-Passes zur Speicheroptimierung: Die Zusammenfassung von Sub-Graphen (z. B. Loop Fusion) reduziert DRAM-Lese- und Schreibzugriffe dramatisch, indem Zwischenergebnisse direkt über Vektorregister weitergeleitet werden.
• 19:49 Compiler-Architektur (Frontend vs. Backend): Das Frontend führt geräteunabhängige Transformationen (CSE, algebraische Vereinfachungen, Layout-Wahl) durch; das Backend erzeugt den gerätespezifischen Maschinencode und steuert das asynchrone Scheduling.
• 24:00 Custom Ops, Statische Speicherallokation & Kompilierzeiten: Custom Ops agieren meist als Blackboxes für Fusions-Passes. Das Fehlen dynamischer Allokation erzwingt Compile-Zeit-Allokationen und lange Schleifenentrollungen (z. B. bei jax.scan).
• 27:14 Steuerung der numerischen Präzision: Diskussion von Genauigkeitsunterschieden bei der Emulation niedriger Bitbreiten (z. B. BF16 via F32) sowie Verweis auf Konfigurations-Flags (XLA_FLAGS, Präzisionsprofile für Matrixmultiplikationen).

Glossar

• XLA (Accelerated Linear Algebra): Ein domänenspezifischer Graph-Compiler für lineares Rechnen, der Machine-Learning-Modelle für TPUs, GPUs und CPUs optimiert und kompiliert.
• HLO (High Level Optimizer): Die primäre, datenflussbasierte Zwischendarstellung (IR) innerhalb des XLA-Compilers.
• MHLO (Meta-HLO): Eine auf dem MLIR-Compiler-Framework basierende Dialekt-Repräsentation von HLO.
• StableHLO: Ein herstellerübergreifender, stabiler IR-Standard auf Basis von HLO mit festgelegter Spezifikation und Abwärts-/Vorwärtskompatibilitätsgarantien.
• VHLO (Versioned StableHLO): Eine versionierte Variante von StableHLO innerhalb von MLIR, die automatische In-Memory-Upgrades und -Downgrades von IR-Instruktionen durchführt.
• CHLO (Client HLO): Ein High-Level-Dialekt zur Erfassung komplexer Framework-Abstraktionen (z. B. komplexe Zahlenrechnung), der vor der Backend-Kompilierung in StableHLO dekomponiert wird.
• PJRT (Parallel Just-In-Time Runtime): Die einheitliche C-Schnittstelle zwischen Machine-Learning-Frameworks (wie JAX oder PyTorch) und Hardware-Backends zur Steuerung von Kompilierung und Ausführung.
• MLIR (Multi-Level Intermediate Representation): Ein von LLVM initiiertes Compiler-Framework zur Erstellung modularer, wiederverwendbarer Zwischendarstellungen (Dialekte).
• Tiling (Kachelung): Eine Optimierungstechnik für Speicher-Layouts, bei der mehrdimensionale Arrays in kachelförmige Blöcke unterteilt werden, um die Vektoreinheiten der Hardware optimal auszulasten.
• Minor-to-Major Layout: Speicherausrichtungsformat in XLA, das die physikalische Anordnung von Array-Dimensionen von der schnellstvariierenden (Minor) zur langsamstvariierenden (Major) Dimension beschreibt.
• Loop Fusion: Eine Compiler-Transformation, die mehrere aufeinanderfolgende elementweise Operationen in einer einzigen Schleife zusammenfasst, um Daten im Register zu halten und Speicherbandbreite zu sparen.
• Custom Call: Eine spezielle HLO-Instruktion, die die Einbettung opaker, externer Code-Blobs oder herstellerspezifischer Hardware-Kernels ermöglicht.
• Token: Ein spezieller HLO-Datentyp zur Modellierung sequenzieller Abhängigkeiten und Nebenwirkungen (wie I/O-Zugriffe) in einem ansonsten rein funktionalen Datenflussgraphen.

Abstract

In dieser Präsentation stellt Arthi vom AMD JAX-Team die Integration von FlyDSL-Kerneln in das JAX-Framework auf AMD-GPUs vor. FlyDSL ist eine Python-basierte Kernel-DSL (Domain-Specific Language), die Entwicklern feingranulare Kontrolle über GPU-Hardwarestrukturen wie Kachelgrößen (Tiling), Speicherlayouts, MFMA-Instruktionen (Matrix Fused Multiply-Add) und Epilog-Fusionen bietet. Dies geht über die automatischen Optimierungsmöglichkeiten des XLA-Compilers hinaus, der bei komplexen Workloads wie Fused Attention oder Mixture-of-Experts (MoE) an seine Grenzen stößt.

Die vorgestellte Architektur nutzt eine Brückenschicht, die JAX-Tensoren via DLPack in memref-Formate konvertiert, XLA Custom Calls registriert und zur Laufzeit einen C-Trampoline-Layer verwendet, um Puffer an GPU-Shards zu übergeben. Am Beispiel eines zwei-stufigen gerouteten GEMM-Kernels für MoE-Modelle (Mixtral 8x7B in FP8) zeigt der Ansatz deutliche Effizienzgewinne: Durch die Fusion von Gather-, Group-GEMM-, SwiGLU- und Scatter-Operationen wird unnötiger Rechenaufwand vermieden. In End-to-End-Inferenz-Testläufen auf MaxText wurde eine 2,5-fache Beschleunigung im Prefill und eine 3,5-fache Beschleunigung im Decode erzielt (Durchsatzsteigerung von 14 auf 50 Tokens/Sekunde). Zukünftige Erweiterungen umfassen die Unterstützung von Rückwärts-Kerneln (Custom VJP) für das Training, dynamische Skalare sowie vorgefertigte Python-Wheels.

Wichtigste Punkte & Zeitstempel

• 0:05 Motivation und Grenzen von XLA: XLA erzeugt zwar effizienten Code für elementweise und Reduktions-Operationen (z. B. Triton-Softmax), stößt jedoch bei leistungskritischen Workloads (MoE GEMMs, Fused Attention) an Grenzen, die explizite Kontrolle über Tiling, Layouts und MFMA-Instruktions-Scheduling erfordern.
• 1:38 FlyDSL-Architektur: FlyDSL bietet ein Python-basiertes „Layout-first“-Programmiermodell (ähnlich CuTe) mit Kontrolle über Block-, Wavelet-, Thread- und Atom/MFMA-Ebenen. Der Kompilierungspfad verläuft von Python über den Fly-Dialekt und AMDOGDL bis zum HSA Code Object.
• 2:41 Kernel-Autorisierung und JAX-Einbindung: Trennung in GPU-Ausführung (@flyc.kernel) und Grid/Block-Launcher (@flyc.jit). Der Wrapper from_jax konvertiert JAX-Tensoren via DLPack in speicherkompatible memref-Strukturen für FlyDSL.
• 4:07 Parameterisierung der fly_call-Schnittstelle: Bietet Übergabe von Eingaben, Launcher-Kernel, Ausgabespezifikationen (Form/Dtype für XLA-Allokation) und statischen Kompilierzeit-Skalaren (aktuell noch keine dynamischen Laufzeit-Skalare).
• 5:05 Tracing-, Kompilierungs- und Laufzeitpfad: Bei der Tracing-Phase wird der Kernel in Fly-IR überführt, als HSA-Code-Objekt kompiliert und als XLA Custom Call registriert. Zur Laufzeit entpackt ein C-Trampolin-Layer Puffer, Skalare und Streams je GPU-Shard.
• 5:56 Optimierung gerouteter MoE-GEMMs: MoE-Workloads beinhalten dynamische Expertengrößen. Der AMD-Standardpfad nutzt rechenintensive maskierte dichte GEMMs; FlyDSL eliminiert Overhead durch die Fusion von Gather-, Group-GEMM-, SwiGLU- und Scatter-Operationen in einen einzigen Kernel.
• 7:31 Standalone-Benchmarks: Die Validierung des FlyDSL-MoE-Kernels gegen die Composable Kernel (CK)-Baseline (FP8, Mixtral 8x7B) belegt Leistungsparität sowie leichte Vorteile in ausgewählten Testfällen.
• 8:14 End-to-End-Integration in MaxText: Einbindung des gerouteten GEMMs in das Mixtral-Inferenzmodell (Top-2-Routing, 16 Shards) inklusive einmaligem Load-Time-Preshuffle; erzielt vollständige numerische Parität mit HuggingFace-Referenzgewichten.
• 9:14 End-to-End-Performance-Ergebnisse: Erreichen einer 2,5-fachen Beschleunigung im Prefill und einer 3,5-fachen Beschleunigung im Decode. Der Decode-Durchsatz steigt von 14 auf 50 Tokens/Sekunde (insbesondere da GFX950/BF16 ohne Group-GEMM-Lowering sonst auf rocBLAS zurückfällt).
• 10:07 Ausblick und Roadmap: Zukünftige Arbeitsschwerpunkte umfassen Custom VJP für Trainingsunterstützung (Rückwärts-Pass), dynamische Laufzeitskalare sowie die Veröffentlichung vorgefertigter Python-Wheels.

Diskussions-Highlights

• 11:31 Fusionskontrolle in FlyDSL: Kernel-Fusionen werden manuell vom Entwickler im Python-Code geschrieben und definiert; der FlyDSL-Compiler führt keine automatischen Fusionsoptimierungen durch.

Glossar

• FlyDSL: Eine von AMD entwickelte domänenspezifische Sprache (DSL) in Python zur Erstellung hochoptimierter GPU-Kernel für AMD-Hardware.
• JAX: Ein von Google entwickeltes Python-Framework für automatisches Differenzieren und hochleistungsfähiges numerisches Rechnen auf GPUs/TPUs.
• XLA (Accelerated Linear Algebra): Ein domänenspezifischer Compiler für maschinelles Lernen, der JAX- oder TensorFlow-Graphen in optimierten Maschinencode übersetzt.
• MFMA (Matrix Fused Multiply-Add): Hardware-Spezialinstruktionen und Recheneinheiten auf AMD-CDNA-GPUs zur beschleunigten Ausführung von Matrixmultiplikationen.
• GEMM (General Matrix Multiply): Die grundlegende mathematische Operation der allgemeinen Matrixmultiplikation in Deep-Learning-Workloads.
• MoE (Mixture of Experts): Eine Modellarchitektur im maschinellen Lernen, bei der Eingaben (Tokens) dynamisch an spezialisierte Teilnetzwerke („Experten“) geroutet werden.
• DLPack: Ein sprach- und framework-übergreifender In-Memory-Tensor-Transfer-Standard zur speichereffizienten Übergabe von Daten zwischen Frameworks ohne Kopieren.
• HSA (Heterogeneous System Architecture) Code Object: Das binäre Ausführungsformat für Kernel, die auf AMD-ROCm-Laufzeitumgebungen ausgeführt werden.
• SwiGLU: Eine mathematische Aktivierungsfunktion (Kombination aus Swish und Gated Linear Unit), die häufig in modernen Large Language Models (LLMs) verwendet wird.
• Group GEMM: Eine Matrixmultiplikationsvariante, die mehrere unabhängige Matrixoperationen unterschiedlicher Dimensionen effizient in einem einzigen GPU-Aufruf gebündelt ausführt.
• CK (Composable Kernel): AMDs C++-basiertes Performance-Framework zur Erstellung hochoptimierter GPGPU-Kernel.
• GFX950: Die interne GPU-Architekturbezeichnung für AMD Instinct-Beschleuniger der neuesten Generation.

Abstract

Dieser Vortrag von Bavani (Software Engineer bei Intel und OpenXLA-Contributor) behandelt die Integration des Intel oneAPI-Backends in XLA, um die Ausführung von JAX-Workloads auf Intel-GPUs zu ermöglichen. Der Schwerpunkt liegt auf der Hardware-Architektur der Intel Arc Pro B70 (Xe2-Architektur) sowie auf dem Zusammenspiel des oneAPI-Software-Stacks (1DNN, 1MKL, 1CCL) mit der Kompilierungs- und Laufzeitpipeline von XLA. Es wird erläutert, wie HLO-Operationen über spezialisierte Pfade – 1DNN für Matrixmultiplikationen und Faltungen, MLIR/SPIR-V für Elementwise- und Reduktions-Ops sowie eine dreistufige Triton-CodeGen-Pipeline – abgesenkt werden. Abschließend werden die Multi-GPU-Skalierung über 1CCL, das PJRT-basierte JAX-Plugin-System und der aktuelle Status im Upstream-Repository samt zukünftiger Roadmap (u. a. SYCL Graph Command Buffers und Crescent Island Support) vorgestellt.

Wichtigste Punkte & Zeitstempel

• 0:00 oneAPI-XLA-Integration: Überblick über die Einbindung des oneAPI-Backends in XLA zur Ausführung von JAX auf Intel-GPUs.
• 0:56 Battlemage-GPU-Architektur: Die Workstation-GPU Arc Pro B70 basiert auf der Xe2-Architektur, bietet 32 GB GDDR6-VRAM für lokale Inferenz großer KI-Modelle und unterstützt Multi-GPU-Deployments unter Linux.
• 1:40 Xe2-Core-Spezifikationen: Ein Xe-Stack kombiniert 32 Xe-Cores; jeder Core besitzt 8 Vektor-Engines (ALUs), 8 Matrix-Engines (XMX) sowie 256 KB dynamisch aufteilbaren L1- und Shared Local Memory.
• 2:26 oneAPI-Software-Stack: Bietet die Voraussetzungen für das JAX-Plugin über plattformübergreifende Bibliotheken: 1MKL (BLAS/FFT), 1DNN (Deep-Learning-Basisfunktionen) und 1CCL (Kollektive Kommunikation).
• 4:07 HLO-Kompilierungspipeline: HLO-Faltungen und Matrixmultiplikationen (Dots) laufen über 1DNN oder optional Triton-Code-Generierung, während Elementwise- und Reduktions-Ops die MLIR-Pipeline durchlaufen.
• 4:46 1DNN-Absenkung & Epilog-Fusion: Der GEMM-Rewriter wandelt Dots in Custom Calls um; der Thunk-Emitter erzeugt Thunks, die zur Laufzeit von 1DNN als gefügte SYCL-Kernels mit Epilog-Post-Ops JIT-kompiliert werden.
• 6:26 MLIR & SPIR-V-CodeGen: Fusionierte Knoten werden in MLIR-Dialekte (XLA, SCF, Arith) übersetzt, optimiert, über LLVM in SPIR-V-Binärdateien konvertiert und mittels Level Zero / Intel Graphics Compiler (IGC) in nativen ISA-Code übersetzt.
• 8:32 Triton-CodeGen-Pipeline: Dreistufiger Ablauf vom geräteagnostischen Triton-IR über den hardwarenahen TTIG-Dialekt (Triton Intel GPU) und Triton GenOps bis hin zur SPIR-V-Generierung.
• 9:58 Triton-Autotuning: Ein GEMM-Fusion-Pass entscheidet per Autotuner zwischen 1DNN und Triton und wählt durch Tiling-Benchmarking zur Kompilierzeit die optimale Konfiguration aus.
• 10:57 SYCL-Laufzeitausführung: Der SYCL-Executor verwaltet Kernel- sowie Speicherlebenszyklen, lädt SPIR-V-Binärdateien in Level-Zero-Module und stellt SYCL-Ressourcen für 1DNN-Thunks bereit.
• 12:09 Multi-GPU-Ausführung via 1CCL: 1CCL-Communicators verwalten Device-Ranks und führen GPU-Kollektivoperationen (z. B. AllReduce, Broadcast) durch direkte Einbindung der SYCL-Queue aus dem XLA-Stream aus.
• 13:21 JAX PJRT-Plugin & Modell-Validierung: Die Anbindung erfolgt über ein dreiteiliges Wheel-Set (jaxlib, jax-oneapi-plugin, jax-oneapi-pjrt); erfolgreich getestete Modelle umfassen MaxText (Llama 2 2B), MaxDiffusion (SDXL), BERT, ViT und Whisper Tiny.
• 14:15 Upstream-Status & Roadmap: Upstream-XLA unterstützt bereits StreamExecutor und SPIR-V-Codegen; in Entwicklung befinden sich PTI-Profiling-Tools, Scaleout-Multi-GPU, Low-Precision-Datentypen, SYCL Graph Command Buffers und Crescent Island Support.

Glossar

• XLA (Accelerated Linear Algebra): Ein spezialisierter Domänen-Compiler für Machine Learning (u. a. für TensorFlow und JAX), der mathematische Graphen in hochoptimierten Maschinencode für Beschleuniger übersetzt.
• JAX: Ein von Google entwickeltes Python-Framework für wissenschaftliches Rechnen und Machine Learning mit automatischer Differenzierung und XLA-Kompilierung.
• oneAPI: Ein von Intel initiiertes, offenes und plattformübergreifendes Programmiermodell, das eine einheitliche Entwicklung über CPUs, GPUs und FPGAs hinweg ermöglicht.
• SYCL: Ein herstellerunabhängiger C++-Standard für heterogene Systeme, auf dem der oneAPI-Laufzeitstack aufbaut.
• 1DNN (oneAPI Deep Neural Network Library): Eine hochoptimierte Bibliothek von Intel, die grundlegende Bausteine für neuronale Netze (z. B. Faltungen, Aktivierungsfunktionen, Matrixmultiplikationen) bereitstellt.
• 1MKL (oneAPI Math Kernel Library): Eine Bibliothek mit mathematischen Routinen für lineare Algebra, Fast Fourier Transformations (FFT) und Zufallsgeneratoren.
• 1CCL (oneAPI Collective Communications Library): Eine Bibliothek für verteilte Trainingsinfrastrukturen, die primitive Kommunikationsmuster (z. B. AllReduce) über mehrere GPUs hinweg beschleunigt.
• MLIR (Multi-Level Intermediate Representation): Ein flexibles Compiler-Framework unter dem LLVM-Schirm, das Programmcode schrittweise über verschiedene Abstraktionsebenen (Dialekte) transformiert.
• SPIR-V: Ein plattformunabhängiges binäres Zwischenformat für Grafik- und Rechen-Kernels, das von Treibern direkt in nativen GPU-Code kompiliert werden kann.
• Level Zero: Eine hardwarenahe, extrem leichtgewichtige Treiberschnittstelle von Intel zur Steuerung von GPU-Ressourcen und zur JIT-Kompilierung.
• PJRT (Plugin Joint Runtime): Die standardisierte Laufzeitschnittstelle in XLA, über die verschiedene Hardware-Backends nahtlos an JAX angebunden werden.
• Triton: Eine Open-Source-Programmiersprache und ein Compiler, mit denen hochleistungsfähige GPU-Kernels (insbesondere für Matrixmultiplikationen) vereinfacht geschrieben werden können.
• TTIG (Triton Intel GPU Dialect): Eine spezielle MLIR-Zwischenrepräsentation für Intel-Hardware, die Matrix-Engines und Subgruppenstrukturen im Triton-Compiler abbildet.
• Xe2-Architektur / XMX (Xe Matrix eXtensions): Intels GPU-Architektur der zweiten Generation mit dedizierten Hardware-Beschleunigern (Matrix-Engines) für KI-Berechnungen.

Abstract

Dieser Vortrag von Rutaj Vidya (Software Engineer bei AMD) stellt Jaxator vor, ein Integrations-Framework, das hochoptimierte AMD-AITer-Kernel in den JAX/XLA-Softwarestack für die AMD ROCm-Plattform einbettet. AITer ist AMDs KI-Operator-Bibliothek, die optimierte und gefügte Kernel (z. B. Attention, GEMM, MoE) über verschiedene Backends wie Composable Kernel (CK), Triton, HIP und Hand-Assembly bereitstellt. Da AITer historisch eng an PyTorch gebunden war, dient Jaxator als entkoppelte Brücke, die JAX-APIs über XLA Custom Calls anbindet. Jaxator kümmert sich um Lowering-Regeln, Custom VJP für automatische Differentiation, Sharding-Regeln für verteilte Multi-GPU-Systeme und die HIP-Laufzeitintegration. Die Leistungsfähigkeit wird anhand zweier Fallstudien demonstriert: Die Integration von Multi-Head Attention (MHA) erzielte einen medianen Geschwindigkeitsgewinn von 4,39 %, während ein experimentelles End-to-End MXFP4-Trainingsrezept (4-Bit-Fließkomma-Quantisierung) in MaxText auf LLaMA 3.0/3.1 bis zu 30 % Durchsatzsteigerung gegenüber einer FP8-Baseline erreichte. Jaxator wird als Open-Source-Wheel bereitgestellt, um native, hochperformante KI-Trainings- und Inferenz-Pipelines auf ROCm zu ermöglichen.

Wichtigste Punkte & Zeitstempel

• 00:05 Zielsetzung von Jaxator: Anbindung von AMD AITer-Kerneln an den JAX/XLA-Workload auf ROCm zur Bereitstellung performanter End-to-End-Pipelines inklusive Rückwärtspfad-Unterstützung und verteiltem Training.
• 02:07 Architektur von AMD AITer: High-Performance-Operator-Bibliothek für ROCm mit optimierten, gefügten Kerneln (Attention, GEMM, MoE, Norms) über Backends wie CK, Triton, HIP und Hand-Assembly, primär mit PyTorch-Abhängigkeiten entwickelt.
• 03:21 Leistungsgewinne durch AITer: Performance-Steigerungen zwischen 2x und 17x auf Operator-Ebene sowie eine Verdopplung des Durchsatzes bei Realwelt-Modellen wie DeepSeek via SGLang.
• 04:26 Funktionsweise des Jaxator-Frameworks: Schnittstelle zwischen JAX-APIs und XLA Custom Calls zur Integration hardwarenaher Kernel mit korrekten Lowering-Regeln, Layouts und Gradienten.
• 06:36 Benutzererfahrung & Integration: Bereitstellung als installierbares Python-Wheel; direkte Nutzung mit JAX-Arrays, Kompatibilität mit JIT-Kompilierung und custom VJP für Differenzierbarkeit.
• 08:00 Niedrigschwellige Systemarchitektur: XLA behandelt den Kernel als opaque Black-Box; Jaxator liefert explizite Custom Calls, VJP-Regeln für das Training sowie Sharding-Regeln für Multi-GPU-Kommunikation.
• 10:04 Fallstudie 1: Multi-Head Attention (MHA): Fokussierung auf MHA als zentralen Rechen-Hotspot beim LLM-Training; Erzielung eines medianen Performance-Gewinns von 4,39 % durch entkoppelte C++-APIs.
• 12:03 Fallstudie 2: MXFP4 End-to-End-Stress-Test: Validierung eines vollständigen 4-Bit-Fließkomma-Trainingsrezepts (MXFP4) in MaxText auf LLaMA 3.0/3.1 inklusive Quantisierung von BF16, Hadamard-Transformation und Stochastic Rounding.
• 13:40 Kommunikation & Sharding bei MXFP4: Notwendigkeit expliziter Kollektiv-Kommunikationsregeln während des Weight-Gradient-Passes (DB), da die Kontraktionsachse M über mehrere GPUs geshardet wird.
• 16:32 MXFP4-Performance-Ergebnisse: Erreichen von 20 % Durchsatzsteigerung bei reinen MLP-Schichten und 30 % bei voller Integration (MLP + QKV/O-Projektionen) gegenüber einer FP8-Baseline bei vergleichbarer Konvergenz.
• 18:09 Erkenntnisse & Zukünftige Entwicklungen: Entkopplung weiterer PyTorch-Abhängigkeiten in reine C++-APIs, Bereitstellung neuer Rezepte, MLPerf-Konvergenzvalidierung und Upstreaming der Quantisierungsmodi in MaxText.
• 23:37 Q&A - Framework-Entkopplung & Backend-Auswahl: Erläuterung der Entfernung von PyTorch-Abhängigkeiten, Funktionsweise des gefügten Quantisierungs-Kernels und automatische Selektion der optimalen AITer-Backends (z. B. Assembly vs. CK) basierend auf Tensor-Shapes.

Glossar

• AITer / AIT: AMD AI Operator Library; eine hochoptimierte Bibliothek von KI-Operatoren für die AMD ROCm-Plattform.
• ROCm: Radeon Open Compute platform; AMDs Open-Source-Software-Stack für GPU-Computing und parallele Programmierung.
• JAX / XLA: JAX ist ein Framework für transformierbaren Python-/NumPy-Code; XLA (Accelerated Linear Algebra) ist der zugrundeliegende Domänen-spezifische Compiler für Machine Learning.
• Custom Call: Eine Schnittstelle in XLA, die es ermöglicht, externe, native C++/GPU-Kernel direkt aufzurufen, ohne dass der XLA-Compiler deren innere Mathematik analysieren muss.
• VJP (Vector-Jacobian Product): Ein mathematischer Operator der automatischen Differentiation, der im Rückwärtspass von neuronalen Netzen zur Gradientenberechnung verwendet wird.
• MHA (Multi-Head Attention): Ein zentraler Aufmerksamkeitsmechanismus in Transformer-Architekturen, der die parallele Verarbeitung verschiedener Sequenzabschnitte ermöglicht.
• MXFP4: Microscaling Floating-Point 4-bit; ein extrem niedrigpräzises 4-Bit-Fließkommaformat mit blockweisem Skalierungsfaktor zur Beschleunigung von KI-Training und Inferenz.
• BF16 / FP8: Bfloat16 (16-Bit) und Float8 (8-Bit); numerische Fließkomma-Formate zur Optimierung von Speicherbedarf und Rechenleistung in Deep-Learning-Workloads.
• Sharding: Die Aufteilung von Daten oder Tensor-Dimensionen auf mehrere Beschleunigerkarten (GPUs) zur verteilten Berechnung.
• CK (Composable Kernel): Eine C++-GPU-Bibliothek von AMD zur Erstellung performanter Rechen-Kernel.
• Hadamard-Transformation & Stochastic Rounding: Mathematische Verfahren zur Reduzierung von Quantisierungsfehlern und Ausreißern bei der Umwandlung von Daten in Niedrigpräzisionsformate wie MXFP4.

Abstract

Dieser technische Vortrag behandelt das Post-Training von Large Language Models (LLMs) unter Verwendung von Tunix („Tune in JAX“), einer Open-Source-Bibliothek für das JAX-Ökosystem. Der Schwerpunkt liegt auf der systematischen Abgrenzung zwischen Pre-Training (generische Wissensakkumulation auf unstrukturierten Daten) und Post-Training (Verhaltensausrichtung, Feinausrichtung an menschlichen Präferenzen und Domänenspezialisierung mittels kuratierter Datensätze).

Der Vortrag analysiert die Architektur von Tunix im Zusammenspiel mit JAX-Bibliotheken wie Flax (neuronaler Netzaufbau), Optax (Gradienten-Optimierung) und Grain (E/A-Datenströme) sowie die Skalierungsebene MaxText. Es werden drei zentrale Post-Training-Methoden detailliert aufgeführt: Supervised Fine-Tuning (SFT) inklusive parameter-effizienter Methoden (LoRA und 4-Bit-quantisiertes QLoRA), Direct Preference Optimization (DPO) für Paar-Präferenzvergleiche sowie Reinforcement Learning mittels Group Relative Policy Optimization (GRPO) zur Steigerung von Reasoning-Fähigkeiten. Abschließend werden Hardware-Optimierungsstrategien zur Vermeidung von Out-Of-Memory-Fehlern (OOM) sowie die kritische Rolle kontinuierlicher Evaluierungspipelines zur Vermeidung von Entropie-Kollaps erörtert.

Wichtigste Punkte & Zeitstempel

• 0:00 Pre-Training vs. Post-Training: Unterscheidung zwischen dem grundlegenden Wissenserwerb (Pre-Training) und der gezielten Verhaltensausrichtung sowie Domänenspezialisierung (Post-Training) mit signifikant kleineren, kuratierten Datensätzen.
• 6:21 Tunix-Architektur im JAX-Ökosystem: Einordnung von Tunix als modulare, JAX-native Bibliothek („Tune in JAX“), die auf Flax, Optax und Grain aufbaut und als Unterbau für verteilte Workflows in MaxText dient.
• 10:53 Supervised Fine-Tuning (SFT) & PEFT: Grundlagen des SFT zur Festlegung des Modellverhaltens sowie der Einsatz von Parameter-Efficient Fine-Tuning (LoRA / QLoRA mit 4-Bit-Quantisierung) zur Reduktion von Rechen- und Speicheraufwand.
• 11:42 Direct Preference Optimization (DPO): Verfeinerung von Antwortpräferenzen durch Paar-Vergleiche (bevorzugte vs. abgelehnte Ausgaben) auf Basis des zuvor erstellten SFT-Checkpoints.
• 13:43 Reinforcement Learning & GRPO: Implementierung von RL-Pipelines mit Algorithmen wie Group Relative Policy Optimization (GRPO) zur spezifischen Steigerung der mathematischen und logischen Schlussfolgerungsfähigkeiten (Reasoning).
• 22:15 Hardware-Optimierung & OOM-Prophylaxe: Strategien zur Adressierung von Out-Of-Memory-Engpässen auf TPUs und GPUs durch Konfiguration globaler Batch-Größen und Hyperparameter.
• 25:08 Entscheidungsmatrix & Datensatzanforderungen: Leitfaden zur Algorithmenauswahl basierend auf der Datenverfügbarkeit; Erläuterung, dass bereits einige tausend hochwertige Beispiele ausreichen, um signifikante Verhaltensänderungen zu erzielen.
• 28:31 Tunix vs. MaxText & Evaluierung: Einsatz von Tunix für feingranulare Low-Level-Steuerung gegenüber MaxText für Cluster-Skalierung; Relevanz begleitender Evaluierungspipelines zur Vermeidung von Entropie-Kollaps.

Glossar

• LLM (Large Language Model): Großes Sprachmodell auf Basis tiefes neuronaler Netzwerke, das zur Verarbeitung und Erzeugung von natürlicher Sprache trainiert wird.
• Pre-Training: Initiale, rechenintensive Trainingsphase eines Sprachmodells auf unstrukturierten Massendaten zur Aneignung von Sprach- und Weltwissen.
• Post-Training: Nachgelagerter Trainingsprozess zur Verhaltensanpassung, Instruktionsbefolgung und Ausrichtung an Anwenderpräferenzen.
• Tunix: Eine auf JAX basierende Open-Source-Bibliothek („Tune in JAX“), die modularisierte Workflows für das Post-Training von Sprachmodellen bereitstellt.
• JAX: Ein von Google entwickeltes Framework für hochperformante numerische Berechnungen und automatische Differentiation im maschinellen Lernen.
• Flax: Eine auf JAX aufbauende Bibliothek zur Definition und zum Training neuronaler Netzwerke.
• Optax: Eine Gradientenverarbeitungs- und Optimierungsbibliothek für das JAX-Ökosystem.
• Grain: Eine Bibliothek innerhalb des JAX-Ökosystems zur effizienten Datenladung und -vorverarbeitung (I/O).
• MaxText: Eine auf Tunix und JAX aufbauende High-Performance-Bibliothek zur Skalierung des LLM-Trainings auf große TPU- und GPU-Cluster.
• SFT (Supervised Fine-Tuning): Überwachtes Feinabstimmen, bei dem ein Basismodell mit Prompt-Antwort-Paaren trainiert wird, um ein bestimmtes Antwortverhalten zu erlernen.
• PEFT (Parameter-Efficient Fine-Tuning): Sammelbezeichnung für Methoden, die nur einen geringen Prozentsatz der Modellparameter anpassen, um Speicherplatz und Rechenzeit zu sparen.
• LoRA (Low-Rank Adaptation): Eine PEFT-Technik, bei der gewichtete Low-Rank-Matrizen in die Aufmerksamkeits-Ebenen des Modells eingefügt werden, anstatt alle Modellgewichte zu aktualisieren.
• QLoRA (Quantized Low-Rank Adaptation): Eine Variante von LoRA, bei der das Basismodell in quantisierter Form (z. B. 4-Bit) im Speicher gehalten wird, um den Speicherbedarf weiter zu senken.
• DPO (Direct Preference Optimization): Ein Algorithmus zur Ausrichtung von Modellen an menschlichen Präferenzen anhand von Datenpaaren (bevorzugt/abgelehnt), ohne dass ein separates Belohnungsmodell trainiert werden muss.
• GRPO (Group Relative Policy Optimization): Ein Reinforcement-Learning-Algorithmus zur Optimierung von Zielparametern, insbesondere zur Steigerung von Reasoning- und Logikfähigkeiten.
• OOM (Out Of Memory): Ein Hardware-Fehlerzustand, bei dem der zur Verfügung stehende Grafik- oder Tensor-Speicher (VRAM/HBM) durch Modellparameter und Aktivierungen überschritten wird.
• Entropie-Kollaps: Ein Phänomen beim Modelltraining, bei dem die Stochastik und Vielfalt der Modellantworten stark degradiert und das Modell zu redundanten oder monotonen Antworten neigt.

Abstract

In diesem Vortrag präsentiert Sayak Paul, Senior Research Engineer bei Hugging Face, die Entwicklung einer TPU-optimierten Implementierung des Bildgenerierungsmodells „Qwen Image“ von Alibaba. Der Beitrag erläutert die Funktionsweise von Latent Diffusion Models (LDM) als mehrstufige Pipelines im Gegensatz zu Monolithen sowie die Integration der Hugging Face diffusers-Bibliothek mit PyTorch XLA. Um das Modell effizient auf Google TPU v6e-8 Hardware auszuführen, wurde das Text-Encoder-Modell (Qwen 2.5 VL) auf der CPU belassen, während der speicherintensive Diffusion Transformer (DiT) mittels einer 2D-Mesh-SPMD-Strategie über 8 TPU-Chips hinweg geshardet wurde. Trotz einer initialen XLA-Kompilierungszeit von 95 Minuten erreichte das System im Steady-State eine Inferenzzeit von 20 Sekunden für 50 Denoising-Schritte und übertraf damit eine ungeshardete A100-Referenzkonfiguration. Abschließend wird der zukünftige Wechsel zu torch_tpu zur vereinfachten Portierbarkeit diskutiert.

Wichtigste Punkte & Zeitstempel

• 0:00 TPU-Portierung von Qwen Image: Vorstellung des Projekts zur Bereitstellung von Alibabas Open-Source-Bildgenerierungsmodell auf Google TPUs unter Nutzung von PyTorch XLA und Hugging Face diffusers.
• 1:33 Architektur von Latent Diffusion Models: Bildgenerierungsmodelle arbeiten vorwiegend in einem niederdimensionalen Latentraum und nutzen iterative Denoising-Schritte anstelle klassischer autoregressiver Verfahren.
• 3:16 Pipeline-Komponenten: Das Qwen-Image-System besteht aus mehreren separaten Teilmodellen: dem Text-Encoder Qwen 2.5 VL, dem iterativen Diffusion Transformer (DiT) und dem Variational Autoencoder (VAE) als Decoder.
• 5:44 PyTorch XLA-Integration in Diffusers: Die Einbindung in PyTorch XLA erfolgt primär über gezielte mark_step()-Aufrufe zur Steuerung der Kompilierungsgrenzen sowie den Einsatz experimenteller XLA FlashAttention.
• 9:19 Hybride CPU-TPU-Ausführungsstrategie: Der Text-Encoder wird auf der CPU ausgeführt, da er nur 1 % der Rechenlast ausmacht und somit Kompilierungsaufwand spart, während DiT und VAE auf dem TPU-Beschleuniger laufen.
• 12:03 SPMD-Sharding auf TPU v6e-8: Da der Diffusion Transformer nicht auf einen einzelnen TPU-Chip passt, wird das Modell mittels 2D-Mesh und Single-Process-SPMD entlang der größten Parameterdimensionen über 8 Chips verteilt.
• 15:16 XLA-Debugging & Fehlerbehebung: Behebung von Fehlern bei der Prompt-Kodierung sowie Eliminierung negativer Indexierungsfehler, die von XLA im Gegensatz zu CUDA/CPU explizit als Laufzeitverletzungen ausgewiesen werden.
• 16:24 Kompilierung & Inferenz-Latenz: Nach einer initialen XLA-Kompilierungszeit von 95 Minuten sinkt die Inferenzdauer auf 20 Sekunden pro Bild (50 Schritte) auf der TPU v6e-8 im Vergleich zu ~35 Sekunden auf einer NVIDIA A100.
• 17:48 Laufzeitanalyse und Profiling: Profiling-Ergebnisse bestätigen hohe Graphdichte, effektive Operationsfusion und minimale Host-to-Device-Transfers im kompilierten Ausführungsgraphen.
• 19:31 Ausblick auf Torch TPU: Künftige Implementierungen zielen auf die Nutzung von torch_tpu ab, um das manuelle Setzen von Kompilierungsgrenzen zu erübrigen und die Code-Portierbarkeit zwischen CUDA und TPU weiter zu vereinfachen.

Glossar

• Latent Diffusion Model (LDM): Ein Bildgenerierungsmodell, das den iterativen Rauschentfernungsprozess (Denoising) in einem komprimierten latenten Vektorraum statt direkt auf hochauflösenden Bildpixeln ausführt.
• Diffusion Transformer (DiT): Eine Modellarchitektur für Diffusionsmodelle, die auf dem Transformer-Prinzip (Selbstaufmerksamkeitsmechanismen) basiert und ältere U-Net-Netzwerke ersetzt.
• PyTorch XLA: Eine Bibliothek, die das PyTorch-Framework mit dem XLA-Compiler (Accelerated Linear Algebra) verbindet, um PyTorch-Code auf Google TPUs und anderen Beschleunigern auszuführen.
• Single Process Multiple Data (SPMD): Ein paralleles Programmiermodell, bei dem derselbe Ausführungscode gleichzeitig auf mehreren Prozessoren läuft, wobei jeder Prozessor unterschiedliche Datensegmente verarbeitet.
• Sharding: Die Aufteilung großer Modellparameter oder Datensätze auf mehrere physische Speicherbausteine bzw. TPU-Chips zur Bewältigung von Arbeitsspeicherbeschränkungen.
• FlashAttention: Ein speichereffizienter und exakter Aufmerksamkeitsalgorithmus (Attention Mechanism), der den E/A-Lese- und Schreibzugriff zwischen GPU/TPU-SRAM und HBM optimiert.
• Variational Autoencoder (VAE): Ein neuronales Netzwerk, das Eingabedaten in einen latenten Raum kodiert und daraus neue Daten rekonstruieren bzw. dekodieren kann.
• Tensor Processing Unit (TPU): Ein von Google entwickelter anwendungsspezifischer integrierter Schaltkreis (ASIC), der speziell für die Beschleunigung von Algorithmen des maschinellen Lernens konzipiert ist.

Zusammenfassung

Dieses Transkript dokumentiert eine Präsentation des Google-Teams (Britney Rockwell, Kuan Kim und Winshin Dong) auf einem DevLab-Event, in der die Skalierung von Large Language Models (LLMs) und Mixture-of-Experts-Modellen (MoE) für die Inferenz auf Google-TPUs (Tensor Processing Units) behandelt wird. Im Fokus stehen die Integration von TPUs in Open-Source-Inferenz-Engines wie vLLM und SGLang (mittels JAX), horizontale und vertikale Optimierungstechniken (wie Instruction Reordering und Kernel-Fusion), die Unterstützung von Reinforcement Learning (RL) über Pathways sowie die kommende Einführung des PyTorch-TPU-Backends im Oktober.

Wichtigste Punkte & Zeitstempel

• 0:30 vLLM- und SGLang-JAX-Backends: Integration von TPUs direkt in Open-Source-Inferenz-Engines zur Steigerung des Preis-Leistungs-Verhältnisses (Perf-per-Dollar), Erhöhung der Hardware-Verfügbarkeit und Minimierung von Entwickler-Reibungsverlusten.
• 3:50 Fokus auf Ironwood und Trillium: Anpassung der Engines an den neuen Ironwood-Chip (XLA-Optimierungen), um speicher- und rechenintensive Reasoning-Modelle im Vergleich zum kleineren Trillium-Chip effizienter auszuführen.
• 5:14 Mixture-of-Experts (MoE) Herausforderungen: Bewältigung von Token-Ungleichgewichten, Expert Parallelism (EP) und Routing-Strategien bei extrem skalierten MoE-Modellen (XL-MoE).
• 7:38 Leistungssteigerung bei Qwen-Modellen: Steigerung der theoretischen Roofline-Effizienz von Qwen-Modellen von 2 % auf bis zu 77 % (ca. 95 % bei realistischen Annahmen) durch neue GMM-Kernel (General Matrix Multiply), Ragged Gather/Scale und Batched-RPA-Kernel.
• 11:02 Horizontale Optimierungen (Token Routing): Ersetzen des ineffizienten All-to-All-Netzwerk-Overheads durch vollständige Token-Replikation und Vorziehen von Reduktionsoperationen, was zu einer Verdopplung der Performance führt.
• 14:31 Vertikale Optimierungen (Kernel-Fusion): Nutzung des großen, exklusiven Vector Memory (VM) pro TPU-Kern für aggressive Fusionsoperationen (Kombination von GMM, Aktivierungsfunktionen wie SiLU sowie In- und Dequantisierung).
• 16:58 Modellübergreifende Code-Wiederverwendung: Einsatz von Compiler-basierten Kostenmodellen und automatisierter Kachelgrößen-Optimierung (Tile Sizing), wodurch die Portierungszeit für Qwen 3.5 auf zwei Monate verkürzt wurde.
• 19:45 Reinforcement Learning (RL) auf TPUs: Winshin Dong erläutert die Orchestrierung von MaxText, Tunx, TPU-Inference (vLLM) und Pathways für skalierbare RL-Infrastrukturen in kolokierten oder disaggregierten Modi.
• 25:12 Native vLLM-Erweiterungen für RL: Implementierung von Data Parallelism, effizientem Gewichts-Sync (Weight Sync) über Cache-Neuzuweisung, Pathways-Integration und Expert Replay zur Beschleunigung der Konvergenz.
• 27:52 GRPO-Skalierungsergebnisse: Erzielung einer Sampler-Throughput-Leistung von 35k Token/Sekunde und einer globalen Schrittzeit von 6 Minuten bei einem GRPO-Lauf (Group Relative Policy Optimization) mit Qwen3-35B auf 128 Ironwood-Chips.
• 28:50 Zukünftige Roadmap & PyTorch-TPU-Launch: Die kommenden sechs Monate konzentrieren sich auf Multimodalität, langen Kontext (Agentic Workloads), Rack-Scale-Multi-Host-Support und den Start des PyTorch-TPU-Backends im Oktober für vLLM und SGLang.

Glossar

• Large Language Model (LLM): Ein fortgeschrittenes KI-Modell, das auf massiven Textdaten trainiert wurde, um menschliche Sprache zu verstehen und zu generieren.
• Tensor Processing Unit (TPU): Ein von Google entwickelter, spezialisierter Mikrochip, der gezielt für die Beschleunigung von rechenintensiven Workloads im Bereich des maschinellen Lernens optimiert ist.
• Inferenz (Inference): Die Betriebsphase eines KI-Modells, in der dieses auf Basis bereits gelerntem Wissens Vorhersagen trifft oder Antworten auf neue Anfragen generiert.
• Mixture-of-Experts (MoE): Eine neuronale Netzwerkarchitektur, bei der ein Modell in mehrere spezialisierte Subnetze ("Experten") unterteilt wird, von denen bei jeder Anfrage nur ein Bruchteil aktiv geschaltet wird, um Rechenleistung zu sparen.
• vLLM & SGLang: Weit verbreitete Open-Source-Inferenz-Engines, die für den Hochdurchsatz und die speichereffiziente Ausführung von Sprachmodellen optimiert sind.
• JAX: Ein von Google entwickeltes Programmierframework für Hochleistungs-Rechenoperationen und automatische Differenzierung, das häufig im modernen KI-Training und -Inferenz-Stack eingesetzt wird.
• XLA (Accelerated Linear Algebra): Ein domainspezifischer Compiler von Google, der Berechnungen linearer Algebren optimiert, um die Ausführungsgeschwindigkeit und Speichereffizienz auf Hardware wie TPUs zu maximieren.
• GMM (General Matrix Multiply): Eine grundlegende mathematische Rechenoperation (Matrixmultiplikation), die den überwiegenden Teil der Rechenlast in neuronalen Netzen darstellt.
• Vector Memory (VM): Ein extrem schneller, lokaler Hochgeschwindigkeitsspeicher, der direkt in einen TPU-Kern integriert ist.
• Pathways: Ein von Google entwickeltes, verteiltes Multi-Host-Orchestrierungssystem, das die Ausführung komplexer, paralleler KI-Workloads über Tausende von TPU-Chips hinweg koordiniert.
• Reinforcement Learning (RL): Ein maschinelles Lernverfahren, bei dem ein Modell durch Versuch und Irrtum (Belohnungen und Bestrafungen) lernt, optimale Handlungsstrategien zu entwickeln.
• GRPO (Group Relative Policy Optimization): Ein effizientes Reinforcement-Learning-Verfahren zur Optimierung von Sprachmodellen, das im Gegensatz zu älteren Methoden ohne ein separates Bewertungsnetzwerk (Critic) auskommt.

Zusammenfassung

David Masher, Forschungsdatenwissenschaftler am Cedars-Sinai Medical Center in Los Angeles, präsentiert in diesem Vortrag einen methodischen Ansatz zur Integration von professionellem medizinischem Fachwissen in Open-Source-Large-Language-Models (LLMs) anhand von Daten aus elektronischen Gesundheitsakten (EHR). Der Vortrag verdeutlicht, dass herkömmliche LLM-Benchmarks, die auf medizinischen Examensfragen basieren, die tatsächliche Komplexität klinischer Aufgaben – wie longitudinale Patientenverläufe, unvollständige Daten und symptomatische Ambiguität – nicht adäquat abbilden. Es wird ein Supervised Fine-Tuning (SFT) des Modells Qwen-34B auf Basis realer Notaufnahmedaten (ED-Notes) erläutert, um das Modul für medizinische Entscheidungsfindung (Medical Decision Making - MDM) zu imitieren. Die Evaluation erfolgte sowohl qualitativ durch verblindete Notärzte (Anhand von fünf Dimensionen) als auch quantitativ mittels ICD-10-Klassifikationsvorhersagen.

Wichtigste Punkte & Zeitstempel

• 0:00 Limitationen von LLM-Examens-Benchmarks: Herkömmliche Benchmarks basieren auf starren Prüfungsfragen und spiegeln nicht die longitudinale Komplexität, unvollständige Diagnostik und symptomatische Ambiguität des klinischen Alltags wider.
• 2:13 Das EHR-Datenproblem: Wertvolles klinisches Erfahrungswissen befindet sich in elektronischen Gesundheitsakten (EHR) hinter strengen Datenschutzbarrieren in Krankenhäusern, während im Internet zugängliche Daten meist nur Leitlinien, klinische Studien oder Lehrbücher umfassen.
• 3:50 Feinabstimmung auf Medical Decision Making (MDM): Einsatz von regulären Ausdrücken (RegEx) zur Segmentierung von Notaufnahmedaten (ED-Notes) bei Cedars-Sinai, um das MDM-Modul zu isolieren und das Open-Source-Modell Qwen-34B darauf zu trainieren, den ärztlichen Entscheidungsfindungsprozess nachzubilden.
• 5:30 Technische Trainingsinfrastruktur: Training auf acht Nvidia A100 (80GB) GPUs mittels Hugging Face TRL (Transformer Reinforcement Learning) unter Verwendung der Cross-Entropy-Loss-Funktion auf Tokens sowie 10 % externer Daten zur Vermeidung von katastrophalem Vergessen; die Trainingsdauer betrug 172 Stunden bei Batch-Größe 1.
• 6:20 Qualitative und Quantitative Evaluation: ED-Ärzte bewerteten das generierte MDM auf fünf Dimensionen (Genauigkeit, Halluzination, Vollständigkeit, Nutzen, Verständlichkeit), wobei das trainierte Modell die Baseline übertraf und eine 16-prozentige Verbesserung des F1-Scores bei der ICD-10-Diagnosevorhersage erzielte.
• 9:00 Generalisierungsaufgabe und Label-Kollaps: Ein Test zur Erkennung seltener kardialer Arrests in Notfallberichten zeigte einen partiellen Label-Kollaps, der eine nachgelagerte Voll-Feinabstimmung (Full Fine-Tuning) zur Fokussierung der Aufmerksamkeit erforderte.
• 10:04 Verlust der Chain-of-Thought-Fähigkeiten: Nach dem Training zeigte sich ein Rückgang der Reasoning-Fähigkeiten des Qwen-Modells, da die ursprünglichen Chain-of-Thought-Spuren im Trainingsdatensatz fehlten und als leere Tags generiert wurden.
• 11:38 Skalierungspläne und Zukünftige Optimierungen: Geplante nächste Schritte umfassen das Hochskalieren auf 14B- und 70B+-Parameter-Modelle mittels TPUs, verfeinerte Regex-Segmentierung zur Rückgewinnung verlorener Notizen, die Beibehaltung numerischer Metadaten bei der De-Identifikation sowie den Übergang zu Reinforcement Learning.

Diskussions-Highlights

• Implizite Diagnostik: Es wurde diskutiert, dass das gezielte Auslassen bestimmter medizinischer Untersuchungen durch Ärzte ein wichtiges Diagnose-Signal darstellt, welches zukünftig im Rahmen des Reinforcement Learnings explizit modelliert werden soll.
• Multi-Site-Datenaggregation: Die genutzten Daten stammen aus mehreren Standorten des Cedars-Sinai-Netzwerks (Hauptkrankenhaus, Pasadena, Del Rey Marina) und umfassen die kollaborative Dokumentation verschiedener Ärzte.
• Evaluierungs-Skalierung: Aufgrund unzureichender Qualität früherer Modellversionen wurde die qualitative Annotation auf 25 Fälle durch zwei verblindete Notärzte beschränkt; bei skalierten Modellen ist eine umfassendere Validierung geplant.

Glossar

• EHR (Electronic Health Record): Elektronische Gesundheitsakte; ein digitales System zur zentralen Speicherung der vollständigen medizinischen Historie, Diagnosen und Behandlungsdaten von Patienten in Kliniken.
• MDM (Medical Decision Making): Medizinische Entscheidungsfindung; ein strukturierter Abschnitt in Notfall- oder Arztbriefen, der die Symptompräsentation, Verdachtsdiagnosen und therapeutische Empfehlungen zusammenfasst.
• Qwen-34B: Ein Open-Source-Large-Language-Model mit 34 Milliarden Parametern, entwickelt von Alibaba.
• Cross-Entropy-Loss: Eine mathematische Verlustfunktion im maschinellen Lernen, die die Diskrepanz zwischen vorhergesagten Wahrscheinlichkeiten und den tatsächlichen Token-Verteilungen misst.
• ICD-10: International Classification of Diseases (10. Revision); ein international standardisiertes System zur Verschlüsselung von Diagnosen und medizinischen Symptomen.
• Chain-of-Thought (CoT): Eine Prompting- oder Trainingsmethode, die ein Sprachmodell dazu anregt, seine logischen Denkschritte explizit schrittweise vor der finalen Antwort auszugeben.
• Katastrophales Vergessen (Catastrophic Forgetting): Ein Phänomen im maschinellen Lernen, bei dem ein neuronales Netz beim Erlernen neuer Fachdomänen (hier Medizin) die Fähigkeiten für zuvor erlernte allgemeine Aufgaben verliert.
• De-ID (De-identification): Anonymisierungsprozess von Patientendaten in Krankenhäusern, bei dem sensible persönliche Identifikatoren (PHI) und in manchen Systemen auch klinische Werte entfernt oder maskiert werden.

Zusammenfassung

Diese Aufzeichnung einer AI-Dev-Lab-Breakout-Session präsentiert MaxText, ein LLM-Referenzimplementierungs-Framework auf Basis von JAX, das für groß angelegte Pre- und Post-Trainings-Workloads auf Google-TPUs optimiert ist. Die Ingenieure erläutern neue Funktionserweiterungen, darunter die native Unterstützung für State-of-the-Art-Open-Source-Modelle wie Gemma 4 und DeepSeek, elastische Trainingsmechanismen zur Fehlertoleranz mittels Pathways sowie optimierte Entwickler-Workflows durch vereinfachte Pip-Installationen und dynamisches Laden von Safetensors.

Wichtigste Punkte & Zeitstempel

• 0:00 MaxText-Ökosystem: Vorstellung des Frameworks für Großformat-Pre-Trainings auf TPUs, aufbauend auf JAX, dem Compiler XLA, benutzerdefinierten Pallas-Kernels sowie Bibliotheken wie Orbax (Checkpointing), Tunix (Post-Training) und Quix (Quantisierung).
• 0:47 Modell- und Skalierungsunterstützung: Integration von OSS-Modellen wie DeepSeek und Gemma, skalierbar von wenigen Hunderten bis zu Zehntausenden TPUs, unterstützt durch automatisierte Kompilierung und fortgeschrittene Parallelitätstechniken.
• 05:00 Gemma 4 & Multimodalität: Native Unterstützung für Gemma 4 (inklusive 31B Dense, 26B MoE sowie 2B- und 4B-Varianten) mit Fokus auf physische KI, Bild- und Texteingaben, Argumentation sowie erweiterte multimodale Sequenzverarbeitung.
• 06:42 Elastisches Training (Pathways): Mechanismus zur dynamischen Bewältigung von TPU-Ausfällen ohne Cluster-weite Neustarts durch Pause-Resume und Replica Resize, steuerbar über spezifische XPK- und MaxText-Konfigurations-Flags.
• 08:41 Leistungskennzahlen & Monitoring: Überwachung der Trainingsinfrastruktur mittels Elastic Bad Put Breakdown, kapazitätsbewusstem Good Put sowie Metriken wie MTTR (Mean Time to Recovery) und MTTI (Mean Time Between Interruptions).
• 10:15 Entwickler-Erfahrung & Bereitstellung: Direkte Pip-Installation ohne Klonen oder komplexe Setup-Skripte, konsistente Abhängigkeiten für Reproduzierbarkeit, vereinfachte Multi-Host-Docker-Builds und On-the-fly-Laden von Hugging-Face-Safetensors.

Glossar

• JAX: Ein von Google entwickeltes Python-Framework für hochleistungsfähige numerische Berechnungen und automatisierte Differenzierung im maschinellen Lernen.
• TPU (Tensor Processing Unit): Ein von Google speziell für beschleunigte Workloads im Bereich des maschinellen Lernens entwickelter ASIC-Prozessor.
• XLA (Accelerated Linear Algebra): Ein domainspezifischer Compiler, der lineare Algebra-Berechnungen optimiert, um die Ausführungsgeschwindigkeit auf Hardware wie TPUs zu erhöhen.
• Orbax: Eine JAX-basierte Bibliothek, die für optimiertes Checkpointing (Sichern von Modellzwischenständen) bei großen Trainingsläufen verwendet wird.
• MoE (Mixture of Experts): Eine Modellarchitektur, bei der verschiedene Teilnetzwerke (Experten) dynamisch für verschiedene Aufgaben eingesetzt werden, um die Effizienz bei enormen Parameterzahlen zu steigern.
• Pathways: Ein von Google entwickeltes verteiltes Rechensystem, das elastisches Training und die nahtlose Verwaltung großer Rechencluster ermöglicht.
• XPK: Ein Bereitstellungstool (Tool for GKE), das das Ausführen von ML-Workloads auf Google Kubernetes Engine und TPUs vereinfacht.
• Safetensors: Ein sicheres und schnelles Dateiformat von Hugging Face zum Speichern von Gewichten maschineller Lernmodelle.
• Good Put / Bad Put: Kennzahlen zur Messung der tatsächlichen, effektiven Rechenleistung (Good Put) im Vergleich zu verlorener Zeit durch Ausfälle oder Re-Kompilierungen (Bad Put).
• MTTR (Mean Time to Recovery): Die durchschnittliche Zeit, die benötigt wird, um ein System nach einem Ausfall wiederherzustellen.
• MTTI (Mean Time Between Interruptions): Die durchschnittliche Zeitspanne zwischen unvorhergesehenen Systemunterbrechungen oder Ausfällen.

Zusammenfassung

Dieser Vortrag präsentiert den aktuellen Entwicklungsstand von Torch TPU, einer kollaborativen Initiative von Google und Meta zur nativen Ausführung von PyTorch-Workloads auf Google-TPUs. Die Vortragenden (Clauddio von Google und Colin von Meta) erläutern die architektonischen Grundlagen, die Integration von Software-Stacks wie Stable HLO und dem XLA-Compiler, neue Eager-Ausführungsmodi (inklusive des hybriden „Defer and Fuse“-Verfahrens) sowie produktionsreife Inferenz- und Trainingsintegrationen mit vLLM und SGLang.

Wichtigste Punkte & Zeitstempel

• 0:00 Motivation & Kollaboration: Google und Meta kooperieren bei Torch TPU, um fragmentierte Hardware-Flotten zu vereinheitlichen und den Entwicklungsaufwand beim Portieren von PyTorch-Modellen auf TPUs zu minimieren.
• 0:25 Code-Integration & Bereitstellung: Demonstration der Ausführung von nativem PyTorch-Code auf TPUs mittels torch.device("tpu"), Feinabstimmung von Llama 3 8B über LoRA und Ausführung über XPK auf 256 Ironwood-TPU-Chips.
• 6:03 Architektur-Stack: Der Software-Stack leitet PyTorch-Operationen über Stable HLO an den bewährten XLA-Compiler weiter, wobei konsequent PyTorch-native Programmierparadigmen bevorzugt werden.
• 9:16 Ausführungsmodi (Strict, Debug & Defer-and-Fuse): Neben dem standardmäßigen Strict-Modus für numerische Stabilität und dem blockierenden Debug-Modus führt Torch TPU den neuen Defer and Fuse-Modus ein, der Operationen heuristisch akkumuliert und durch Fusion über 50 % Leistungssteigerung liefert.
• 11:39 Kompilierungs-Cache & torch.compile: Implementierung eines persistenten Kompilierungscaches zur Vermeidung redundanter Neukompilierungen, wobei torch.compile über Dynamo und ein "Defer-all"-Verfahren direkt in den XLA-Stack eingespeist wird.
• 18:00 Benutzerdefinierte Kernel & Verteilung: Unterstützung von Palace- und Jax-Kerneln, Validierung von verteilten Paradigmen inklusive SPMD und MPND, FSDP v2 sowie Optimierungen für Sparse Cores.
• 36:11 Inferenz-Integrationen (vLLM & SGLang): Enger Zusammenschluss mit Inferenz-Engines durch Wiederverwendung von Scheduling, kontinuierlichem Batching und KV-Cache-Managern sowie Nutzung statischer Shape-Bucketisierung.
• 42:53 Benchmarks & Modellabdeckung: Umfangreiche Tests mit der Qwen 3-Familie (Dense, MoE, Vision-Language-Modelle), wobei auf Qwen 3.0 eine Leistungsparität oder -überlegenheit gegenüber dem reinen vLLM-TPU-Stack erreicht wird.
• 48:47 Trainings-Optimierungen: Einsatz eines make_effects-basierten Dispatch-Tracers für das Training, um vollständige Graphen inklusive Optimizer-Schritten direkt zu optimieren.
• 55:03 Ausblick & Open-Source-Pläne: Geplante Open-Source-Freigabe der Repositories und Integrationen zum Jahresende zeitgleich mit dem offiziellen Torch-TPU-Release.

Glossar

• TPU (Tensor Processing Unit): Von Google entwickelte, spezialisierte ASIC-Prozessoreinheit, die für beschleunigte Machine-Learning- und Deep-Learning-Workloads optimiert ist.
• PyTorch: Ein weit verbreitetes Open-Source-Framework für maschinelles Lernen, das auf Python basiert und für dynamische Berechnungsdiagramme bekannt ist.
• Stable HLO: Eine stabile High-Level-Optimierungs-Zwischenrepräsentation (IR) für Machine-Learning-Compiler, die als standardisierte Schnittstelle zwischen Frontend-Frameworks und Backend-Compilern dient.
• XLA (Accelerated Linear Algebra): Ein domainspezifischer Compiler für lineare Algebra von Google, der Berechnungsdiagramme von Machine-Learning-Frameworks optimiert und in ausführbare Binärdateien übersetzt.
• LoRA (Low-Rank Adaptation): Eine effiziente Fine-Tuning-Technik für große Sprachmodelle, bei der nur eine minimale Anzahl von Gewichtsparametern durch Rang-Zerlegung angepasst wird.
• Vertex AI: Eine von Google Cloud bereitgestellte Managed-Machine-Learning-Plattform zum Trainieren, Verwalten und Bereitstellen von KI-Modellen.
• XPK: Ein Open-Source-Tool zur Orchestrierung und Bereitstellung von Workloads auf Google Cloud TPUs im Großmaßstab.
• Ironwood: Eine spezifische Generation von Google-TPU-Chips.
• Eager Mode: Ein Ausführungsmodus in Frameworks wie PyTorch, bei dem Operationen sofort beim Aufruf ausgeführt werden, anstatt sie vorab zu kompilieren.
• Dynamo (TorchDynamo): Eine PyTorch-Frontend-Komponente, die Python-Bytecode zur Laufzeit abfängt und in ein FX-Berechnungsdiagramm zur Kompilierung umwandelt.
• Inductor: Der standardmäßige Code-Generator und PyTorch-Compiler für torch.compile, der primär für GPUs optimiert ist.
• Palace Kernel: Eine domänenspezifische Programmiersprache (DSL) bzw. ein Kernel-Schreibstil zur Erstellung von benutzerdefinierten Low-Level-Tuning-Kerneln auf TPUs.
• SPMD (Single Program, Multiple Data): Ein paralleles Programmiermodell, bei dem dasselbe Programm auf mehreren Prozessoren mit unterschiedlichen Daten ausgeführt wird.
• MPND (Multi-Program, Multiple Data): Ein paralleles Programmiermodell, bei dem verschiedene Programme auf verschiedenen Prozessoren mit unterschiedlichen Daten ausgeführt werden.
• FSDP (Fully Sharded Data Parallel): Eine Speicher- und Rechenoptimierungstechnik im verteilten Training, bei der Modellparameter, Gradienten und Optimizer-Zustände über mehrere Beschleuniger gesharded (aufgeteilt) werden.
• vLLM / SGLang: Beliebte Open-Source-Inferenz-Engines für Large Language Models, die auf hohem Durchsatz, Paged Attention und effizienter Speicherverwaltung basieren.
• KV-Cache (Key-Value Cache): Ein Zwischenspeicher für Schlüssel- und Wertvektoren bei der Transformer-Inferenz, der die Neuberechnung bereits generierter Token-Historien vermeidet.
• MoE (Mixture of Experts): Eine Architektur für neuronale Netze, bei der nur eine kleine Teilmenge von Teilnetzwerken ("Experten") dynamisch für jeden Token aktiviert wird.
• Qwen 3 / Llama 3: Bekannte Open-Weights-Familien von großen Sprachmodellen.

Zusammenfassung

Dieser Vortrag präsentiert ein KI-gestütztes System zur Echtzeitübersetzung der indonesischen Gebärdensprache (Bisindo) mithilfe von Motion-Grids und On-Device-Kameras. Das System kombiniert ein multimodales Visions-Sprachmodell mit einem 3-Milliarden-Parameter-Backbone, LoRA-Adaptern und Gemma-Modellen zur grammatikalischen Nachbearbeitung. Durch die Umwandlung von Videosequenzen in 3x3-Motion-Grids werden dynamische Handzeichen als Bild-zu-Text-Aufgabe verarbeitet. Das System ist für eine plattformübergreifende Bereitstellung auf Mobilgeräten, Web-Anwendungen und Cloud-Endpunkten ausgelegt.

Wichtigste Punkte & Zeitstempel

• 0:00 Forschungsprojekt zur Gebärdensprachübersetzung: Vorstellung der Initiative zur Echtzeitübersetzung der indonesischen Gebärdensprache (Bisindo) für gehörlose Menschen in Indonesien.
• 2:19 Linguistische Struktur der Bisindo: Die Grammatik verwendet eine Thema-Kommentar-Struktur, zeitliche Indikatoren am Anfang und Fragewörter am Ende, ergänzt durch mimische Ausdrücke (z. B. Augenbrauen).
• 5:14 Bisindo versus SIBI: Vergleich zwischen der offiziellen, akademischen, einhändigen SIBI-Gebärdensprache und der in der Community weit verbreiteten, intuitiveren Bisindo-Gebärdensprache.
• 8:50 Datenerfassung & Motion Grids: Nutzung von 154 Videoclips eines staatlich anerkannten Experten, umgewandelt in 3x3-Montagen (1.848 Trainingsbeispiele), um temporäre Bewegungen in statische Rasterbilder für das Visionsmodell zu überführen.
• 10:31 Modellarchitektur & Training: Einsatz eines 3-Milliarden-Parameter-Modells mit LoRA-Adaptern (Rang 16, Ziel: Attention und MLP, ca. 0,3 % trainierbare Parameter) auf Google Cloud Vertex AI mit Nvidia L4 GPUs, konvergiert in rund 2 Stunden.
• 13:08 Ingestierung durch CLIP Encoder: Verwendung eines multimodalen generativen Visions-Sprachmodells anstelle eines einfachen Klassifizifiers zur Erzielung starker Prä-trainierter Priors.
• 15:48 Validierungsergebnisse: Erreichung einer Genauigkeit von 99,4 % (153 von 154 gehaltenen Grids) bei internen Validierungstests mit demselben Unterzeichner.
• 16:54 Multi-Plattform-Export: Komprimierung des multimodalen Modells zu einer ca. 3 GB großen Datei zur Ausführung auf lokalen Laptops, Web-GPUs (via Worker), Cloud-Endpunkten und nativen Mobilgeräten.
• 17:51 Flutter-Integration: Verwendung einer Flutter-Codebasis mit bedingten Importen, um plattformspezifische Runtimes (LiteRT für Mobile, Web-GPU für Browser) anzusprechen.
• 21:19 Live-Demonstration & Syntaxkorrektur: Echtzeit-Demo der Gebärdenerkennung, bei der das multimodale Modell die Rohzeichen liefert und nachfolgende Gemma-Modelle die Grammatik in korrekte indonesische Sätze umwandeln.

Glossar

• Bisindo (Bahasa Isyarat Indonesia): Die von der Gehörlosengemeinschaft in Indonesien entwickelte, auf visueller Logik basierende Gebärdensprache.
• SIBI (Sistem Isyarat Bahasa Indonesia): Das offizielle, akademische, einhändige System der indonesischen Gebärdensprache, das sich eng am gesprochenen Indonesisch orientiert.
• Motion Grid: Eine Rasterdarstellung (hier 3x3), die mehrere sequenzielle Einzelbilder eines Videos zu einem einzigen kombinierten Bild zusammenfasst, um zeitliche Dynamiken für Bildmodelle darstellbar zu machen.
• PaliGemma: Ein offenes, multimodales Visions-Sprachmodell (Vision-Language Model), das Bild- und Text-Inputs verarbeitet (im Transkript phonetisch als „Pal Jig 2“ oder „Bali Jema“ erfasst).
• LoRA (Low-Rank Adaptation): Eine Methode zum effizienten Fine-Tuning großer neuronaler Netzwerke, bei der nur eine kleine Untermenge an Parametern über Rang-Adapter trainiert wird.
• MLP (Multi-Layer Perceptron): Ein klassischer Typ eines Feedforward-künstlichen neuronalen Netzwerks mit mehreren Schichten.
• Vertex AI: Die verwaltete Machine-Learning-Plattform von Google Cloud zum Trainieren, Evaluieren und Bereitstellen von Modellen.
• Nvidia L4: Eine für Deep-Learning-Inferenz, Grafik und Video-Workloads optimierte Rechen-GPU.
• CLIP Encoder: Ein Encoder-Modell, das visuelle und textuelle Daten in einen gemeinsamen Vektorraum abbildet.
• LiteRT: Ein leichtgewichtogenes Framework zur Ausführung von Modellen auf mobilen und Edge-Geräten (früher TensorFlow Lite).
• Flutter: Ein von Google entwickeltes UI-Toolkit zur Erstellung plattformübergreifender nativer Anwendungen aus einer einzigen Codebasis.

Abstract

Diese Präsentation stellt zwei agentenbasierte Lösungen vor – Max Code und Max Kernel –, die entwickelt wurden, um die Entwicklung und Bereitstellung von Modellen für maschinelles Lernen auf Google-TPUs zu optimieren. Max Code ist ein KI-gesteuertes Migrationstool, das PyTorch-Modelle automatisiert nach JAX übersetzt, um Entwicklungszeiten (Model Bring-up) durch Erreichung hoher numerischer Äquivalenz drastisch zu verkürzen. Max Kernel agiert als KI-gestütztes Framework zur Erstellung und Optimierung von Pallas-Kerneln, um Hardware-Ressourcen optimal auszunutzen. Evaluiert an einem maßgeschneiderten Benchmark von 50 Kerneln, erreicht Max Kernel eine numerische Korrektheit von 96 % und liefert erhebliche Geschwindigkeitsvorteile, wie etwa einen 4,7-fachen Speedup für Qwen-3-Kernel, während es bis zu 60 % der Leistung handoptimierter Kernel erreicht.

Wichtigste Punkte & Zeitstempel

• 0:00 Einführung zu ML-Entwicklungsagenten: Die Moderatoren Deepak und Sithu präsentieren agentenbasierte Ansätze zur Bewältigung von Herausforderungen bei Training und Inferenz mithilfe von Max Code und Max Kernel.
• 0:35 Herausforderungen im ML-Workload: Entwickler leiden unter langwierigen Modellkonvertierungen von PyTorch nach JAX sowie unter dem extremen Zeit- und Fachwissenaufwand für die Erstellung von Pallas-Kerneln zur Erreichung von Spitzenleistungen.
• 1:37 JAX- und MaxText-Technologie-Stack: Überblick über den Software-Stack, bestehend aus dem JAX-Framework, Pallas für benutzerdefinierte TPU-Kernel, dem Compiler XLA sowie Hilfsbibliotheken wie Orbax (Checkpoints) und MaxText (LLM-Referenzimplementierung).
• 2:27 Typischer ML-Entwicklungszyklus: Der Ablauf umfasst Modell- und Checkpoint-Konvertierung, Roofline-Analysen zur Bestimmung des Hardware-Limits, Profiling via XPRO und zielgerichtete Optimierungen bezüglich Sharding und Quantisierung.
• 5:07 Max Code (PyTorch-zu-JAX-Migration): Ein agentenbasierter Code-Migrationsexperte, der den Übergang von PyTorch zu JAX automatisiert, um die Lernkurve zu glätten und die Bereitstellung neuer Modellarchitekturen zu beschleunigen.
• 7:09 Vierstufige Max-Code-Architektur: Nutzt einen Orchestrierungsagenten für die Abhängigkeitsanalyse, einen Code-Generierungsagenten (gestützt auf kuratierte MaxText-Beispiele), einen Evaluierungsagenten sowie einen Selbst-Debugging-Agenten.
• 11:04 Max Kernel (TPU-Kernel-Optimierung): Ein Framework mit fünf spezialisierten Sub-Agenten (Planung, Generierung, Testen, Tuning, Profiling), das Pallas-Code für TPUs anstelle von CUDA optimiert.
• 12:18 Jaxbench-Evaluierungs-Benchmark: Ein eigens entwickelter Benchmark aus 50 Kerneln (33 aus KernelBench, 17 aus MaxText) zur präzisen Messung der Agentenleistung an anspruchsvollen Betriebspunkten.
• 14:31 Iteratives Refinement und Leistungssteigerung: Durch die Integration eines Profilers in eine iterative Feedbackschleife verdoppelt sich der geometrische Mittelwert der Leistungssteigerung von 18 % auf 35 %, bei einer Korrektheit von 96 % (48/50 Kerneln).
• 15:42 Auto-Optimierung und Suchstrategien: Ein compiler-inspirierter Beam-Search-Ansatz führt granulare Optimierungsschritte aus und erzielt bis zu 60 % der Leistung manuell getunter Referenzkernel.
• 17:13 Reale Leistungsmetriken: Die Anwendung zeigt signifikante Produktivitäts- und Leistungsgewinne, darunter einen bis zu 4,7-fachen Speedup für Gated-Delta-Kernel (Qwen 3) und rund 10 % Durchsatzverbesserung bei MLA (Multi-Head Latent Attention).
• 18:01 Open-Source-Verfügbarkeit: Die Repositories für Max Code, Max Kernel und Jaxbench sind öffentlich auf GitHub verfügbar, wobei ein starker Fokus auf numerischer Korrektheit und Hardware-Awareness liegt.

Glossar

• JAX: Ein Hochleistungs-Framework für numerische Berechnungen und maschinelles Lernen in Python, das automatische Differenzierung und JIT-Kompilierung unterstützt.
• PyTorch: Ein weit verbreitetes Open-Source-Framework für maschinelles Lernen und Deep Learning, das primär für Forschung und Produktion eingesetzt wird.
• TPU (Tensor Processing Unit): Ein von Google entwickelter, anwendungsspezifischer integrierter Schaltkreis (ASIC), der speziell zur Beschleunigung von neuronalen Netzen und KI-Workloads entwickelt wurde.
• Pallas: Eine domain-spezifische Programmierschnittstelle innerhalb von JAX, die es Entwicklern ermöglicht, hocheffiziente, maßgeschneiderte Hardware-Kernel für TPUs zu schreiben.
• XLA (Accelerated Linear Algebra): Ein domänenspezifischer Compiler, der lineare Algebra-Graphen optimiert und kompiliert, um Code auf CPUs, GPUs und TPUs effizient auszuführen.
• MFU (Model Flops Utilization): Eine Leistungskennzahl, die angibt, wie viel Prozent der maximal theoretischen Rechenleistung (FLOPS) der Hardware während des Modelltrainings tatsächlich genutzt werden.
• Roofline-Analyse: Ein analytisches Leistungsmodell, das die Obergrenze der erreichbaren Rechenleistung basierend auf Rechenintensität und Speicherbandbreite visualisiert.
• Pytorch Oracle: Eine als unveränderliche Wahrheit dienende Referenzimplementierung in PyTorch, gegen die der generierte JAX-Code auf numerische Korrektheit geprüft wird.
• KernelBench / Jaxbench: Benchmark-Suiten zur standardisierten Bewertung der Korrektheit und Ausführungsgeschwindigkeit von manuell oder KI-generierten Kerneln.
• MLA (Multi-Head Latent Attention): Eine optimierte Attention-Architektur in modernen Sprachmodellen, die den Speicherbedarf für den Key-Value-Cache drastisch reduziert.