MaxText README

MaxText is Google's reference JAX LLM library and was JetStream's JAX engine implementation. Facts: written in pure Python/JAX, targets Cloud TPU AND GPUs, claims high MFU from single host to tens of thousands of chips while staying 'largely optimization-free thanks to the power of JAX and the XLA compiler'. Models: Gemma, Llama, DeepSeek, Qwen, Mistral; pre-training plus post-training (SFT, GRPO, GSPO). PyPI-installable as maxtext, with an explicit warning that main is NOT expected to be production-ready and that Python 3.12 is the primary supported version. Recent TPU-specific work in the news section: Streaming DiLoCo distributed training across multi-slice TPU clusters on low-bandwidth DCN/WAN (2026-08-19), DeepSeek V4 Flash 284B support (2026-08-05), Tokamax GMM v2 MoE kernel (2026-08-03). Has a 'decoupled mode' for running without any GCP dependencies. Repo metadata via GitHub API 2026-10-04: created 2023-02-28, pushed 2026-10-04, 2436 stars, 371 open issues, Apache-2.0.