MaxText
An open-source LLM training library using Python and JAX, with pre-training and fine-tuning on TPUs and GPUs and a mode without GCP dependencies.
MaxText is an LLM training library for researchers and engineers building models from scratch or adapting existing ones. It provides a reference implementation teams can modify for their own model designs, with training that scales from a single host to large clusters.
- Model training: Train dense and mixture-of-experts models from families including Gemma, Llama, DeepSeek, Qwen, Mistral, Mixtral and GPT-OSS.
- Fine-tuning and reinforcement learning: Adapt proprietary or open-source models through supervised fine-tuning, knowledge distillation and reinforcement learning methods such as GRPO and GSPO, using Tunix for post-training.
- Multimodal models: Train vision-language models from the Gemma 3, Gemma 4 and Llama 4 families; Kimi-K2.5 and Kimi-K2.6 support is text-only.
- Checkpoint evaluation: Run lm-eval, evalchemy or custom benchmarks against MaxText checkpoints to compare model results.
MaxText uses Python and JAX with the XLA compiler. It targets Google Cloud TPUs and GPUs, while its decoupled mode removes GCP dependencies. Reinforcement learning sampling uses vLLM on TPU. The project is open source under Apache 2.0 and distributes Python packages and prebuilt container images.
