From c55c36b93898abddf47f2080356f95256f8cc74e Mon Sep 17 00:00:00 2001 From: Branden Vandermoon Date: Fri, 6 Mar 2026 18:22:34 +0000 Subject: [PATCH] MaxText 0.2.0 release notes and installation instructions --- README.md | 7 ++++--- docs/install_maxtext.md | 46 +++++++++++++++++++++++++++++++++-------- docs/release_notes.md | 24 +++++++++++++++++++++ 3 files changed, 65 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 65768e8df2..679a39681a 100644 --- a/README.md +++ b/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. diff --git a/docs/install_maxtext.md b/docs/install_maxtext.md index 0a10bdbdb8..3b6a18fcdf 100644 --- a/docs/install_maxtext.md +++ b/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: diff --git a/docs/release_notes.md b/docs/release_notes.md index a656658ef7..674426f2f7 100644 --- a/docs/release_notes.md +++ b/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.