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.