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.