MaxText Documentation

MaxText provides scalable LLM (Large Language Model) and VLM (Vision Language Model) post-training on TPUs. The stack is composed of several integrated components:

• MaxText model library: JAX LLMs optimized for TPUs.
• Tunix: Provides the latest algorithms and post-training techniques.
• vLLM on TPU: Used for high-performance sampling (inference) during Reinforcement Learning (RL).
• Pathways: A single controller JAX runtime used for multi-host inference (sampling) and efficient weight transfer.

Supported techniques include:

• SFT (Supervised Fine-Tuning)
• LoRA (Low-Rank Adaptation)
• DPO (Direct Preference Optimization) and ORPO (Odds-Ratio Policy Optimization)
• Multimodal SFT
• Reinforcement Learning (RL)

MaxText is a high-performance LLM framework written in Python/JAX designed for Google Cloud TPUs and NVIDIA GPUs. It focuses on scalability and high Model FLOPs Utilization (MFU).

Key Features:

• Supported Precisions: FP32, BF16, INT8, and FP8.
• Ahead-of-Time Compilation (AOT): Enables faster prototyping and earlier Out-of-Memory (OOM) detection.
• Quantization: Recommended via Qwix (supports QAT and PTQ); AQT is deprecated.
• Diagnostics: Logging via max_logging, profiling with XProf, and visualization in TensorBoard.
• Multi-Token Prediction (MTP): Improves training efficiency for specific architectures.
• Elastic Training: Supports fault-tolerant, dynamic scaling on Cloud TPUs using Pathways.
• Flexible Remat Policy: Allows memory-compute trade-offs using pre-defined policies ('full', 'minimal') or a 'custom' policy.

GRPO is a reinforcement learning algorithm implemented in MaxText designed to optimize language model policies using a reward signal. It enables training models to perform specific tasks beyond standard language modeling objectives.

Key architectural characteristics:

• Training: Uses Fully Sharded Data Parallelism (FSDP) to distribute model parameters across multiple devices.
• Inference: Uses a combination of Data Parallelism (DP) and Tensor Parallelism (TP) to distribute data and model shards across devices for efficient sampling.
• Orchestration: Designed to work with Pathways, a system for orchestrating computations across multiple TPU slices.

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

Key Capabilities

• Model Support: Includes Gemma, Llama, DeepSeek, Qwen, and Mistral.
• Training Scales: Supports pre-training (up to tens of thousands of chips) and scalable post-training.
• Post-Training Techniques: Supports Supervised Fine-Tuning (SFT), Group Relative Policy Optimization (GRPO), and Group Sequence Policy Optimization (GSPO).
• Performance: Achieves high Model FLOPs Utilization (MFU) and high tokens/second from single hosts to large clusters using JAX and the XLA compiler.

MaxText provides several pathways for maximizing performance and Model Flops Utilization (MFU) on TPUs. Optimization efforts can be categorized into four main areas:

1. Customizing Model Configs: Tailoring LLM model configurations specifically for higher performance on TPU hardware.
2. Sharding Strategies: Selecting efficient parallelization strategies such as FSDP (Fully Sharded Data Parallel), TP (Tensor Parallel), EP (Expert Parallel), and PP (Pipeline Parallel) based on Roofline Analysis and arithmetic intensity.
3. Pallas Kernels: Using Pallas kernels to achieve fine-grained control over hardware execution for specialized performance gains.
4. Benchmarking & Tuning: Setting up performance benchmarks, performing tuning cycles, and analyzing performance metrics to identify bottlenecks.

MaxText is built on a curated stack of JAX libraries designed for explicitness, composability, and performance at scale. The core components are:

• Flax (NNX): Used for functional and object-oriented model definition.
• Optax: Used for composable gradient processing and optimization.
• Orbax: Used for robust, asynchronous checkpointing and format conversion.
• Grain: Used for deterministic, multi-host data loading and global shuffling.
• Qwix: Used for native JAX quantization (QAT and PTQ).
• Tunix: Used for post-training tasks like SFT, RL, and PEFT (e.g., LoRA).

MaxText is built on a specific stack of technologies that handle the complexity of high-performance LLM training. Understanding these components is essential for moving from basic training to advanced research and production:

Multi-tier checkpointing is a solution for large-scale ML training (thousands of nodes) designed to increase Goodput and reduce mean-time-to-recovery (MTTR).

Instead of saving directly to slow persistent storage (which blocks training), it uses a tiered architecture:

1. RAM (in-memory): Fast, local saves to each node's ramdisk for low latency.
2. In-cluster (peer replication): Checkpoints are replicated to other nodes/slices within the cluster.
3. GCS (persistent storage): Background backup to GCS for long-term durability.

A managed GKE component handles replication and GCS backups asynchronously, meaning the training process is not blocked by the device-to-host data transfer to GCS.

Online distillation runs both the teacher and student models within the same training process. In each step, both models perform a forward pass on the same batch, and the student is updated to match a combination of:

1. KL-divergence soft loss: The teacher's softened logit distribution (controlled by distill_alpha).
2. Cross-entropy hard loss: The ground-truth labels.
3. Feature loss (Optional): The teacher's intermediate attention activations (controlled by distill_beta).

The trainer entry point is maxtext.trainers.post_train.distillation.train_distill, which is built on Tunix.