Merge pull request #3281 from AI-Hypercomputer:anisha-update-nb

PiperOrigin-RevId: 878111193

Author: Google-ML-Automation

.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

dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt

@@ -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

docs/tutorials/posttraining/rl.md

@@ -44,73 +44,45 @@ Let's get started!
 
 ## Create virtual environment and Install MaxText dependencies
 
-If you have already completed the
-[MaxText installation](../../install_maxtext.md), you can skip to the next
-section for post-training dependencies installations. Otherwise, please install
-`MaxText` using the following commands before proceeding.
-
 ```bash
-# 1. Clone the repository
-git clone https://github.com/AI-Hypercomputer/maxtext.git
-cd maxtext
-
-# 2. Create virtual environment
+# Create a virtual environment
 export VENV_NAME=<your virtual env name> # e.g., maxtext_venv
 pip install uv
 uv venv --python 3.12 --seed $VENV_NAME
 source $VENV_NAME/bin/activate
-
-# 3. Install dependencies in editable mode
-uv pip install -e .[tpu] --resolution=lowest
-install_maxtext_github_deps
 ```
 
-## Install Post-Training dependencies
+### Option 1: From PyPI releases (Recommended)
 
-### Option 1: From PyPI releases
+Run the following commands to get all the necessary installations.
 
-> **Caution:** RL in MaxText is currently broken with PyPI releases of
-> post-training dependencies. We are working on fixing this and recommend
-> following [Option 2: From Github](#option-2-from-github) in the meantime.
+```bash
+uv pip install maxtext[tpu-post-train] --resolution=lowest
+install_maxtext_tpu_post_train_extra_deps
+```
 
-Next, run the following bash script to get all the necessary installations
-inside the virtual environment (for e.g., `maxtext_venv`). This will take few
-minutes. Follow along the installation logs and look out for any issues!
+It installs MaxText and then for post-training, it installs primarily the following:
 
-```
-bash tools/setup/setup_post_training_requirements.sh
-```
+a. [Tunix](https://github.com/google/tunix) as the LLM Post-Training Library, and
 
-Primarily, it installs `Tunix`, and `vllm-tpu` which is
+b. `vllm-tpu` which is
 [vllm](https://github.com/vllm-project/vllm) and
 [tpu-inference](https://github.com/vllm-project/tpu-inference) and thereby
 providing TPU inference for vLLM, with unified JAX and PyTorch support.
 
 ### Option 2: From Github
 
-You can also locally git clone [tunix](https://github.com/google/tunix) and
-install using the instructions
-[here](https://github.com/google/tunix?tab=readme-ov-file#installation).
-Similarly install [vllm](https://github.com/vllm-project/vllm) and
-[tpu-inference](https://github.com/vllm-project/tpu-inference) from source
-following the instructions
-[here](https://docs.vllm.ai/projects/tpu/en/latest/getting_started/installation/#install-from-source).
-To get a set of compatible commit IDs for `maxtext`, `tunix`, `tpu-inference`,
-and `vllm`, follow these steps:
-
-1. Navigate to the
-   [MaxText Package Tests](https://github.com/AI-Hypercomputer/maxtext/actions/workflows/build_and_test_maxtext.yml?query=event%3Aschedule)
-   GitHub Actions workflow.
-
-2. Select the latest successful run.
+For using a version newer than the latest PyPI release, you could also install the latest vetted versions of the dependencies from MaxText in the following way:
 
-3. Within the workflow run, find and click on the `maxtext_jupyter_notebooks (py312)` job, then expand the `run` job.
-
-4. Locate the `Record Commit IDs` step. The commit SHAs for `maxtext`, `tunix`,
-   `tpu-inference`, and `vllm` that were used in that successful run are listed
-   in the logs of this step.
+```bash
+# 1. Clone the repository
+git clone https://github.com/AI-Hypercomputer/maxtext.git
+cd maxtext
 
-5. Prior to installation, ensure that the `maxtext`, `tunix`, `vllm`, and `tpu-inference` repositories are synchronized to the specific commits recorded from the CI logs. For each repository, use the following command to switch to the correct commit: `git checkout <commit_id>`.
+# 2. Install dependencies in editable mode
+uv pip install -e .[tpu-post-train] --resolution=lowest
+install_maxtext_tpu_post_train_extra_deps
+```
 
 ## Setup environment variables
 

docs/tutorials/posttraining/sft.md

@@ -24,22 +24,21 @@ We use [Tunix](https://github.com/google/tunix), a JAX-based library designed fo
 
 In this tutorial we use a single host TPU VM such as `v6e-8/v5p-8`. Let's get started!
 
-## Install dependencies
+## Install MaxText and Post-Training dependencies
 
-```sh
-# 1. Clone the repository
-git clone https://github.com/AI-Hypercomputer/maxtext.git
-cd maxtext
-
-# 2. Create virtual environment
+```bash
+# Create a virtual environment
 export VENV_NAME=<your virtual env name> # e.g., maxtext_venv
 pip install uv
 uv venv --python 3.12 --seed $VENV_NAME
 source $VENV_NAME/bin/activate
+```
+
+Run the following commands to get all the necessary installations.
 
-# 3. Install dependencies in editable mode
-uv pip install -e .[tpu] --resolution=lowest
-bash tools/setup/setup_post_training_requirements.sh
+```bash
+uv pip install maxtext[tpu-post-train] --resolution=lowest
+install_maxtext_tpu_post_train_extra_deps
 ```
 
 ## Setup environment variables

src/install_maxtext_extra_deps/install_post_train_extra_deps.py

@@ -35,7 +35,7 @@ def main():
   """
   script_dir = Path(__file__).resolve().parent
 
-  os.environ['VLLM_TARGET_DEVICE'] = 'tpu'
+  os.environ["VLLM_TARGET_DEVICE"] = "tpu"
 
   # Adjust this path if your extra_post_train_deps_from_github.txt is in a different location,
   # e.g., script_dir / "data" / "extra_post_train_deps_from_github.txt"
@@ -67,12 +67,12 @@ def main():
 
   local_vllm_install_command = [
       sys.executable,  # Use the current Python executable's pip to ensure the correct environment
-        "-m",
-        "uv",
-        "pip",
-        "install",
-        "src/maxtext/integration/vllm",
-        "--no-deps",
+      "-m",
+      "uv",
+      "pip",
+      "install",
+      "src/maxtext/integration/vllm",  # MaxText on vllm installations
+      "--no-deps",
   ]
 
   print(f"Installing extra dependencies from '{extra_deps_file}' using uv...")