Merge branch 'main' into sanbao/format

Author: susanbao

.github/CODEOWNERS

@@ -1,4 +1,4 @@
-* @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google
+* @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google @dipannita08 @igorts-git
 
 # Model bring-up
 src/MaxText/assets @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande

.github/workflows/pypi_release.yml

@@ -44,7 +44,9 @@ jobs:
 
   publish_maxtext_package_to_pypi:
     name: Publish MaxText package to PyPI
-    needs: [build_and_test_maxtext_package]
+    # Temporarily only require release_approval for a one-time upload.
+    # Immediately revert this to `needs: [build_and_test_maxtext_package]`.
+    needs: [release_approval]
     runs-on: ubuntu-latest
     environment: release
     steps:

.github/workflows/run_jupyter_notebooks.yml

@@ -57,33 +57,18 @@ jobs:
       - name: Install MaxText and Dependencies
         shell: bash
         run: |
+          # 1. Create virtual environment
           python3 -m uv venv --seed
           source .venv/bin/activate
-
-          # Install MaxText package
           maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
-          uv pip install ${maxtext_wheel}[${MAXTEXT_PACKAGE_EXTRA}] --resolution=lowest
-          uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt
 
-          # Install dependencies for running notebooks
-          uv pip install papermill ipykernel ipywidgets
+          # 2. Install MaxText package and all the post training dependencies
+          uv pip install ${maxtext_wheel}[tpu-post-train] --resolution=lowest
+          #TODO: @mazumdera:  replace this with the following after release
+          # uv pip install maxtext[tpu-post-train] --resolution=lowest
+          install_maxtext_tpu_post_train_extra_deps
           .venv/bin/python3 -m ipykernel install --user --name maxtext_venv
-
-          # Install Tunix for post-training notebooks
-          git clone https://github.com/google/tunix
-          uv pip install ./tunix
           
-          # Install vllm for post-training notebooks
-          git clone https://github.com/vllm-project/vllm.git
-          VLLM_TARGET_DEVICE="tpu" uv pip install ./vllm
-
-          # Install tpu-inference for post-training notebooks
-          git clone https://github.com/vllm-project/tpu-inference.git
-          uv pip install ./tpu-inference
-
-          uv pip install --no-deps qwix==0.1.4
-          uv pip install --no-deps protobuf==5.29.5
-          uv pip install math-verify==0.9.0
           python3 -m pip freeze
       - name: Run Post-Training Notebooks
         shell: bash

.gitignore

@@ -148,3 +148,6 @@ dmypy.json
 # Gemini CLI
 .gemini/
 gha-creds-*.json
+
+# vscode workspace
+maxtext.code-workspace

README.md

@@ -41,7 +41,8 @@ See our guide on running MaxText in decoupled mode, without any GCP dependencies
 
 ## 🔥 Latest news 🔥
 
-* \[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.
+* \[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. 
+* \[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.

benchmarks/globals.py

@@ -17,7 +17,7 @@
 import os.path
 
 # This is the MaxText root: with "max_utils.py"; &etc. TODO: Replace `os.path.basename` with `os.path.abspath`
-MAXTEXT_PKG_DIR = os.environ.get("MAXTEXT_PKG_DIR", "src/MaxText")
+MAXTEXT_PKG_DIR = os.environ.get("MAXTEXT_PKG_DIR", "src/maxtext")
 
 # This is the maxtext repo root: with ".git" folder; "README.md"; "pyproject.toml"; &etc.
 MAXTEXT_REPO_ROOT = os.environ.get(

benchmarks/maxtext_xpk_runner.py

@@ -35,7 +35,7 @@
 import omegaconf
 
 import benchmarks.maxtext_trillium_model_configs as model_configs
-from benchmarks.globals import MAXTEXT_CONFIGS_DIR
+from benchmarks.globals import MAXTEXT_PKG_DIR
 from benchmarks.command_utils import run_command_with_updates
 import benchmarks.xla_flags_library as xla_flags
 from benchmarks.disruption_management.disruption_handler import DisruptionConfig
@@ -107,7 +107,7 @@ class WorkloadConfig:
   generate_metrics_and_upload_to_big_query: bool = True
   hardware_id: str = "v6e"
   metrics_gcs_file: str = ""
-  base_config: str = os.path.join(MAXTEXT_CONFIGS_DIR, "base.yml")
+  base_config: str = os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")
   topology: str = dataclasses.field(init=False)
   num_devices_per_slice: int = dataclasses.field(init=False)
   db_project: str = ""
@@ -354,7 +354,7 @@ def _build_args_from_config(wl_config: WorkloadConfig) -> dict:
       "xla_flags": f"'{xla_flags_str}'",
       "dataset": dataset,
       "run_type": "maxtext-xpk",
-      "config_file": os.path.join(MAXTEXT_CONFIGS_DIR, "base.yml"),
+      "config_file": os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"),
       "topology": wl_config.topology,
       "tuning_params": f"'{tuning_params_str}'",
       "db_project": wl_config.db_project,
@@ -439,8 +439,8 @@ def build_user_command(
           "export ENABLE_PATHWAYS_PERSISTENCE=1 &&",
           f"export JAX_PLATFORMS={jax_platforms} &&",
           "export ENABLE_PJRT_COMPATIBILITY=true &&",
-          "export MAXTEXT_ASSETS_ROOT=/deps/src/maxtext/assets MAXTEXT_PKG_DIR=/deps/src/MaxText MAXTEXT_REPO_ROOT=/deps &&"
-          f'{hlo_dump} python3 -m maxtext.trainers.pre_train.train {os.path.join(MAXTEXT_CONFIGS_DIR, "base.yml")}',
+          "export MAXTEXT_ASSETS_ROOT=/deps/src/maxtext/assets MAXTEXT_PKG_DIR=/deps/src/maxtext MAXTEXT_REPO_ROOT=/deps &&"
+          f'{hlo_dump} python3 -m maxtext.trainers.pre_train.train {os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")}',
           f"{config_tuning_params}",
           f"steps={wl_config.num_steps}",
           f"model_name={wl_config.model.model_type}",

dependencies/dockerfiles/maxtext_post_training_dependencies.Dockerfile

@@ -23,6 +23,9 @@ ENV PATH="/usr/local/google-cloud-sdk/bin:/usr/local/bin/python3.12:${PATH}"
 ARG MODE
 ENV ENV_MODE=$MODE
 
+ARG WORKFLOW
+ENV ENV_WORKFLOW=$WORKFLOW
+
 ARG JAX_VERSION
 ENV ENV_JAX_VERSION=$JAX_VERSION
 
@@ -34,7 +37,7 @@ ENV ENV_DEVICE=$DEVICE
 
 ENV MAXTEXT_ASSETS_ROOT=/deps/src/maxtext/assets
 ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/tests/assets
-ENV MAXTEXT_PKG_DIR=/deps/src/MaxText
+ENV MAXTEXT_PKG_DIR=/deps/src/maxtext
 ENV MAXTEXT_REPO_ROOT=/deps
 
 # Set the working directory in the container
@@ -43,14 +46,15 @@ WORKDIR /deps
 # Copy setup files and dependency files separately for better caching
 COPY tools/setup tools/setup/
 COPY dependencies/requirements/ dependencies/requirements/
-COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt src/install_maxtext_extra_deps/
+COPY src/install_maxtext_extra_deps/ src/install_maxtext_extra_deps/
+COPY src/maxtext/integration/vllm/ src/maxtext/integration/vllm/
 
 # Copy the custom libtpu.so file if it exists inside maxtext repository
 COPY libtpu.so* /root/custom_libtpu/
 
 # Install dependencies - these steps are cached unless the copied files change
-RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION LIBTPU_VERSION=$ENV_LIBTPU_VERSION DEVICE=${ENV_DEVICE}"
-RUN --mount=type=cache,target=/root/.cache/pip bash /deps/tools/setup/setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} LIBTPU_VERSION=${ENV_LIBTPU_VERSION} DEVICE=${ENV_DEVICE}
+RUN echo "Running command: bash setup.sh MODE=$ENV_MODE WORKFLOW=$ENV_WORKFLOW JAX_VERSION=$ENV_JAX_VERSION LIBTPU_VERSION=$ENV_LIBTPU_VERSION DEVICE=${ENV_DEVICE}"
+RUN --mount=type=cache,target=/root/.cache/pip bash /deps/tools/setup/setup.sh MODE=${ENV_MODE} WORKFLOW=${ENV_WORKFLOW} JAX_VERSION=${ENV_JAX_VERSION} LIBTPU_VERSION=${ENV_LIBTPU_VERSION} DEVICE=${ENV_DEVICE}
 
 # Now copy the remaining code (source files that may change frequently)
 COPY . .

dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile

@@ -11,6 +11,7 @@ antlr4-python3-runtime>=4.9.3
 anyio>=4.11.0
 aqtp>=0.9.0
 array_record>=0.8.3
+asttokens>=3.0.1
 astor>=0.8.1
 astroid>=4.0.2
 astunparse>=1.6.3
@@ -24,7 +25,7 @@ botocore>=1.42.56
 build>=1.3.0
 cachetools>=6.2.2
 cbor2>=5.8.0
-certifi>=2025.11.12
+certifi>=2026.2.25
 cffi>=2.0.0
 cfgv>=3.5.0
 charset-normalizer>=3.4.4
@@ -38,6 +39,7 @@ clu>=0.0.12
 cmake>=4.2.1
 colorama>=0.4.6
 colorful>=0.5.8
+comm>=0.2.3
 compressed-tensors>=0.13.0
 contourpy>=1.3.3
 coverage>=7.12.0
@@ -46,6 +48,7 @@ cycler>=0.12.1
 dacite>=1.9.2
 dataclasses-json>=0.6.7
 datasets>=4.6.0
+debugpy>=1.8.20
 decorator>=5.2.1
 depyf>=0.20.0
 dill>=0.4.0
@@ -60,13 +63,16 @@ editdistance>=0.8.1
 einops>=0.8.1
 einshape>=1.0
 email-validator>=2.3.0
+entrypoints>=0.4
 etils>=1.13.0
 evaluate>=0.4.6
 execnet>=2.1.2
+executing>=2.2.1
 fastapi>=0.122.0
 fastapi-cli>=0.0.24
 fastapi-cloud-cli>=0.13.0
 fastar>=0.8.0
+fastjsonschema>=2.21.2
 filelock>=3.20.0
 flatbuffers>=25.9.23
 flax>=0.12.4
@@ -126,13 +132,18 @@ importlib_metadata>=8.7.0
 importlib_resources>=6.5.2
 iniconfig>=2.3.0
 interegular>=0.3.3
+ipykernel>=7.2.0
+ipython>=9.10.0
+ipython_pygments_lexers>=1.1.1
+ipywidgets>=8.1.8
 isort>=7.0.0
 jaraco.classes>=3.4.0
 jaraco.context>=6.1.0
 jaraco.functools>=4.3.0
 jax>=0.8.3
 jaxlib>=0.8.3
 jaxtyping>=0.3.3
+jedi>=0.19.2
 jeepney>=0.9.0
 Jinja2>=3.1.6
 jiter>=0.13.0
@@ -141,6 +152,9 @@ joblib>=1.5.2
 jsonlines>=4.0.0
 jsonschema>=4.26.0
 jsonschema-specifications>=2025.9.1
+jupyter_client>=8.8.0
+jupyter_core>=5.9.1
+jupyterlab_widgets>=3.0.16
 kagglehub>=0.3.13
 keras>=3.12.0
 keyring>=25.7.0
@@ -162,6 +176,7 @@ MarkupSafe>=3.0.3
 marshmallow>=3.26.2
 math-verify>=0.9.0
 matplotlib>=3.10.7
+matplotlib-inline>=0.2.1
 mccabe>=0.7.0
 mcp>=1.26.0
 mdurl>=0.1.2
@@ -178,6 +193,8 @@ multidict>=6.7.0
 multiprocess>=0.70.18
 mypy_extensions>=1.1.0
 namex>=0.1.0
+nbclient>=0.10.4
+nbformat>=5.10.4
 nest-asyncio>=1.6.0
 networkx>=3.6
 ninja>=1.13.0
@@ -186,7 +203,7 @@ nltk>=3.9.2
 nodeenv>=1.9.1
 numba>=0.62.1
 numpy>=2.2.6
-numpy-typing-compat>=20250818.2.0
+numpy-typing-compat>=20250818.2.2
 nvidia-cublas-cu12>=12.8.4.1
 nvidia-cuda-cupti-cu12>=12.8.90
 nvidia-cuda-nvrtc-cu12>=12.8.93
@@ -218,7 +235,7 @@ opentelemetry-exporter-prometheus>=0.60b1
 opentelemetry-proto>=1.39.1
 opentelemetry-sdk>=1.39.1
 opentelemetry-semantic-conventions>=0.60b1
-opentelemetry-semantic-conventions-ai>=0.4.14
+opentelemetry-semantic-conventions-ai>=0.4.15
 opt_einsum>=3.4.0
 optax>=0.2.6
 optree>=0.18.0
@@ -228,23 +245,29 @@ orbax-export>=0.0.8
 outlines_core>=0.2.11
 packaging>=26.0
 pandas>=2.3.3
+papermill>=2.7.0
 parameterized>=0.9.0
+parso>=0.8.6
 partial-json-parser>=0.2.1.1.post7
 pathspec>=0.12.1
 pathwaysutils>=0.1.4
 perfetto>=0.16.0
+pexpect>=4.9.0
 pillow>=12.0.0
-platformdirs>=4.5.0
+platformdirs>=4.9.2
 pluggy>=1.6.0
 portpicker>=1.6.0
 pre_commit>=4.5.0
 prometheus-fastapi-instrumentator>=7.1.0
 prometheus_client>=0.23.1
 promise>=2.3
+prompt_toolkit>=3.0.52
 propcache>=0.4.1
 proto-plus>=1.26.1
 protobuf>=5.29.6
-psutil>=7.1.3
+psutil>=7.2.2
+ptyprocess>=0.7.0
+pure_eval>=0.2.3
 py-cpuinfo>=9.0.0
 py-spy>=0.4.1
 pyarrow>=22.0.0
@@ -316,18 +339,19 @@ sniffio>=1.3.1
 sortedcontainers>=2.4.0
 sse-starlette>=3.2.0
 starlette>=0.50.0
+stack-data>=0.6.3
 supervisor>=4.3.0
 sympy>=1.14.0
 tabulate>=0.9.0
-tenacity>=9.1.2
-tensorboard>=2.19.0
+tenacity>=9.1.4
+tensorboard>=2.20.0
 tensorboard-data-server>=0.7.2
 tensorboard-plugin-profile>=2.13.0
 tensorboardX>=2.6.4
-tensorflow>=2.19.1
+tensorflow>=2.20.0
 tensorflow-datasets>=4.9.9
 tensorflow-metadata>=1.17.2
-tensorflow-text>=2.19.0
+tensorflow-text>=2.20.0
 tensorstore>=0.1.79
 termcolor>=3.2.0
 tiktoken>=0.12.0
@@ -339,8 +363,10 @@ toolz>=1.1.0
 torch>=2.9.0
 torchax>=0.0.11
 torchvision>=0.24.0
+tornado>=6.5.4
 tpu-info>=0.7.1
-tqdm>=4.67.1
+tqdm>=4.67.3
+traitlets>=5.14.3
 transformers>=4.57.1
 treescope>=0.1.10
 triton>=3.5.0
@@ -358,9 +384,11 @@ uvloop>=0.22.1
 virtualenv>=20.35.4
 wadler_lindig>=0.1.7
 watchfiles>=1.1.1
+wcwidth>=0.6.0
 websockets>=15.0.1
 Werkzeug>=3.1.3
 wheel>=0.46.3
+widgetsnbextension>=4.0.15
 wrapt>=2.0.1
 xgrammar>=0.1.29
 xprof>=2.21.1