### Index # MaxText ```{raw} html --- file: index.html --- ```
```{include} ../README.md --- start-after: end-before: --- ```
```{toctree} --- maxdepth: 2 hidden: --- getting_started install_maxtext tutorials run_maxtext guides reference development release_notes ``` --- ### Development ```{include} ../CONTRIBUTING.md ``` ```{toctree} --- hidden: --- development/update_dependencies.md development/contribute_docs.md development/hlo_diff_testing.md ``` --- ### Getting Started (getting-started)= # Getting Started Welcome to MaxText! This guide will help you get started with running your first MaxText workloads. Whether you are working on a single host or scaling up to a multihost environment using Cloud TPUs or NVIDIA GPUs, this page provides the starting point for your journey. Follow the steps below to install MaxText, train your first model, and run inference. ## Prerequisites 1. To store logs and checkpoints, [create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) in your project. To run MaxText, the TPU or GPU VMs must have read/write permissions for the bucket. These permissions are granted by service account roles, such as the `STORAGE ADMIN` role. 2. MaxText reads a yaml file for configuration. We also recommend reviewing the configurable options in `configs/base.yml`. This file includes a decoder-only model of ~1B parameters. The configurable options can be overwritten from the command line. For instance, you can change the `steps` or `log_period` by either modifying `configs/base.yml` or by passing in `steps` and `log_period` as additional arguments to the `train.py` call. Set `base_output_directory` to a folder in the bucket you just created. 3. **Checkpoint Conversion**: In order to run MaxText on HuggingFace checkpoints, you must convert them to the MaxText/Orbax format first. For detailed instructions, see the [Checkpoint Conversion Guide](guides/checkpointing_solutions/convert_checkpoint.md). ## Running MaxText on a Single Host This procedure describes how to run MaxText on a single GPU or TPU host. ### 1. Installation Before running MaxText, you must install it on your VM. - For detailed installation instructions, see the [Installation Guide](install_maxtext.md). - For TPU VMs, install `maxtext[tpu]` for pre-training, or `maxtext[tpu-post-train]` for post-training. - For GPU VMs, ensure you install `maxtext[cuda12]`. ### 2. Running Pre-training To get started with training your first model, refer to the [Pre-training Tutorial](tutorials/pretraining.md). ### 3. Running Post-training To fine-tune your model or apply post-training techniques (such as SFT or RL), refer to the [Post-training Tutorial](tutorials/post_training_index.md). This guide covers various post-training workflows. ### 4. Running Inference To run inference (decoding) using MaxText models, refer to the [Inference Tutorial](tutorials/inference.md). This guide covers offline and online inference, as well as integration with vLLM. ## Running MaxText on Multiple Hosts Google Kubernetes Engine (GKE) is the recommended way to run MaxText on multiple hosts. It provides a managed environment for deploying and scaling containerized applications, including those that require TPUs or GPUs. See [Running MaxText with Cluster Toolkit](run_maxtext/run_maxtext_via_cluster_toolkit.md) or [Running Maxtext with XPK](run_maxtext/run_maxtext_via_xpk.md) for details. ## Running MaxText in Notebooks You can run MaxText interactively using Jupyter notebooks, Google Colab, or Visual Studio Code. Refer to the [Notebook Guide](guides/run_python_notebook.md) for instructions on setting up your notebook environment on TPUs. ## Next steps: preflight optimizations After you get workloads running, there are optimizations you can apply to improve performance. For more information, see [Optimization Tips](guides/optimization.md). --- ### Guides # How-to guides Explore our how-to guides for optimizing, debugging, and managing your MaxText workloads. ::::{grid} 1 2 2 2 :gutter: 2 :::{grid-item-card} ⚡ Optimization :link: guides/optimization :link-type: doc Techniques for maximizing performance, including sharding strategies, Pallas kernels, and benchmarking. ::: :::{grid-item-card} 💾 Data Pipelines :link: guides/data_input_pipeline :link-type: doc Configure input pipelines using **Grain** (recommended for determinism), **HuggingFace**, or **TFDS**. ::: :::{grid-item-card} 🔄 Checkpointing :link: guides/checkpointing_solutions :link-type: doc Manage GCS checkpoints, handle preemption with emergency checkpointing, and configure multi-tier storage. ::: :::{grid-item-card} 🔍 Monitoring & Debugging :link: guides/monitoring_and_debugging :link-type: doc Tools for observability: goodput monitoring, hung job debugging, and Vertex AI TensorBoard integration. ::: :::{grid-item-card} 🐍 Python Notebooks :link: guides/run_python_notebook :link-type: doc Interactive development guides for running MaxText on Google Colab or local JupyterLab environments. ::: :::{grid-item-card} 🌱 Model Bringup :link: guides/model_bringup :link-type: doc A step-by-step guide for the community to help expand MaxText's model library. ::: :::{grid-item-card} 🎗️ LoRA Model Bringup :link: guides/lora_model_bringup :link-type: doc Learn how to integrate Low-Rank Adaptation (LoRA) support for a new model architecture. ::: :::{grid-item-card} 🎓 Distillation :link: guides/distillation :link-type: doc How online distillation works in MaxText: loss anatomy, α / β / temperature schedule tuning, layer indices, monitoring metrics, and troubleshooting. ::: :::{grid-item-card} 📊 Evaluation :link: guides/eval_framework :link-type: doc Run benchmark evaluation (lm-eval, evalchemy, custom datasets) against MaxText checkpoints using the vLLM-native eval framework. ::: :::: ```{toctree} --- hidden: maxdepth: 1 --- guides/optimization.md guides/data_input_pipeline.md guides/checkpointing_solutions.md guides/monitoring_and_debugging.md guides/run_python_notebook.md guides/model_bringup.md guides/lora_model_bringup.md guides/distillation.md guides/eval_framework.md ``` --- ### Reference # Reference documentation Deep dive into MaxText architecture, models, and core concepts. ````{grid} 1 2 2 2 --- gutter: 2 --- ```{grid-item-card} 📊 Performance Metrics :link: reference/performance_metrics :link-type: doc Understanding Model Flops Utilization (MFU), calculation methods, and why it matters for performance optimization. ``` ```{grid-item-card} 🤖 Models :link: reference/models :link-type: doc Supported models and architectures, including Llama, Qwen, and Mixtral. Details on tiering and new additions. ``` ```{grid-item-card} 🏗️ Architecture :link: reference/architecture :link-type: doc High-level overview of MaxText design, JAX/XLA choices, and how components interact. ``` ```{grid-item-card} 💡 Core Concepts :link: reference/core_concepts :link-type: doc Key concepts including checkpointing strategies, quantization, tiling, and Mixture of Experts (MoE) configuration. ``` ```` ## 📚 API Reference Find comprehensive API documentation for MaxText modules, classes, and functions in the [API Reference page](reference/api.rst). ```{toctree} --- hidden: maxdepth: 1 --- reference/performance_metrics reference/models reference/architecture reference/core_concepts reference/api.rst ``` --- ### Release Notes # MaxText release notes ## PyPI Package MaxText is [available in PyPI](https://pypi.org/project/maxtext/) and can be installed through pip. Please see our [MaxText Installation Guide](install_maxtext.md) for setup instructions. ## Releases ### v0.2.3 #### Changes - Upgraded JAX to version 0.10.0 for pre-training and 0.10.1 for post-training. - **New vLLM-Powered Evaluation Framework**: Introduced an eval framework for running lm-eval, evalchemy, and custom benchmarking against MaxText checkpoints. See the [evaluation guide](https://maxtext.readthedocs.io/en/latest/guides/eval_framework.html) for details. - Added support for pre-training new models: - **Qwen3.5**: Qwen3.5 35B & 397B is now [supported](https://github.com/AI-Hypercomputer/maxtext/blob/d938b91acaa3baaaf32956e21677bd29e14549a1/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md). - **Qwen3-Omni**: Support for multimodal SFT ([PR #3863](https://github.com/AI-Hypercomputer/maxtext/pull/3863)). - **Direct Preference Optimization (DPO/ORPO) Support**: Full support for DPO and ORPO alignment pipelines. See the [DPO tutorial](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/dpo.html) for details. - **Reinforcement Learning (RL) Recipe**: Added a pre-configured [RL recipe for Qwen3-30b-a3b](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl_qwen3_30b.html). - **Iterative Quality Monitoring (RL)**: Added intermediate evaluation hooks to automatically run quality benchmarks during RL training (every `eval_interval` steps), optimized with a new `eval_batch_size` configuration knob. - **Developer Extensibility**: Added `dataset_processor_path` CLI knob for custom dataset integration, and refactored shared post-training hooks to simplify custom SFT, DPO, and RL workflow development. - **Generalized Learn-to-Init (LTI) for Distillation**: Enhanced post-training distillation capabilities with generalized LTI support. - Added support for recording elastic goodput events during training to track efficiency ([PR #3901](https://github.com/AI-Hypercomputer/maxtext/pull/3901)). - **Installation Updates**: Updated the `[tpu-post-train]` installation command to require `UV_TORCH_BACKEND=cpu`(see [Installation Guide](install_maxtext.md)). - **Zero1 AOT Compilation**: Added zero1 support to Ahead-Of-Time (AOT) compilation in train compile, improving compilation capabilities for zero1 config. - **MoE Performance Optimization**: Integrated ragged gather reduce into Mixture of Experts (MoE) layers to optimize memory and performance by replacing ragged scatter and supporting backward pass. - Added [E2E scripts](https://github.com/AI-Hypercomputer/maxtext/tree/main/tests/end_to_end/tpu/gemma3/4b) to run checkpoint conversion, pre-training and post-training (SFT, RL) with Gemma3-4B model. - **Bug Fixes and Usability Enhancements**: - **Attention Masking Fix in RL**: Fixed an issue in `TunixMaxTextAdapter` where queries at non-pad positions could attend to pad-position keys during training, which was corrupting log-probabilities and affecting GRPO training reward trajectories ([PR #4016](https://github.com/AI-Hypercomputer/maxtext/pull/4016)). - **JAX/NNX Gradient Mutation Fix**: Refactored post-training loops (`train_distill`, `train_sft`, `train_rl`) to use `jax.value_and_grad` with explicit NNX state split/merge instead of nesting `nnx.value_and_grad` inside `nnx.jit` ([PR #3652](https://github.com/AI-Hypercomputer/maxtext/pull/3652)). - **Qwen3-MoE Checkpoint Conversion**: Fixed checkpoint conversion issues for Qwen3-MoE models ([PR #3868](https://github.com/AI-Hypercomputer/maxtext/pull/3868)). - **Duplicate Configuration Failures Fix**: Allowed identical config overrides and handled configuration exceptions cleanly ([PR #3933](https://github.com/AI-Hypercomputer/maxtext/pull/3933)). - **Documentation Improvements**: Updated [Getting started](https://maxtext.readthedocs.io/en/latest/getting_started.html) guide, including new guides for the [evaluation framework](https://maxtext.readthedocs.io/en/latest/guides/eval_framework.html) and the [DPO tutorial](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/dpo.html). #### Deprecations - Deleted [legacy DPO implementation](https://github.com/AI-Hypercomputer/maxtext/pull/3997) in favor of the integrated [DPO trainer](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/dpo.html). - Removed stack trace collection feature. ### v0.2.2 #### Changes - Upgraded JAX to version 0.9.2, improving support for both pre-training and post-training. - Introduced simplified APIs for accessing MaxText models. - Included [maxtext_with_gepa.ipynb](https://github.com/AI-Hypercomputer/maxtext/blob/3c7d8d27864fc12cccac07786f02bd0e5262c982/src/maxtext/examples/maxtext_with_gepa.ipynb), a new notebook demonstrating AIME prompt optimization using the GEPA framework within MaxText. - Added support for Kimi-K2 models and the MuonClip optimizer. Users can explore this with the [kimi-k2-1t](https://github.com/AI-Hypercomputer/maxtext/blob/fa5b5ebf9a8e4f7a33bd88eae051dc21f3147791/src/maxtext/configs/models/kimi-k2-1t.yml) config (see [user guide](https://github.com/AI-Hypercomputer/maxtext/blob/fa5b5ebf9a8e4f7a33bd88eae051dc21f3147791/tests/end_to_end/tpu/kimi/Run_Kimi.md) for details). - Kimi-K2-Thinking, Kimi-K2.5 (text), and Kimi-K2.6 (text) are now supported. See [Run_Kimi.md](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/kimi/Run_Kimi.md#quantized-variants-k2-thinking-k25-k26) for details. - [DeepSeek-V3.2](https://arxiv.org/pdf/2512.02556) is now supported, including DeepSeek Sparse Attention for handling long contexts. Use the [deepseek3.2-671b](https://github.com/AI-Hypercomputer/maxtext/blob/20d93f62a91899dbbb8f23562973d75104411d3a/src/maxtext/configs/models/deepseek3.2-671b.yml) config to try it out (refer to the [user guide](https://github.com/AI-Hypercomputer/maxtext/blob/20d93f62a91899dbbb8f23562973d75104411d3a/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md) for more information). - Support has been added for Gemma 4 multi-modal models (26B MoE and 31B dense). These can be used with the [gemma4-26b](https://github.com/AI-Hypercomputer/maxtext/blob/cdc587f0935a5e2d6f8287b96669cf2e87a0acdc/src/maxtext/configs/models/gemma4-26b.yml) and [gemma4-31b](https://github.com/AI-Hypercomputer/maxtext/blob/cdc587f0935a5e2d6f8287b96669cf2e87a0acdc/src/maxtext/configs/models/gemma4-31b.yml) configs. See [Run_Gemma4.md](https://github.com/AI-Hypercomputer/maxtext/blob/cdc587f0935a5e2d6f8287b96669cf2e87a0acdc/tests/end_to_end/tpu/gemma4/Run_Gemma4.md) for further details. - Support has been added for Gemma 4 inference using [MaxText on vLLM plugin](tutorials/inference.md). - Enhanced RL capabilities with support for the `open-r1/OpenR1-Math-220k` dataset and `nvidia/OpenMathReasoning`. - Added more evaluation modes for RL like majority voting and pass@1 estimation. - Sync weights to vllm prior to pre RL evaluation. - More robust usage of math-verify in RL. - MaxText's Supervised Fine-Tuning (SFT) now supports non-instruct models. - Added support for tensor parallelism using the Fused MoE kernel for MaxText on vLLM inference. - Added support for MaxText to vllm converters for Qwen3 and Gemma4 family of models. - [validate_converter.py](https://github.com/AI-Hypercomputer/maxtext/blob/472f53b70089e661be399ad3905c05a53a172ec5/src/maxtext/integration/vllm/torchax_converter/validate_converter.py#L108) now runs on multislice environment to test larger models with utilities to compare maxtext and vllm weights. #### Deprecations - Legacy `MaxText.*` shims have been removed. Please refer to [src/MaxText/README.md](https://github.com/AI-Hypercomputer/maxtext/blob/0536605a8ca116087ed93178433a67e905be566c/src/MaxText/README.md) for details on the new command locations and how to migrate. - Sequence parallelism has been deprecated, please use context parallelism instead. - The flag `expert_shard_attention_option` is deprecated, use `custom_mesh_and_rule=ep-as-cp` for the same functionality. ### v0.2.1 #### Changes - Use the new `maxtext[runner]` installation option to build Docker images without cloning the repository. This can be used for scheduling jobs through XPK. See the [MaxText installation instructions](build-docker) for more info. - Config can now be inferred for most MaxText commands. If you choose not to provide a config, MaxText will now [select an appropriate one](https://github.com/AI-Hypercomputer/maxtext/blob/9e786c888cc7acdfc00a8f73064e285017e80b86/src/maxtext/configs/pyconfig.py#L51-L67). - Configs in MaxText PyPI will now be picked up without storing them locally. - New features from DeepSeek-AI are now supported: Conditional Memory via Scalable Lookup ([Engram](https://arxiv.org/abs/2601.07372)) and Manifold-Constrained Hyper-Connections ([mHC](https://arxiv.org/abs/2512.24880)). Try them out with our [deepseek-custom](https://github.com/AI-Hypercomputer/maxtext/blob/9e786c888cc7acdfc00a8f73064e285017e80b86/src/maxtext/configs/models/deepseek-custom.yml) starter config. - MaxText now supports customizing your own mesh and logical rules. Two examples guiding how to use your own mesh and rules for sharding are provided in the [custom_mesh_and_rule](https://github.com/AI-Hypercomputer/maxtext/tree/9e786c888cc7acdfc00a8f73064e285017e80b86/src/maxtext/configs/custom_mesh_and_rule) directory. ### v0.2.0 #### Changes - New `tpu-post-train` target in PyPI. Please also use this installation option for running vllm_decode. See the [MaxText installation instructions](install_maxtext.md) for more info. - [Qwen3-Next](https://github.com/AI-Hypercomputer/maxtext/blob/7656eb8d1c9eb0dd91e617a6fdf6ad805221221a/tests/end_to_end/tpu/qwen/next/run_qwen3_next.md) is now supported. - New MaxText structure! MaxText has been restructured according to [RESTRUCTURE.md](https://github.com/AI-Hypercomputer/maxtext/blob/1b9e38aa0a19b6018feb3aed757406126b6953a1/RESTRUCTURE.md). Please feel free to share your thoughts and feedback. - [Muon optimizer](https://kellerjordan.github.io/posts/muon) is now supported. - DeepSeek V3.1 is now supported. Use existing configs for [DeepSeek V3 671B](https://github.com/AI-Hypercomputer/maxtext/blob/7656eb8d1c9eb0dd91e617a6fdf6ad805221221a/src/maxtext/configs/models/deepseek3-671b.yml) and load in V3.1 checkpoint to use model. - [New RL and SFT Notebook tutorials](https://github.com/AI-Hypercomputer/maxtext/tree/7656eb8d1c9eb0dd91e617a6fdf6ad805221221a/src/maxtext/examples) are available. - The [ReadTheDocs documentation site](index.md) has been reorganized. - Multi-host support for GSPO and GRPO is now available via [new RL tutorials](tutorials/posttraining/rl_on_multi_host.md). - A new guide, [What is Post Training in MaxText?](tutorials/post_training_index.md), is now available. - Ironwood TPU co-designed AI stack announced. Read the [blog post on its co-design with MaxText](https://cloud.google.com/blog/products/compute/inside-the-ironwood-tpu-codesigned-ai-stack?e=48754805). - [Optimized models tiering documentation](reference/models/tiering.md) has been refreshed. - Added Versioning. Check out our [first set of release notes](release_notes.md)! - Post-Training (SFT, RL) via [Tunix](https://github.com/google/tunix) is now available. - Vocabulary tiling ([PR](https://github.com/AI-Hypercomputer/maxtext/pull/2242)) is now supported in MaxText! Adjust config `num_vocab_tiling` to unlock more efficient memory usage. - The GPT-OSS family of models (20B, 120B) is now supported. #### Deprecations - Many MaxText modules have changed locations. Core commands like train, decode, sft, etc. will still work as expected temporarily. Please update your commands to the latest file locations - install_maxtext_github_deps installation script replaced with install_maxtext_tpu_github_deps - `tools/setup/setup_post_training_requirements.sh` for post training dependency installation is deprecated in favor of [pip installation](install_maxtext.md) ### v0.1.0 Our first MaxText PyPI package is here! MaxText is a high performance, highly scalable, open-source LLM library and reference implementation written in pure Python/JAX and targeting Google Cloud TPUs and GPUs for training. We are excited to make it easier than ever to get started. Users can now install MaxText through pip, both for local development and through stable PyPI builds. Please see our [MaxText Installation Guide](install_maxtext.md) for more setup details. Going forward, this page will document notable changes as we release new versions of MaxText. --- ### Run Maxtext # Run MaxText Choose your environment and orchestration method to run MaxText. ::::{grid} 1 2 2 2 :gutter: 2 :::{grid-item-card} 💻 Localhost / Single VM :link: run_maxtext/run_maxtext_localhost :link-type: doc Get started quickly on a single machine. Clone the repo, install dependencies, and run your first training job on a single TPU or GPU VM. ::: :::{grid-item-card} 🎮 Single-host GPU :link: run_maxtext/run_maxtext_single_host_gpu :link-type: doc Run MaxText on single-host NVIDIA GPUs (e.g., A3 High/Mega). Includes Docker setup, NVIDIA Container Toolkit installation, and 1B/7B model training examples. ::: :::{grid-item-card} 🚀 At scale with Cluster Toolkit (GKE) :link: run_maxtext/run_maxtext_via_cluster_toolkit :link-type: doc Deploy to Google Kubernetes Engine (GKE) using Cluster Toolkit's `gcluster` CLI. Package and run multi-host JAX workloads with on-the-fly container builds. ::: :::{grid-item-card} 🏗️ At scale with XPK (GKE) :link: run_maxtext/run_maxtext_via_xpk :link-type: doc Deploy to Google Kubernetes Engine (GKE) using XPK. Orchestrate large-scale training jobs on TPU or GPU clusters with simple CLI commands. ::: :::{grid-item-card} 🌐 Multi-host via Pathways :link: run_maxtext/run_maxtext_via_pathways :link-type: doc Run large-scale JAX jobs on TPUs using Pathways. Supports batch and headless (interactive) workloads on GKE. ::: :::{grid-item-card} 🔌 Decoupled Mode :link: run_maxtext/decoupled_mode :link-type: doc Run tests and local development without Google Cloud dependencies (no `gcloud`, GCS, or Vertex AI required). ::: :::{grid-item-card} ♻️ Elastic training (demo) :link: run_maxtext/run_maxtext_elastic_training :link-type: doc Demonstrate fault-tolerant training with Pathways on GKE: lose a TPU slice mid-run and recover in-process from the last checkpoint, no job restart. ::: :::: ```{toctree} --- hidden: maxdepth: 1 --- run_maxtext/run_maxtext_localhost.md run_maxtext/run_maxtext_single_host_gpu.md run_maxtext/run_maxtext_via_cluster_toolkit.md run_maxtext/run_maxtext_via_xpk.md run_maxtext/run_maxtext_via_pathways.md run_maxtext/decoupled_mode.md run_maxtext/run_maxtext_elastic_training.md ``` ---