MaxText

Documentation | Release Notes

MaxText is a high performance, highly scalable, open-source LLM library and reference implementation written in pure Python/JAX and targeting Google Cloud TPUs and GPUs for training.

MaxText provides a library of high performance models to choose from, including Gemma, Llama, DeepSeek, Qwen, and Mistral. For each of these models, MaxText supports pre-training (up to tens of thousands of chips) and scalable post-training, with popular techniques like Supervised Fine-Tuning (SFT) and Group Relative Policy Optimization (GRPO, a type of Reinforcement Learning) and Group Sequence Policy Optimization (GSPO, a type of Reinforcement Learning).

MaxText achieves high Model FLOPs Utilization (MFU) and tokens/second from single host to very large clusters while staying simple and largely "optimization-free" thanks to the power of JAX and the XLA compiler.

MaxText is the launching point for ambitious LLM projects both in research and production. We encourage you to start by experimenting with MaxText out of the box and then fork and modify MaxText to meet your needs.

Check out our Read The Docs site or directly Get Started with your first MaxText run. If you’re interested in Diffusion models (Wan 2.1, Flux, etc), see the MaxDiffusion repository in our AI Hypercomputer GitHub organization.

Installation

See our installation guide to install MaxText with pip from PyPI.

NOTE: We recommend customers use the latest PyPI release (see installation guide) or corresponding pre-built container images, while the main branch is not expected to be production-ready.

We recommend running MaxText with Python 3.12, as it is our primary supported version. Other Python versions may encounter compatibility issues.

Decoupled mode

See our guide on running MaxText in decoupled mode, without any GCP dependencies in Decoupled Mode Guide.

This is the beginning of the README. Read the rest at the original. Read the full README ↗