Revise Colab setup instructions for TPU: streamline steps for importing notebooks and connecting to local Jupyter Lab.

Author: RexBearIU

docs/tutorials/how_to_run_colabs.md

@@ -27,40 +27,28 @@ Before starting, make sure you have:
 
 This is the fastest way to run MaxText without managing infrastructure.
 
-### Step 1: Open Google Colab
-
-1. Go to [Google Colab](https://colab.research.google.com/)
-2. Sign in → New Notebook
-
-### Step 2: Enable TPU Runtime
-
+### Step 1: Choose an Example
+1. Visit the MaxText Examples directory on GitHub:
+   [https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/examples](https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/examples)
+2. Find the notebook you want to run (e.g., `sft_qwen3_demo.ipynb`) and copy its URL.
+
+### Step 2: Import into Colab
+1. Go to [Google Colab](https://colab.research.google.com/) and sign in.
+2. Select **File** -> **Open Notebook**.
+3. Select the **GitHub** tab.
+4. Paste the target `.ipynb` link you copied in Step 1 and press Enter.
+5. Click on the notebook name to open it.
+
+### Step 3: Enable TPU Runtime
 1. **Runtime** → **Change runtime type**
 2. Set **Hardware accelerator** → **TPU**
 3. Select TPU version:
    - **v5e-8** → recommended for most MaxText examples, but it's a paid option
    - **v5e-1** → free tier option (slower, but works for Qwen-0.6B demos)
 4. Click **Save**
 
-### Step 3: Upload & Prepare MaxText
-
-Upload notebooks or mount your GitHub repo
-
-> **Note:** In Colab, the repo root will usually be `/content/maxtext`
-
-**Example:**
-```bash
-!git clone https://github.com/AI-Hypercomputer/maxtext.git
-%cd maxtext
-```
-
-### Step 4: Run Examples
-
-1. Open `src/MaxText/examples/`
-2. Try:
-   - `sft_qwen3_demo.ipynb`
-   - `sft_llama3_demo.ipynb`
-   - `rl_llama3_demo.ipynb` (GRPO/GSPO training)
-
+### Step 4: Run the Notebook
+Follow the instructions within the notebook cells to install dependencies and run the training/inference.
 
 > ⚡ **Tip:** If Colab disconnects, re-enable TPU and re-run setup cells. Save checkpoints to GCS or Drive.
 
@@ -80,12 +68,16 @@ In Google Cloud Console:
    - **TPU type:** `v5e-8` (or `v6p-8` for newer hardware)
    - **Runtime Version:** `tpu-ubuntu-alpha-*` (matches your VM image)
 
-### Step 2: Connect to TPU VM
+### Step 2: Connect with Port Forwarding
+Run this command to SSH into the TPU VM while simultaneously creating a secure tunnel for Jupyter.
+> **Note**: The `--` separator before the `-L` flag is required. This tunnels the remote port 8888 to your local machine securely.
 
 ```bash
-gcloud compute tpus tpu-vm ssh maxtext-tpu-node --zone=YOUR_ZONE
+gcloud compute tpus tpu-vm ssh maxtext-tpu-node --zone=YOUR_ZONE -- -L 8888:localhost:8888
 ```
 
+> **Note**: If you get a "bind: Address already in use" error, it means port 8888 is busy on your local computer. Change the first number to a different port, e.g., -L 9999:localhost:8888. You will then access Jupyter at localhost:9999.
+
 ### Step 3: Install Dependencies
 
 ```bash
@@ -100,23 +92,23 @@ pip3 install jupyterlab
 jupyter lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root
 ```
 
-Copy the URL with token from terminal
-
-### Step 5: Secure Access
-
-#### Option A: SSH Tunnel (Recommended)
-
-```bash
-gcloud compute tpus tpu-vm ssh maxtext-tpu-node --zone=YOUR_ZONE -- -L 8888:localhost:8888
-```
-
-Then open → `http://localhost:8888`
+### Step 5: Access the Notebook
+1. Look at the terminal output for a URL that looks like: `http://127.0.0.1:8888/lab?token=...`
+2. Copy that URL.
+3. 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=...`
 
 
 ## Method 3: Colab + Local Jupyter Lab Hybrid
 
-Set up Jupyter Lab as in step 2.
-Use the link for Jupyter Lab as a link for "Connect to a local runtime" in Collab - at the dropdown where you select the runtime.
+Set up Jupyter Lab as in Method 2 (complete Steps 1-5).
+
+### To connect Colab to your local Jupyter Lab:
+1. Copy the Jupyter Lab URL from Method 2, Step 5 output (format: `http://127.0.0.1:8888/lab?token=...` or the custom port you used).
+2. In Google Colab, click the **arrow next to the RAM/Disk indicator** in the top right corner.
+3. Select **"Connect to a local runtime"**.
+4. Paste the Jupyter Lab URL and click **Connect**.
 
 ## Available Examples
 
@@ -133,8 +125,13 @@ Use the link for Jupyter Lab as a link for "Connect to a local runtime" in Colla
 
 For interactive GRPO or GSPO training in Google Colab or Jupyter:
 
-1. **Open** `src/MaxText/examples/rl_llama3_demo.ipynb`
-2. **Enable TPU runtime** (Runtime → Change runtime type → TPU)
+1. **Open the Notebook**
+   * **Google Colab:** Go to **File** → **Open Notebook** → **GitHub**. Paste the URL for `src/MaxText/examples/rl_llama3_demo.ipynb`.
+   * **Jupyter Lab:** In the file browser, navigate to `src/MaxText/examples/` and open `rl_llama3_demo.ipynb`.
+
+2. **Enable TPU Runtime**
+   * **Google Colab:** Go to **Runtime** → **Change runtime type** → Select **TPU**.
+   * **Jupyter Lab:** Skip this step (the VM is already connected to the TPU).
 3. **Set `LOSS_ALGO`** to `"grpo"` for GRPO or `"gspo-token"` for GSPO
 4. **Run cells** to train Llama3.1-8B with GRPO or GSPO on GSM8K dataset