Merge pull request #3338 from AI-Hypercomputer:bvandermoon-install-maxtext

PiperOrigin-RevId: 879766150

Author: Google-ML-Automation

README.md

@@ -41,11 +41,12 @@ See our guide on running MaxText in decoupled mode, without any GCP dependencies
 
 ## 🔥 Latest news 🔥
 
-* \[March 5, 2026\] [Qwen3-Next](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen/next/run_qwen3_next.md) is now supported. 
+* \[March 5, 2026\] New `tpu-post-train` [target in PyPI](https://pypi.org/project/maxtext). Please also use this installation option for running vllm_decode. See the [MaxText installation instructions](https://maxtext.readthedocs.io/en/latest/install_maxtext.html) for more info.
+* \[March 5, 2026\] [Qwen3-Next](https://github.com/AI-Hypercomputer/maxtext/blob/7656eb8d1c9eb0dd91e617a6fdf6ad805221221a/tests/end_to_end/tpu/qwen/next/run_qwen3_next.md) is now supported.
 * \[February 27, 2026\] 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. 
 * \[December 22, 2025\] [Muon optimizer](https://kellerjordan.github.io/posts/muon) is now supported.
-* \[December 10, 2025\] DeepSeek V3.1 is now supported. Use existing configs for [DeepSeek V3 671B](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/configs/models/deepseek3-671b.yml) and load in V3.1 checkpoint to use model.
-* \[December 9, 2025\] [New RL and SFT Notebook tutorials](https://github.com/AI-Hypercomputer/maxtext/tree/main/src/maxtext/examples) are available.
+* \[December 10, 2025\] 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.
+* \[December 9, 2025\] [New RL and SFT Notebook tutorials](https://github.com/AI-Hypercomputer/maxtext/tree/7656eb8d1c9eb0dd91e617a6fdf6ad805221221a/src/maxtext/examples) are available.
 * \[December 4, 2025\] The [ReadTheDocs documentation site](https://maxtext.readthedocs.io/en/latest/index.html) has been reorganized.
 * \[December 3, 2025\] Multi-host support for GSPO and GRPO is now available via [new RL tutorials](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl_on_multi_host.html).
 * \[November 20, 2025\] A new guide, [What is Post Training in MaxText?](https://maxtext.readthedocs.io/en/latest/tutorials/post_training_index.html), is now available.

docs/install_maxtext.md

@@ -17,6 +17,11 @@
 # Install MaxText
 
 This document discusses how to install MaxText. We recommend installing MaxText inside a Python virtual environment.
+MaxText offers three installation modes:
+
+1. maxtext[tpu]. Used for pre-training and decode on TPUs.
+2. maxtext[cuda12]. Used for pre-training and decode on GPUs.
+3. maxtext[tpu-post-train]. Used for post-training on TPUs. Currently, this option should also be used for running vllm_decode on TPUs.
 
 ## From PyPI (Recommended)
 This is the easiest way to get started with the latest stable version.
@@ -29,11 +34,24 @@ pip install uv
 uv venv --python 3.12 --seed maxtext_venv
 source maxtext_venv/bin/activate
 
-# 3. Install MaxText and its dependencies
-uv pip install maxtext --resolution=lowest
-install_maxtext_github_deps
+# 3. Install MaxText and its dependencies. Choose a single
+#      installation option from this list to fit your use case.
+
+# Option 1: Installing maxtext[tpu]
+uv pip install maxtext[tpu] --resolution=lowest
+install_maxtext_tpu_github_deps
+
+# Option 2: Installing maxtext[cuda12]
+uv pip install maxtext[cuda12] --resolution=lowest
+install_maxtext_cuda12_github_dep
+
+# Option 3: Installing maxtext[tpu-post-train]
+uv pip install maxtext[tpu-post-train] --resolution=lowest
+install_maxtext_tpu_post_train_extra_deps
 ```
-> **Note:** The `install_maxtext_github_deps` command is temporarily required to install dependencies directly from GitHub that are not yet available on PyPI.
+> **Note:** The `install_maxtext_tpu_github_deps`, `install_maxtext_cuda12_github_dep`, and
+`install_maxtext_tpu_post_train_extra_deps` commands are temporarily required to install dependencies directly from GitHub
+that are not yet available on PyPI. As shown above, choose the one that corresponds to your use case.
 
 > **Note:** The maxtext package contains a comprehensive list of all direct and transitive dependencies, with lower bounds, generated by [seed-env](https://github.com/google-ml-infra/actions/tree/main/python_seed_env). We highly recommend the `--resolution=lowest` flag. It instructs `uv` to install the specific, tested versions of dependencies defined by MaxText, rather than the latest available ones. This ensures a consistent and reproducible environment, which is critical for stable performance and for running benchmarks.
 
@@ -50,12 +68,20 @@ pip install uv
 uv venv --python 3.12 --seed maxtext_venv
 source maxtext_venv/bin/activate
 
-# 3. Install dependencies in editable mode
-# install the tpu package
+# 3. Install dependencies in editable mode. Choose a single
+#      installation option from this list to fit your use case.
+
+# Option 1: Installing .[tpu]
 uv pip install -e .[tpu] --resolution=lowest
-# or install the gpu package by running the following line
-# uv pip install -e .[cuda12] --resolution=lowest
-install_maxtext_github_deps
+install_maxtext_tpu_github_deps
+
+# Option 2: Installing .[cuda12]
+uv pip install -e .[cuda12] --resolution=lowest
+install_maxtext_cuda12_github_dep
+
+# Option 3: Installing .[tpu-post-train]
+uv pip install -e .[tpu-post-train] --resolution=lowest
+install_maxtext_tpu_post_train_extra_deps
 ```
 
 After installation, you can verify the package is available with `python3 -c "import maxtext"` and run training jobs with `python3 -m maxtext.trainers.pre_train.train ...`.
@@ -66,6 +92,8 @@ After installation, you can verify the package is available with `python3 -c "im
 
 This document provides a guide to updating dependencies in MaxText using the `seed-env` tool. `seed-env` helps generate deterministic and reproducible Python environments by creating fully-pinned `requirements.txt` files from a base set of requirements.
 
+Please keep dependencies updated throughout development. This will allow each commit to work properly from both a feature and dependency perspective. We will periodically upload commits to PyPI for stable releases. But it is also critical to keep dependencies in sync for users installing MaxText from source.
+
 ## Overview of the Process
 
 To update dependencies, you will follow these general steps:

docs/release_notes.md

@@ -22,6 +22,30 @@ MaxText is [available in PyPI](https://pypi.org/project/maxtext/) and can be ins
 
 ## Releases
 
+### 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](https://maxtext.readthedocs.io/en/latest/install_maxtext.html) 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](https://maxtext.readthedocs.io/en/latest/index.html) has been reorganized.
+* Multi-host support for GSPO and GRPO is now available via [new RL tutorials](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl_on_multi_host.html).
+* A new guide, [What is Post Training in MaxText?](https://maxtext.readthedocs.io/en/latest/tutorials/post_training_index.html), 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](https://maxtext.readthedocs.io/en/latest/reference/models/tiering.html) has been refreshed.
+* Added Versioning. Check out our [first set of release notes](https://maxtext.readthedocs.io/en/latest/release_notes.html)!
+* 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](https://maxtext.readthedocs.io/en/latest/install_maxtext.html)
+
 ### 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.