Most LLM frameworks you meet on this blog start from PyTorch and bend toward accelerators. MaxText starts from the opposite shore: it is a pure Python / JAX reference implementation built first for Google Cloud TPUs and GPUs, and it treats the compiler — XLA — as the place where performance comes from. There is no hand-written CUDA everywhere, no bespoke scheduler; instead there is a very carefully organized set of model definitions, a declarative sharding layer, and a training loop that trusts jax.jit to fuse everything. If gemma-cpp showed us what it takes to squeeze a Gemma-class model into a laptop CPU, MaxText is the other pole: the same family of models expressed so they can scale to tens of thousands of chips.
The project recently went through a structural reorganization (tracked in its RESTRUCTURE.md), so the layout we tour today is the modern one: the Python package lives under src/maxtext/, with subfolders like models/, layers/, trainers/pre_train/, common/, utils/, and inference/. Old tutorials that reference a top-level MaxText/ directory are now stale; the module you import is maxtext, and entry points are python3 -m maxtext.trainers.pre_train.train for training and python3 -m maxtext.inference.decode for sampling. Keeping that map in your head makes the rest of the tour much easier.
What makes MaxText worth a source tour is how few moving parts it needs to hit state-of-the-art model FLOPs utilization. Configs are YAML files. Sharding is described with logical axis names that get resolved against a physical device mesh. The training step is one function. Checkpointing delegates to Orbax. Serving is an engine class with four verbs — load_params, prefill, insert, generate — that JetStream-style servers call. Let’s walk through the whole flow with two diagrams.
The overview graph: pyconfig builds HyperParameters, the training loop wraps a jitted train_step, logical sharding resolves against a physical device mesh, the Transformer decomposes into decoder/attention/MLP layers, and MaxEngine serves decode on the same model.
Reading the overview from left to right:
- pyconfig.initialize, L415 — loads
base.yml, a model YAML, then CLI/env overrides, and hands back a validatedHyperParametersobject. - create_device_mesh, L2028 — turns
ici_parallelism/dcn_parallelismnumbers into a concretejax.Meshover your devices. - create_sharding, L489 — resolves logical names like
activation_batchandembedinto aNamedShardingon that mesh. - train_step, L499 — the jitted heart of pre-training: value-and-grad over
loss_fn, gradient casting, optimizer update. - Transformer, L108 — a Flax NNX module wiring an
Embedtable to anNNXDecoderstack. - NNXDecoder, L549 — scans layers in a loop or splits them across mesh stages for staged data flow.
- Attention, L93 — QKV projection, RoPE, and a dispatch table of fused kernels.
- MaxEngine, L110 — the serving engine:
load_params,prefill,insert,generate. - main, L89 — the CLI decode loop that exercises all of the above in one process.
The Config Machine
Everything in MaxText begins with pyconfig. The function pyconfig.initialize, L415 is a thin wrapper around _initialize_pydantic, which does something subtly clever: it loads src/maxtext/configs/base.yml (about a thousand keys of documented defaults), then merges a config path — either passed explicitly or inferred from model_name — and finally applies CLI overrides via OmegaConf. When you type model_name=llama3-8b, MaxText finds src/maxtext/configs/models/llama3-8b.yml, which itself inherits from base.yml, and layers the deltas on top.
Validation happens through a Pydantic schema: MaxTextConfig, L4189 in configs/types.py declares every knob with a type, a default, and (often) a comment about units. The runtime object, HyperParameters at L333, wraps that Pydantic model and forwards attribute access, so config.per_device_batch_size just works. This design has a pleasant property: the same code path serves training, compilation, and inference, and every one of them sees an identical, already-validated view of the world. The cost is sprawl — base.yml is enormous, and a surprising fraction of onboarding is learning which of the hundred-plus keys actually matter for your run.
Sharding: Physical Meshes, Logical Names
MaxText’s sharding story is the part most worth stealing for your own JAX projects. The entry point is create_device_mesh, L2028: given ici_parallelism (how axes split inside a TPU slice) and dcn_parallelism (how slices split across the data center network), it builds a physical mesh, filling in any unspecified axes automatically. Multi-slice runs get a hybrid mesh; single slices can opt into split physical axes with allow_split_physical_axes.
On top of that physical mesh sits a logical vocabulary. Layers never write PartitionSpec("data", "tensor"); they write names like activation_batch, activation_norm_length, activation_embed, activation_mlp. The resolver logical_to_mesh_axes, L452 maps those logical names through logical_axis_rules to real mesh axes, and create_sharding, L489 returns a ready-to-use NamedSharding. Because the mapping is data, you can re-target a model from an 8-chip host to a 3-dimensional megascale cluster by editing rules rather than touching layer code. It is the same philosophical bet as ZeRO-style state sharding in deepspeed — push partitioning decisions into a declarative layer — but expressed natively in JAX’s sharding semantics rather than as a PyTorch plugin.
Inside the Training Step
The training entry point trainers/pre_train/train.py is refreshingly linear: main parses argv, initialize builds the config, and train_loop at L1192 calls setup_train_loop, L358 to assemble the mesh, optimizer, sharded state, checkpoint manager, and the jitted step function via jit_train_step, L235. State lives in TrainStateNNX, L26, which unifies the Flax NNX model state and the Optax optimizer state in one module so a single jit boundary captures everything.
The step itself, train_step, L499, splits the NNX model into params and rest, computes jax.value_and_grad over loss_fn, L133 — a shared forward used by both train and eval — casts gradients to grad_dtype, applies gradient accumulation over microbatches when configured, and hands the result to the optimizer. MoE models get extra machinery: router-bias updates, overflow rollback, and auxiliary load-balancing losses threaded through the same function.
Durability is a first-class concern at ten-thousand-chip scale, and it shows. save_checkpoint, L1565 wraps Orbax’s async CheckpointManager — created earlier by create_orbax_checkpoint_manager — blocking only when metrics require it, saving in a layout that keeps NNX and legacy Linen checkpoints interchangeable. On restore, load_state_if_possible, L730 implements the priority chain you can read at the top of base.yml: resume the latest full checkpoint if one exists, else load_parameters_path, else load_full_state_path, else initialize fresh. Preemption is expected, not exceptional.
Serving: Prefill, Insert, Generate
Serving reuses the exact same model code through a different facade. MaxEngine in inference/maxengine/maxengine.py creates its own device mesh at L121, loads or converts params, and exposes the JetStream-compatible verb set. A request flows through three phases: prefill, L774 runs the prompt through the model in one shot and returns the first token plus a populated cache; insert, L1505 places that result into a decode slot; then generate, L1136 steps the whole batch one token per call, with the heavy function jitted and the cache donated in place.
The KV cache itself is worth a look: KVCache, L153 manages paged, layer-stacked arrays with configurable axis orders for prefill versus autoregressive mode, plus quantization hooks for the cache. Around the engine sit several shells: maxengine_server.py wraps it in a gRPC model server, offline_engine.py builds a threaded worker pool around the same verbs for bulk jobs, and vllm_decode.py connects MaxText to vLLM so the sampler and scheduler come from there while JAX computes the forward pass. The CLI script inference/decode.py wires it all together for a single stream — encode the prompt, prefill, insert, generate in a loop, decode back to text, and assert the output matches autoregressive_decode_assert so the end-to-end path is testable.
The layer zoo underneath serving is shared with training: Embed, L60 covers the vocab table and logit head, while the RotaryEmbedding family at L231 implements RoPE and its Yaarn-style variants for different model families. Attention kernels are pluggable — the attention config selects among fused implementations, several of which trace back to the same ideas popularized by flash-attention — and quantization (AqtQuantization for weights/activations, separate kv-quant) is threaded through the projections. Converting HuggingFace checkpoints into Orbax format, documented in checkpoint_conversion/, closes the loop with transformers so you can train here and evaluate elsewhere.
From Install to Your First Model
MaxText is Linux-first and Python 3.12-first, and the docs now recommend uv throughout. From a fresh clone of the repository:
git clone https://github.com/google/maxtext.git
cd maxtext
uv venv --python 3.12 --seed maxtext-env
source maxtext-env/bin/activate
# from-source install (TPU flavor; use .[cuda12] for GPUs)
uv pip install -e .[tpu] --resolution=lowest
install_tpu_pre_train_extra_deps
# verify the environment
uv pip check
python3 -c "import maxtext"
python3 -m maxtext.trainers.pre_train.train --help
Two details matter. The --resolution=lowest flag is not decoration: it pins every dependency to the exact versions the team tested, which is the practical answer to the JAX/jaxlib version-matrix problem on accelerators. And the install_*_extra_deps entry points fetch the accelerator-specific extras (a --with-tf flag adds optional TensorFlow/TFDS and JetStream dependencies) that are deliberately kept out of the base wheel.
With the environment healthy, the smallest real run is a synthetic-data smoke test — no dataset, no checkpoint, ten steps:
python3 -m maxtext.trainers.pre_train.train \
model_name=llama3-8b \
run_name=my-first-run \
base_output_directory=gs://YOUR-BUCKET/maxtext-outputs \
dataset_type=synthetic \
steps=10
On a GPU box you’d instead pass the GPU YAML as the leading positional config, e.g. src/maxtext/configs/models/gpu/llama3-8b.yml. If you just want to poke at the config layer from Python, this snippet uses only verified API surface:
from maxtext.configs import pyconfig
config = pyconfig.initialize([
"prog", # argv[0] placeholder
"model_name=llama3-8b", # inferred model config
"run_name=smoke",
"steps=10",
])
print(config.base_emb_dim, config.steps) # 4096 10
Finally, generate text from a (randomly initialized) model to exercise the serving path: python3 -m maxtext.inference.decode run_name=my-first-run base_output_directory=... per_device_batch_size=1. The output will be gibberish without a converted checkpoint via load_parameters_path, but it proves the whole prefill/insert/generate loop compiles and runs.
Honest Limits
The trade-offs are real and the maintainers don’t hide them. First, this is TPU-first software: the happy path, the benchmark numbers, and most of the docs assume Google Cloud infrastructure — GCS buckets for output, Enjoyed this post? Never miss out on future posts by following us gcloud tooling for workloads. GPU support is genuine (the cuda12 extra, GPU configs, FSDP-style prefetch variants), but a “decoupled mode” environment variable dance exists precisely because GCP assumptions are baked in. Second, config sprawl is the tax for flexibility: with a Pydantic schema of well over a thousand fields, finding which five knobs control your throughput is an archaeology exercise, and the restructure that moved MaxText/ to src/maxtext/ means a lot of blog posts and Stack Overflow answers from 2024 and earlier point at ghosts. Third, the whole design presumes multi-host orchestration — XLA wants to see the full device mesh, so single-node experimentation with large models is awkward, and running real workloads means XPK, Cluster Toolkit, or Pathways-style setups rather than a bare torchrun. Fourth, dependency pinning cuts both ways: --resolution=lowest buys reproducibility at the price of staying off the bleeding edge, and mixing JAX/jaxlib versions from different extras can break your environment in ways that are easier to nuke-and-recreate than to debug. None of these are disqualifying — they are the honest cost of a codebase that treats “train a frontier model on a pod of TPUs” as the normal case. For everyone else, it remains one of the cleanest blueprints for reading how compiler-driven LLM training is actually organized. 
