MaxText Documentation — Scalable LLM Pre-Training and Post-Training in JAX
by Google (AI Hypercomputer)
Google's reference JAX codebase and docs for pre-training, fine-tuning and RL on Gemma, Qwen, DeepSeek, Kimi and Llama at TPU scale.
Overview
MaxText is the documentation for AI-Hypercomputer/maxtext, Google's high-performance, open-source LLM library written in pure Python/JAX for Cloud TPUs and GPUs (Apache-2.0, about 2.4k GitHub stars). The docs are split into getting started, install, tutorials, run guides, how-to guides, reference, development and release notes. Six tutorials cover building and uploading Docker images, pre-training on real datasets such as C4 through Hugging Face, Grain or TFDS, DiLoCo and Streaming DiLoCo multi-slice low-communication training, post-training (SFT and RL on TPU), inference through vLLM, and decision models that score allowed answers using a Qwen3.5-9B LoRA example. Nine guides cover optimisation (sharding strategies, Pallas kernels, benchmarking), data pipelines, checkpointing including emergency checkpoints for preemption, goodput monitoring and hung-job debugging, notebooks on TPUs, community model bring-up, adding LoRA support for a new architecture, online distillation, and a vLLM-based eval framework running lm-eval and evalchemy against checkpoints. Supported models span Gemma 1 to Gemma 4, Qwen 2.5 to Qwen 3.5 MoE, DeepSeek V2 to V3.2 and V4 Flash, Kimi K2 to K2.6 text, Llama 2 to Llama 4, GPT-OSS and Mixtral. The project is moving fast: the news log runs from Gemma 4 support in April 2026 through the Flax NNX migration in June, RL for Qwen3 30B and GPT-OSS 20B in July, and decision scoring on vLLM on October 8, 2026.
At a Glance
- Topic
- Models
- Level
- Advanced
- Format
- Documentation
- Cost
- Free
- Duration
- ~half a day for install and a first training run; one to two weeks to stand up a multi-slice pre-training or RL job
- Provider
- Google (AI Hypercomputer)
- Hands-on
- Yes — code/exercises
- Certificate
- None
What You’ll Learn
- ✓Build MaxText Docker images and launch a pre-training run on C4 with Grain, Hugging Face or TFDS input
- ✓Pick sharding strategies and Pallas kernels to raise MFU and benchmark training throughput
- ✓Run multi-slice training with DiLoCo and Streaming DiLoCo to cut cross-slice communication
- ✓Fine-tune and run RL post-training on TPU for models such as Qwen3 30B and GPT-OSS 20B
- ✓Configure checkpointing, emergency checkpoints and multi-tier storage so preempted jobs resume
- ✓Monitor goodput and diagnose hung multi-host jobs with the documented debugging tools
- ✓Serve MaxText checkpoints through vLLM and evaluate them with lm-eval and evalchemy
- ✓Bring up a new model architecture or add LoRA support by following the contributor guides
Highlights
- •One codebase covers pre-training, SFT, RL, distillation, inference via vLLM and evaluation for frontier-scale open models
- •Unusually broad model coverage, including DeepSeek V3.2 671B, Kimi K2 1T and Qwen 3.5 397B MoE
- •Production concerns are documented: emergency checkpointing, goodput monitoring and hung-job debugging
- •Very active: the dated news log shows new models and features every few weeks through October 2026
Who It’s For
Best For
- ✓ML infrastructure engineers pre-training or continued-pre-training LLMs on Cloud TPU pods
- ✓Teams running RL or distillation on large MoE models who need a reference JAX implementation
- ✓Engineers pairing MaxText with Tunix or vLLM for an end-to-end JAX training and serving stack
Prerequisites
- •Strong Python and working JAX knowledge, including sharding and jit compilation
- •Understanding of transformer training: data parallelism, model parallelism and mixed precision
- •A Google Cloud project with TPU or GPU quota for anything beyond the smallest runs
FAQ
What is MaxText Documentation — Scalable LLM Pre-Training and Post-Training in JAX?
MaxText is Google's open-source documentation and reference implementation for training large language models in pure Python and JAX on Cloud TPUs and GPUs. It is for ML engineers running pre-training, SFT, RL or distillation at multi-host scale, who will finish able to launch, shard, checkpoint and evaluate a model such as Gemma 4 or Qwen 3.5.
Is MaxText Documentation — Scalable LLM Pre-Training and Post-Training in JAX free?
MaxText Documentation — Scalable LLM Pre-Training and Post-Training in JAX is free to access.
What level is MaxText Documentation — Scalable LLM Pre-Training and Post-Training in JAX for?
MaxText Documentation — Scalable LLM Pre-Training and Post-Training in JAX is aimed at a advanced audience. Recommended background: Strong Python and working JAX knowledge, including sharding and jit compilation, Understanding of transformer training: data parallelism, model parallelism and mixed precision, A Google Cloud project with TPU or GPU quota for anything beyond the smallest runs.
How long does MaxText Documentation — Scalable LLM Pre-Training and Post-Training in JAX take?
Expect roughly ~half a day for install and a first training run; one to two weeks to stand up a multi-slice pre-training or RL job. Most learners work through it at their own pace.
What will I learn from MaxText Documentation — Scalable LLM Pre-Training and Post-Training in JAX?
You'll learn: Build MaxText Docker images and launch a pre-training run on C4 with Grain, Hugging Face or TFDS input; Pick sharding strategies and Pallas kernels to raise MFU and benchmark training throughput; Run multi-slice training with DiLoCo and Streaming DiLoCo to cut cross-slice communication; Fine-tune and run RL post-training on TPU for models such as Qwen3 30B and GPT-OSS 20B; Configure checkpointing, emergency checkpoints and multi-tier storage so preempted jobs resume; Monitor goodput and diagnose hung multi-host jobs with the documented debugging tools; Serve MaxText checkpoints through vLLM and evaluate them with lm-eval and evalchemy; Bring up a new model architecture or add LoRA support by following the contributor guides.
Topics
Sources
This page was written from 4 sources, 1 on domains other than maxtext.readthedocs.io.