MaxText
A high-performance, open-source large language model library and reference implementation written in pure Python and JAX, aimed at training and fine-tuning models such as Gemma, Llama, DeepSeek, Qwen and Mistral at scale on Google Cloud TPUs and GPUs.
MaxText is an open-source LLM library and reference implementation written in pure Python and JAX, targeting Google Cloud TPUs and GPUs for training. It provides a library of high-performance model implementations, including Gemma, Llama, DeepSeek, Qwen and Mistral, each supporting pre-training at the scale of tens of thousands of chips and post-training with techniques such as supervised fine-tuning, GRPO and GSPO reinforcement learning.
It aims for high model FLOPs utilization from a single host up to very large clusters while staying simple, relying on JAX and the XLA compiler rather than manual optimization. It is intended as a starting point for both research and production LLM projects, meant to be forked and modified, and can also run in a decoupled mode without Google Cloud dependencies. A companion project, MaxDiffusion, covers diffusion models such as Wan 2.1 and Flux.