Fix MaxText installation instructions in all jupyter notebooks and update doc

SurbhiJainUSC · SurbhiJainUSC · commit ed3c5698cae5 · 2026-03-13T18:03:13.000Z
diff --git a/.github/workflows/run_jupyter_notebooks.yml b/.github/workflows/run_jupyter_notebooks.yml
@@ -64,8 +64,6 @@ jobs:
 
           # 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
           
diff --git a/docs/guides/run_python_notebook.md b/docs/guides/run_python_notebook.md
@@ -103,7 +103,7 @@ To install, click the `Extensions` icon on the left sidebar (or press `Ctrl+Shif
 
 ### Step 4: Install MaxText and Dependencies
 
-To execute post-training notebooks on your TPU-VM, follow the official [MaxText installation guides](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl.html#create-virtual-environment-and-install-maxtext-dependencies) to install MaxText and its dependencies inside a dedicated virtual environment.
+To execute post-training notebooks on your TPU-VM, follow the official [MaxText installation guides](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source) to install MaxText and its dependencies inside a dedicated virtual environment.
 
 ### Step 5: Install the necessary library for Jupyter
 
@@ -162,7 +162,7 @@ pip3 install jupyterlab
 
 ### Step 4: Install MaxText and Dependencies
 
-To execute post-training notebooks on your TPU-VM, follow the official [MaxText installation guides](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl.html#create-virtual-environment-and-install-maxtext-dependencies) to install MaxText and its dependencies inside a dedicated virtual environment.
+To execute post-training notebooks on your TPU-VM, follow the official [MaxText installation guides](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source) to install MaxText and its dependencies inside a dedicated virtual environment.
 
 ### Step 5: Register virtual environment as a Jupyter Kernel
 
diff --git a/docs/install_maxtext.md b/docs/install_maxtext.md
@@ -38,15 +38,15 @@ source maxtext_venv/bin/activate
 #      installation option from this list to fit your use case.
 
 # Option 1: Installing maxtext[tpu]
-uv pip install maxtext[tpu] --resolution=lowest
+uv pip install "maxtext[tpu]>=0.2.0" --resolution=lowest
 install_maxtext_tpu_github_deps
 
 # Option 2: Installing maxtext[cuda12]
-uv pip install maxtext[cuda12] --resolution=lowest
+uv pip install "maxtext[cuda12]>=0.2.0" --resolution=lowest
 install_maxtext_cuda12_github_dep
 
 # Option 3: Installing maxtext[tpu-post-train]
-uv pip install maxtext[tpu-post-train] --resolution=lowest
+uv pip install "maxtext[tpu-post-train]>=0.2.0" --resolution=lowest
 install_maxtext_tpu_post_train_extra_deps
 ```
 > **Note:** The `install_maxtext_tpu_github_deps`, `install_maxtext_cuda12_github_dep`, and
diff --git a/docs/tutorials/posttraining/rl.md b/docs/tutorials/posttraining/rl.md
@@ -57,7 +57,7 @@ source ${VENV_NAME?}/bin/activate
 Run the following commands to get all the necessary installations.
 
 ```bash
-uv pip install maxtext[tpu-post-train] --resolution=lowest
+uv pip install "maxtext[tpu-post-train]>=0.2.0" --resolution=lowest
 install_maxtext_tpu_post_train_extra_deps
 ```
 
diff --git a/docs/tutorials/posttraining/sft.md b/docs/tutorials/posttraining/sft.md
@@ -37,7 +37,7 @@ source ${VENV_NAME?}/bin/activate
 Run the following commands to get all the necessary installations.
 
 ```bash
-uv pip install maxtext[tpu-post-train] --resolution=lowest
+uv pip install "maxtext[tpu-post-train]>=0.2.0" --resolution=lowest
 install_maxtext_tpu_post_train_extra_deps
 ```
 
diff --git a/src/maxtext/examples/rl_llama3_demo.ipynb b/src/maxtext/examples/rl_llama3_demo.ipynb
@@ -89,17 +89,38 @@
    "source": [
     "## Installation: MaxText and Post training Dependencies\n",
     "\n",
-    "**Running the notebook on Visual Studio or JupyterLab:** Create an virtual environment and install dependencies outside of the notebook using the commands in [MaxText installation and dependency setup guide](../../../docs/guides/run_python_notebook.md#step-4-install-maxtext-and-dependencies) before proceeding. And then run the notebook using that virtual environment.\n",
-    "\n",
-    "\n",
-    "**Running the notebook on Google Colab:** Run the above-mentioned commands in the Colab and then restart session\n",
+    "**Running the notebook on Visual Studio or JupyterLab:** Create an virtual environment and install post-training dependencies outside of the notebook using the commands in [MaxText installation and dependency setup guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source) before proceeding. And then run the notebook using that virtual environment."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if IN_COLAB:\n",
+    "    # Clone the MaxText repository\n",
+    "    !git clone https://github.com/AI-Hypercomputer/maxtext.git\n",
+    "    %cd maxtext\n",
     "\n",
+    "    # Install uv, a fast Python package installer\n",
+    "    !pip install uv\n",
+    "    \n",
+    "    # Install MaxText and post-training dependencies\n",
+    "    !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
+    "    !install_maxtext_tpu_post_train_extra_deps"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
     "**Session restart Instructions for Colab:**\n",
     "1.  Navigate to the menu at the top of the screen.\n",
     "2.  Click on **Runtime**.\n",
     "3.  Select **Restart session** from the dropdown menu.\n",
     "\n",
-    "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**.\n"
+    "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
    ]
   },
   {
@@ -119,20 +140,19 @@
     "import os\n",
     "import sys\n",
     "from pathlib import Path\n",
-    "import MaxText\n",
     "from huggingface_hub import login\n",
     "from etils import epath\n",
     "import jax\n",
     "\n",
     "from maxtext.trainers.post_train.rl.train_rl import rl_train, setup_configs_and_devices\n",
+    "from maxtext.utils.globals import MAXTEXT_REPO_ROOT, MAXTEXT_PKG_DIR\n",
     "\n",
     "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"0\"\n",
     "os.environ[\"SKIP_JAX_PRECOMPILE\"] = \"1\"  # Faster startup for vLLM\n",
     "# Suppress vLLM logging with a severity level below ERROR\n",
     "os.environ[\"VLLM_LOGGING_LEVEL\"] = \"ERROR\"\n",
     "\n",
-    "MAXTEXT_PKG_DIR = os.path.dirname(MaxText.__file__)\n",
-    "MAXTEXT_REPO_ROOT = os.sep.join([\"maxtext\" if p == \"MaxText\" else p for p in MAXTEXT_PKG_DIR.split(os.sep)])\n",
+    "\n",
     "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")"
    ]
   },
@@ -154,9 +174,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "if IN_COLAB:\n",
+    "try:\n",
+    "    from google.colab import userdata\n",
     "    HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
-    "else:\n",
+    "except ImportError:\n",
     "    HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
     "\n",
     "# If not found in the environment, prompt the user for input securely\n",
diff --git a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb
@@ -101,17 +101,38 @@
       "source": [
         "## Installation: MaxText and Post training Dependencies\n",
         "\n",
-        "**Running the notebook on Visual Studio or JupyterLab:** Create an virtual environment and install dependencies outside of the notebook using the commands in [MaxText installation and dependency setup guide](../../../docs/guides/run_python_notebook.md#step-4-install-maxtext-and-dependencies) before proceeding. And then run the notebook using that virtual environment.\n",
-        "\n",
-        "\n",
-        "**Running the notebook on Google Colab:** Run the above-mentioned commands in the Colab and then restart session\n",
-        "\n",
+        "**Running the notebook on Visual Studio or JupyterLab:** Create an virtual environment and install post-training dependencies outside of the notebook using the commands in [MaxText installation and dependency setup guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source) before proceeding. And then run the notebook using that virtual environment."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "if IN_COLAB:\n",
+        "    # Clone the MaxText repository\n",
+        "    !git clone https://github.com/AI-Hypercomputer/maxtext.git\n",
+        "    %cd maxtext\n",
+        "\n",
+        "    # Install uv, a fast Python package installer\n",
+        "    !pip install uv\n",
+        "    \n",
+        "    # Install MaxText and post-training dependencies\n",
+        "    !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
+        "    !install_maxtext_tpu_post_train_extra_deps"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
         "**Session restart Instructions for Colab:**\n",
         "1.  Navigate to the menu at the top of the screen.\n",
         "2.  Click on **Runtime**.\n",
         "3.  Select **Restart session** from the dropdown menu.\n",
         "\n",
-        "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**.\n"
+        "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
       ]
     },
     {
@@ -166,9 +187,10 @@
       },
       "outputs": [],
       "source": [
-        "if IN_COLAB:\n",
+        "try:\n",
+        "    from google.colab import userdata\n",
         "    HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
-        "else:\n",
+        "except ImportError:\n",
         "    HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
         "\n",
         "# If not found in the environment, prompt the user for input securely\n",
@@ -250,7 +272,6 @@
         "    print(\"Converting checkpoint from HuggingFace...\")\n",
         "    env = os.environ.copy()\n",
         "    env[\"JAX_PLATFORMS\"] = \"cpu\"\n",
-        "    # env[\"PYTHONPATH\"] = MAXTEXT_PKG_DIR\n",
         "\n",
         "    subprocess.run(\n",
         "        [\n",
diff --git a/src/maxtext/examples/sft_qwen3_demo.ipynb b/src/maxtext/examples/sft_qwen3_demo.ipynb
@@ -95,40 +95,46 @@
     "    IN_COLAB = False"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "D9ms-jTSZUQL"
+   },
+   "source": [
+    "## Installation: MaxText and Post training Dependencies\n",
+    "\n",
+    "**Running the notebook on Visual Studio or JupyterLab:** Create an virtual environment and install post-training dependencies outside of the notebook using the commands in [MaxText installation and dependency setup guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source) before proceeding. And then run the notebook using that virtual environment."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "try:\n",
-    "    import google.colab\n",
-    "    print(\"Running the notebook in Google Colab\")\n",
-    "    IN_COLAB = True\n",
-    "except ImportError:\n",
-    "    print(\"Running the notebook on JupyterLab\")\n",
-    "    IN_COLAB = False"
+    "if IN_COLAB:\n",
+    "    # Clone the MaxText repository\n",
+    "    !git clone https://github.com/AI-Hypercomputer/maxtext.git\n",
+    "    %cd maxtext\n",
+    "\n",
+    "    # Install uv, a fast Python package installer\n",
+    "    !pip install uv\n",
+    "    \n",
+    "    # Install MaxText and post-training dependencies\n",
+    "    !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
+    "    !install_maxtext_tpu_post_train_extra_deps"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "id": "D9ms-jTSZUQL"
-   },
+   "metadata": {},
    "source": [
-    "## Installation: MaxText and Post training Dependencies\n",
-    "\n",
-    "**Running the notebook on Visual Studio or JupyterLab:** Create an virtual environment and install dependencies outside of the notebook using the commands in [MaxText installation and dependency setup guide](../../../docs/guides/run_python_notebook.md#step-4-install-maxtext-and-dependencies) before proceeding. And then run the notebook using that virtual environment.\n",
-    "\n",
-    "\n",
-    "**Running the notebook on Google Colab:** Run the above-mentioned commands in the Colab and then restart session\n",
-    "\n",
     "**Session restart Instructions for Colab:**\n",
     "1.  Navigate to the menu at the top of the screen.\n",
     "2.  Click on **Runtime**.\n",
     "3.  Select **Restart session** from the dropdown menu.\n",
     "\n",
-    "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**.\n"
+    "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
    ]
   },
   {
@@ -153,7 +159,6 @@
     "import sys\n",
     "import transformers\n",
     "\n",
-    "import MaxText\n",
     "from maxtext.configs import pyconfig\n",
     "from maxtext.examples.sft_train_and_evaluate import evaluate_model, get_test_dataset\n",
     "from maxtext.integration.tunix.tunix_adapter import TunixMaxTextAdapter\n",
diff --git a/src/maxtext/examples/sft_train_and_evaluate.py b/src/maxtext/examples/sft_train_and_evaluate.py
@@ -30,7 +30,7 @@
 
 # Run the following commands to get all the necessary installations.
 
-uv pip install maxtext[tpu-post-train] --resolution=lowest
+uv pip install "maxtext[tpu-post-train]>=0.2.0" --resolution=lowest
 install_maxtext_tpu_post_train_extra_deps
 
 
