MaxText im Detail: JAX-Training von Gemma 4 bis DeepSeek V4 Flash

Warum MaxText die Optimierung an XLA auslagert

MaxText verfolgt für ein Trainingsframework dieser Größe einen ungewöhnlichen Ansatz. Der Code besteht laut README aus reinem Python und JAX, handgeschriebene Kernel sucht man darin nicht. Die Optimierung überlässt das Projekt dem XLA-Compiler, weshalb sich die README selbst als weitgehend optimierungsfrei beschreibt. Als Ergebnis nennt sie eine hohe Model FLOPs Utilization und einen hohen Tokendurchsatz, vom einzelnen Host bis zu sehr großen Clustern. Konkrete Zahlen stehen nicht im Text, die Angabe ist damit eine Selbstauskunft, die sich erst mit einem eigenen Lauf prüfen lässt.

Zielplattformen sind Google Cloud TPUs und GPUs. Das Repository liegt bei 2407 Sternen und 593 Forks, bei 310 offenen Issues, und ist nicht archiviert. Die README empfiehlt ausdrücklich, MaxText zunächst unverändert auszuprobieren und danach zu forken und anzupassen. Genau dafür ist die Referenzimplementierung gedacht: Sie soll als Ausgangspunkt für eigene Projekte dienen, nicht als fertiges Produkt mit Supportzusage. Die Apache License 2.0 erlaubt Forks und kommerzielle Nutzung, enthält aber keine Gewährleistung.

Die Modellbibliothek: gemma4-26b.yml, deepseek3.2-671b.yml und kimi-k2-1t.yml

Die Modellauswahl ist der Teil von MaxText, der am schnellsten wächst. In src/maxtext/configs/models/ liegen die Konfigurationsdateien, unter anderem gemma4-26b.yml, gemma4-31b.yml, deepseek3.2-671b.yml und kimi-k2-1t.yml. Gemma 4 kam am 2. April 2026 hinzu, mit einem 26B-MoE-Modell und einer 31B-Variante als dichtes Modell, beide multimodal. DeepSeek ist über mehrere Generationen vertreten: V2 in 16B und 236B, dazu V3, V3.1, V3.2 mit 671B und R1 0528 mit 671B. Am 5. August 2026 folgte DeepSeek V4 Flash mit 284B Parametern, inklusive Hinweis auf Checkpoint-Konvertierung und Evaluierung in Run_DeepSeek.md.

Aus dem Qwen-Umfeld listet das Projekt Qwen 2.5, Qwen 3 Dense in sechs Größen von 0.6B bis 32B, Qwen 3 MoE mit 30B und 235B, Qwen 3 MoE 2507 mit 235B und 480B, Qwen 3 Next mit 80B sowie seit dem 9. Juni 2026 Qwen 3.5 mit 35B und 397B. Moonshot AI steuert Kimi K2 mit 1T Parametern bei, samt MuonClip-Optimierer, dazu K2-Thinking und die Textmodelle K2.5 und K2.6. Llama ist von Llama 2 über Llama 3 mit 405B bis zu Llama 4 Scout und Mistral ist mit 7B sowie Mixtral 8x7B und 8x22B vertreten.

Pre-Training bis zehntausende Chips, Post-Training mit SFT, GRPO und GSPO

Für das Vortraining gibt die README eine Obergrenze von bis zu Zehntausenden von Chips an. Der Nachtrainingspfad ist feiner ausgearbeitet: Supervised Fine-Tuning ist dokumentiert, dazu zwei Varianten des bestärkenden Lernens, Group Relative Policy Optimization und Group Sequence Policy Optimization. Die Pipeline nutzt Tunix für das RL-Training und vLLM für das Sampling. Für Multi-Host-Szenarien wird Pathways erwähnt, und zwar ausdrücklich als Zukunftsperspektive und nicht als fertiger Pfad.

Seit dem 15. Juli 2026 ist RL für Qwen3 30B und GPT-OSS 20B freigegeben, mit eigenen Tutorien unter rl_qwen3_30b.html und rl_gptoss_20b.html auf Read the Docs. Wer sich für Prompt-Optimierung interessiert, findet mit maxtext_with_gepa.ipynb ein Notebook, das AIME-Prompts mit dem GEPA-Framework optimiert. Multimodales Training ist für die VLM-Varianten von Gemma 3, Gemma 4 und Llama 4 vorgesehen. Wie sich die einzelnen Post-Training-Verfahren auf die Modellqualität auswirken, beziffert die README nicht.

Flax NNX statt Linen: was die Migration vom 29. Juni 2026 bedeutet

Zum 29. Juni 2026 ist MaxText von Flax Linen auf Flax NNX umgestellt. Die README formuliert die Erwartung, dass bestehende Workloads weiterlaufen wie bisher, und bittet um Fehlermeldungen über die Vorlage bug_report.yml. Für neue Nutzer ist der Wechsel unsichtbar. Wer einen älteren Fork pflegt und eigene Module geschrieben hat, sollte die eigenen Erweiterungen gegen die NNX-API prüfen.

Kurze Zeit vorher, am 14. April 2026, wurden die alten Post-Training-Shims mit dem Namenspräfix MaxText.* entfernt. Die neuen Befehlspositionen stehen in src/MaxText/README.md. Beide Änderungen zusammen zeigen, dass sich die Verzeichnisstruktur und die Aufrufwege innerhalb kurzer Zeit verschoben haben. Skripte, die MaxText-Befehle mit festen Pfaden aufrufen, sind davon betroffen, und die Migrationshinweise in src/MaxText/README.md sind dafür die maßgebliche Quelle.

Installation: pip aus PyPI, Python 3.12 und der Decoupled Mode ohne GCP

Der Installationsweg führt über pip aus PyPI, ergänzt durch vorgefertigte Container-Images, die in docs/tutorials/build_maxtext.md beschrieben sind. Die README rät Kunden ausdrücklich zum jeweils neuesten PyPI-Release und warnt, dass der main-Branch nicht als produktionsreif gilt. Als primär unterstützte Python-Version nennt das Projekt 3.12, andere Versionen können laut README auf Kompatibilitätsprobleme stoßen. Wer in einer bestehenden Umgebung arbeitet, führt pip install -e . im Wurzelverzeichnis aus, seit das Repository auf ein src-Layout umgestellt wurde. Für das Paket existiert zusätzlich das Extra tpu-post-train, das für vllm_decode gebraucht wird.

Eine Besonderheit ist der Decoupled Mode, den docs auf Read the Docs unter run_maxtext/decoupled_mode.html beschreibt. Er erlaubt den Betrieb ohne GCP-Abhängigkeiten. Wie weit dieser Modus funktioniert und welche Funktionen dabei wegfallen, steht nicht im README, sondern nur im verlinkten Leitfaden. Für lokale Experimente auf einer einzelnen GPU ist das der Einstieg, der ohne Cloud-Projekt auskommt.

Tokamax GMM v2 für MoE und das Eval-Framework für lm-eval

Am 3. August 2026 wurde Tokamax GMM v2 als MoE-Kernel verfügbar. Aktiviert wird er über die Einstellungen, die docs/reference/core_concepts/moe_configuration.md beschreibt. Für Mixture-of-Experts-Modelle wie DeepSeek V3 oder Qwen 3 MoE ist das der Teil der Konfiguration, der am ehesten über Durchsatz und Speicherverbrauch bestimmt.

Seit dem 29. Mai 2026 existiert ein Evaluierungsframework, das lm-eval, evalchemy und eigene Benchmarks gegen MaxText-Checkpoints laufen lässt. Der Einstieg liegt in guides/eval_framework.html. Damit lässt sich ein trainierter Checkpoint bewerten, ohne die Auswertung selbst zusammenzubauen. Welche Benchmark-Suites vorkonfiguriert sind und wie die Checkpoint-Konvertierung für die jeweilige Modellfamilie aussieht, ist in der README nicht aufgeführt.

Grenzen von MaxText und wann MaxDiffusion die passendere Adresse ist

Die deutlichste Einschränkung steht in der README selbst: Empfohlen werden der PyPI-Release und die Container-Images, der main-Branch gilt nicht als produktionsreif. Dazu kommen 310 offene Issues und die Bindung an Google Cloud TPUs und GPUs als Zielplattform. Wer auf Hardware anderer Anbieter trainiert, findet in der Dokumentation keine Angaben. Auch die MFU-Angaben bleiben ohne Zahlen, und die README empfiehlt Python 3.12, was in Umgebungen mit älteren Python-Versionen zu Nacharbeit führt.

Wer keine Sprachmodelle, sondern Diffusionsmodelle trainieren will, ist bei MaxText an der falschen Adresse. Das Projekt verweist für Wan 2.1 und Flux auf das Schwester-Repository AI-Hypercomputer/maxdiffusion. Der Unterschied liegt im Workload, nicht in der Qualität: MaxDiffusion deckt Bild- und Videogenerierung ab, MaxText konzentriert sich auf Textmodelle inklusive multimodaler VLM-Varianten von Gemma und Llama. Für reine Text-Pipelines führt der Weg über MaxText, für Diffusionsmodelle über MaxDiffusion.

Redaktionelles Fazit

MaxText passt zu Teams, die auf Google Cloud TPUs trainieren und eine lesbare Codebasis zum Forken suchen: reines Python plus JAX, Konfigurationen über yml-Dateien in src/maxtext/configs/models/, Post-Training mit Tunix und vLLM. Wer eine stabile Produktionsplattform ohne eigene Nacharbeit erwartet, sollte stattdessen den empfohlenen PyPI-Release maxtext-v0.2.4 und die vorgefertigten Container-Images nutzen und den main-Branch meiden. Vor dem Einstieg klären: Läuft die eigene Python-Version (empfohlen 3.12), ist der Zielchip TPU oder GPU, und existiert für das Wunschmodell bereits eine Konfigurationsdatei oder ein Run_-Leitfaden unter tests/end_to_end/tpu/.

Offizielle Quellen

1. Official documentation
2. Official README
3. Project repository
4. Release notes