Merge pull request #3000 from AI-Hypercomputer:chengnuojin-pr-3000

PiperOrigin-RevId: 863399027

Author: SurbhiJainUSC

.github/workflows/run_jupyter_notebooks.yml

@@ -87,6 +87,7 @@ jobs:
       - name: Run Post-Training Notebooks
         shell: bash
         env:
+          PYTHONPATH: "${{ github.workspace }}/src"
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
           MAXTEXT_REPO_ROOT=$(pwd)

docs/guides/run_python_notebook.md

@@ -19,6 +19,7 @@ Before starting, make sure you have:
 - ✅ Basic familiarity with Jupyter, Python, and Git
 
 **For Method 2 (Visual Studio Code) and Method 3 (Local Jupyter Lab) only:**
+
 - ✅ A Google Cloud Platform (GCP) account with billing enabled
 - ✅ TPU quota available in your region (check under IAM & Admin → Quotas)
 - ✅ `tpu.nodes.create` permission to create a TPU VM
@@ -36,16 +37,18 @@ Currently, this method only supports the **`sft_qwen3_demo.ipynb`** notebook, wh
 
 Before proceeding, please verify that the specific notebook you are running works reliably on the free-tier TPU resources. If you encounter frequent disconnections or resource limitations, you may need to:
 
-* Upgrade to a Colab Pro or Pro+ subscription for more stable and powerful TPU access.
+- Upgrade to a Colab Pro or Pro+ subscription for more stable and powerful TPU access.
 
-* Move to local Jupyter Lab setup method with access to a powerful TPU machine.
+- Move to local Jupyter Lab setup method with access to a powerful TPU machine.
 
 ### Step 1: Choose an Example
+
 1.a. Visit the [MaxText examples directory](https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/examples) on Github.
 
 1.b. Find the notebook you want to run (e.g., `sft_qwen3_demo.ipynb`) and copy its URL.
 
 ### Step 2: Import into Colab
+
 2.a. Go to [Google Colab](https://colab.research.google.com/) and sign in.
 
 2.b. Select **File** -> **Open Notebook**.
@@ -63,9 +66,11 @@ Before proceeding, please verify that the specific notebook you are running work
 3.c. Click **Save**
 
 ### Step 4: Run the Notebook
+
 Follow the instructions within the notebook cells to install dependencies and run the training/inference.
 
 ## Method 2: Visual Studio Code with TPU (Recommended)
+
 Running Jupyter notebooks in Visual Studio Code (VS Code) provides a powerful, interactive environment that combines the flexibility of notebooks with the robust features of a code editor. Follow these steps to get your environment up and running.
 
 ### Step 1: Set Up TPU VM
@@ -75,9 +80,10 @@ In Google Cloud Console, create a standalone TPU VM:
 1.a. **Compute Engine** → **TPUs** → **Create TPU**
 
 1.b. Example config:
-   - **Name:** `maxtext-tpu-node`
-   - **TPU type:** Choose your desired TPU type
-   - **Runtime Version:** `tpu-ubuntu2204-base` (or other compatible runtime)
+
+- **Name:** `maxtext-tpu-node`
+- **TPU type:** Choose your desired TPU type
+- **Runtime Version:** `tpu-ubuntu2204-base` (or other compatible runtime)
 
 ### Step 2: SSH to TPU-VM via VS Code
 
@@ -86,11 +92,12 @@ In Google Cloud Console, create a standalone TPU VM:
 2.b. Follow [Connect to a remote host](https://code.visualstudio.com/docs/remote/ssh#_connect-to-a-remote-host) guide to connect to your TPU-VM via VS Code.
 
 ### Step 3. Install Necessary Extensions on VS Code
+
 To enable notebook support, you must install two official extensions from the VS Code Marketplace:
 
-* Python Extension: Provides support for the Python language.
+- Python Extension: Provides support for the Python language.
 
-* Jupyter Extension: Enables you to create, edit, and run `.ipynb` files directly inside VS Code.
+- Jupyter Extension: Enables you to create, edit, and run `.ipynb` files directly inside VS Code.
 
 To install, click the `Extensions` icon on the left sidebar (or press `Ctrl+Shift+X` or `Cmd+Shift+X`), search for `Jupyter` and `Python`, and click `Install`.
 
@@ -99,6 +106,7 @@ To install, click the `Extensions` icon on the left sidebar (or press `Ctrl+Shif
 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.
 
 ### Step 5: Install the necessary library for Jupyter
+
 Jupyter requires a kernel to execute code. This kernel is tied to a specific Python environment. Open your terminal inside VS Code and run:
 
 ```bash
@@ -110,9 +118,9 @@ uv pip install ipykernel
 Before you can run the notebook, you must tell VS Code which Python environment to use.
 
 1. Look at the top-right corner of the notebook editor.
-2. Click `Select Kernel`.
-3. Choose Python Environments and select the virtual environment you created in Step 4.
-4. Open [available post-training notebooks in MaxText](#available-examples) inside VS Code and run the jupyter notebook cells.
+1. Click `Select Kernel`.
+1. Choose Python Environments and select the virtual environment you created in Step 4.
+1. Open [available post-training notebooks in MaxText](#available-examples) inside VS Code and run the jupyter notebook cells.
 
 ## Method 3: Local Jupyter Lab with TPU (Recommended)
 
@@ -125,12 +133,15 @@ In Google Cloud Console, create a standalone TPU VM:
 1.a. **Compute Engine** → **TPUs** → **Create TPU**
 
 1.b. Example config:
-   - **Name:** `maxtext-tpu-node`
-   - **TPU type:** Choose your desired TPU type
-   - **Runtime Version:** `tpu-ubuntu2204-base` (or other compatible runtime)
+
+- **Name:** `maxtext-tpu-node`
+- **TPU type:** Choose your desired TPU type
+- **Runtime Version:** `tpu-ubuntu2204-base` (or other compatible runtime)
 
 ### Step 2: Connect with Port Forwarding
+
 Run the following command on your local machine:
+
 > **Note**: The `--` separator before the `-L` flag is required. This tunnels the remote port 8888 to your local machine securely.
 
 ```bash
@@ -170,13 +181,15 @@ jupyter lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root
 ```
 
 ### Step 7: Access the Notebook
+
 7.a. Look at the terminal output for a URL that looks like: `http://127.0.0.1:8888/lab?token=...`.
 
 7.b. Copy that URL.
 
 7.c. Paste it into your **local computer's browser**.
-   * **Important:** If you changed the port in Step 2 (e.g., to `9999`), you must manually replace `8888` in the URL with `9999`.
-   * *Example:* `http://127.0.0.1:9999/lab?token=...`
+
+- **Important:** If you changed the port in Step 2 (e.g., to `9999`), you must manually replace `8888` in the URL with `9999`.
+- *Example:* `http://127.0.0.1:9999/lab?token=...`
 
 7.d. Once the interface opens in your browser, Click on the current kernel name (e.g., `Python 3 (ipykernel)`).
 
@@ -197,13 +210,13 @@ jupyter lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root
 
 ## Common Pitfalls & Debugging
 
-| Issue | Solution |
-|-------|----------|
-| ❌ TPU runtime mismatch | Check TPU runtime version matches VM image |
-| ❌ Colab disconnects | Save checkpoints to GCS or Drive regularly |
-| ❌ "RESOURCE_EXHAUSTED" errors | Use smaller batch size or v5e-8 instead of v5e-1 |
-| ❌ Firewall blocked | Ensure port 8888 open, or always use SSH tunneling |
-| ❌ Path confusion | In Colab use `/content/maxtext`; in TPU VM use `~/maxtext` |
+| Issue                          | Solution                                                   |
+| ------------------------------ | ---------------------------------------------------------- |
+| ❌ TPU runtime mismatch        | Check TPU runtime version matches VM image                 |
+| ❌ Colab disconnects           | Save checkpoints to GCS or Drive regularly                 |
+| ❌ "RESOURCE_EXHAUSTED" errors | Use smaller batch size or v5e-8 instead of v5e-1           |
+| ❌ Firewall blocked            | Ensure port 8888 open, or always use SSH tunneling         |
+| ❌ Path confusion              | In Colab use `/content/maxtext`; in TPU VM use `~/maxtext` |
 
 ## Support and Resources
 
@@ -217,9 +230,9 @@ jupyter lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root
 If you encounter issues or have improvements for this guide, please:
 
 1. Open an issue on the MaxText repository
-2. Submit a pull request with your improvements
-3. Share your experience in the discussions
+1. Submit a pull request with your improvements
+1. Share your experience in the discussions
 
----
+______________________________________________________________________
 
-**Happy Training! 🚀**
+**Happy Training! 🚀**