MaxText

Written in pure Python and JAX, MaxText is Google's open reference framework for training LLMs on Cloud TPUs and GPUs, part of the AI Hypercomputer stack. It leans on the XLA compiler rather than hand-written kernels to reach high model FLOPs utilization, and Google has demonstrated pretraining runs spanning tens of thousands of chips. Model coverage includes the Gemma, Llama, DeepSeek, Qwen, Mistral, Mixtral, and GPT-OSS families along with Moonshot Kimi models, with both pretraining and post-training paths: supervised fine-tuning recipes plus GRPO and GSPO reinforcement learning, including RL for Qwen3-scale MoE models in recent releases. The framework installs from PyPI, configuration is YAML driven, and documentation at maxtext.readthedocs.io walks through single-host getting-started runs up to multi-pod cluster jobs. Everything is Apache-2.0 licensed. Around 2,400 GitHub stars understate its influence: it serves as the canonical demonstration of JAX scaling practice and a starting point that teams fork to build their own TPU or GPU training stacks.