maxtext

Why we included this project

MaxText is a good starting point if you want to train or fine-tune large language models in JAX and would rather not build the whole pipeline yourself. It ships working reference implementations for several open models, including Llama, Gemma, DeepSeek, Qwen, and Mistral, and covers both pre-training and post-training work such as SFT and reinforcement learning, scaling up to clusters of tens of thousands of chips. Because it leans on the JAX/XLA compiler instead of hand-tuned kernels, it gets strong hardware utilization while staying plain Python that is easy to read and modify. It is designed around Google Cloud TPUs and also runs on NVIDIA GPUs, so teams already on Google Cloud infrastructure will feel at home. Since it doubles as a reference implementation, it is also a practical way to study how production-scale LLM training is actually structured before building your own.