Skip to content

Commit 6f3cf71

Browse files
SurbhiJainUSCSurbhiJainUSC
committed
Fix MaxText installation instructions in all jupyter notebooks and update doc
1 parent ef98ae5 commit 6f3cf71

5 files changed

Lines changed: 74 additions & 40 deletions

File tree

‎.github/workflows/run_jupyter_notebooks.yml‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -60,12 +60,9 @@ jobs:
6060
# 1. Create virtual environment
6161
python3 -m uv venv --seed
6262
source .venv/bin/activate
63-
maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
6463
6564
# 2. Install MaxText package and all the post training dependencies
66-
uv pip install ${maxtext_wheel}[tpu-post-train] --resolution=lowest
67-
#TODO: @mazumdera: replace this with the following after release
68-
# uv pip install maxtext[tpu-post-train] --resolution=lowest
65+
uv pip install maxtext[tpu-post-train] --resolution=lowest
6966
install_maxtext_tpu_post_train_extra_deps
7067
.venv/bin/python3 -m ipykernel install --user --name maxtext_venv
7168
@@ -84,7 +81,7 @@ jobs:
8481
8582
for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/{sft,rl}*.ipynb; do
8683
filename=$(basename "$notebook")
87-
if [[ "$filename" == "sft_qwen3_demo.ipynb" || "$filename" == "sft_llama3_demo_gpu.ipynb" ]]; then
84+
if [[ "$filename" == "sft_llama3_demo_gpu.ipynb" ]]; then
8885
echo "Skipping $filename"
8986
continue
9087
fi

‎docs/guides/run_python_notebook.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ To install, click the `Extensions` icon on the left sidebar (or press `Ctrl+Shif
103103

104104
### Step 4: Install MaxText and Dependencies
105105

106-
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.
106+
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-pypi-recommended) to install MaxText and its dependencies inside a dedicated virtual environment.
107107

108108
### Step 5: Install the necessary library for Jupyter
109109

@@ -162,7 +162,7 @@ pip3 install jupyterlab
162162

163163
### Step 4: Install MaxText and Dependencies
164164

165-
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.
165+
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-pypi-recommended) to install MaxText and its dependencies inside a dedicated virtual environment.
166166

167167
### Step 5: Register virtual environment as a Jupyter Kernel
168168

‎src/maxtext/examples/rl_llama3_demo.ipynb‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -89,17 +89,34 @@
8989
"source": [
9090
"## Installation: MaxText and Post training Dependencies\n",
9191
"\n",
92-
"**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",
93-
"\n",
94-
"\n",
95-
"**Running the notebook on Google Colab:** Run the above-mentioned commands in the Colab and then restart session\n",
96-
"\n",
92+
"**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-pypi-recommended) before proceeding. And then run the notebook using that virtual environment."
93+
]
94+
},
95+
{
96+
"cell_type": "code",
97+
"execution_count": null,
98+
"metadata": {},
99+
"outputs": [],
100+
"source": [
101+
"if IN_COLAB:\n",
102+
" # Install uv, a fast Python package installer\n",
103+
" !pip install uv\n",
104+
" \n",
105+
" # Install MaxText and post-training dependencies\n",
106+
" !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
107+
" !python3 -m MaxText.install_maxtext_extra_deps"
108+
]
109+
},
110+
{
111+
"cell_type": "markdown",
112+
"metadata": {},
113+
"source": [
97114
"**Session restart Instructions for Colab:**\n",
98115
"1. Navigate to the menu at the top of the screen.\n",
99116
"2. Click on **Runtime**.\n",
100117
"3. Select **Restart session** from the dropdown menu.\n",
101118
"\n",
102-
"You will be asked to confirm the action in a pop-up dialog. Click on **Yes**.\n"
119+
"You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
103120
]
104121
},
105122
{

‎src/maxtext/examples/sft_llama3_demo_tpu.ipynb‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -101,17 +101,34 @@
101101
"source": [
102102
"## Installation: MaxText and Post training Dependencies\n",
103103
"\n",
104-
"**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",
105-
"\n",
106-
"\n",
107-
"**Running the notebook on Google Colab:** Run the above-mentioned commands in the Colab and then restart session\n",
108-
"\n",
104+
"**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-pypi-recommended) before proceeding. And then run the notebook using that virtual environment."
105+
]
106+
},
107+
{
108+
"cell_type": "code",
109+
"execution_count": null,
110+
"metadata": {},
111+
"outputs": [],
112+
"source": [
113+
"if IN_COLAB:\n",
114+
" # Install uv, a fast Python package installer\n",
115+
" !pip install uv\n",
116+
" \n",
117+
" # Install MaxText and post-training dependencies\n",
118+
" !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
119+
" !python3 -m MaxText.install_maxtext_extra_deps"
120+
]
121+
},
122+
{
123+
"cell_type": "markdown",
124+
"metadata": {},
125+
"source": [
109126
"**Session restart Instructions for Colab:**\n",
110127
"1. Navigate to the menu at the top of the screen.\n",
111128
"2. Click on **Runtime**.\n",
112129
"3. Select **Restart session** from the dropdown menu.\n",
113130
"\n",
114-
"You will be asked to confirm the action in a pop-up dialog. Click on **Yes**.\n"
131+
"You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
115132
]
116133
},
117134
{

‎src/maxtext/examples/sft_qwen3_demo.ipynb‎

Lines changed: 24 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -95,40 +95,42 @@
9595
" IN_COLAB = False"
9696
]
9797
},
98+
{
99+
"cell_type": "markdown",
100+
"metadata": {
101+
"id": "D9ms-jTSZUQL"
102+
},
103+
"source": [
104+
"## Installation: MaxText and Post training Dependencies\n",
105+
"\n",
106+
"**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-pypi-recommended) before proceeding. And then run the notebook using that virtual environment."
107+
]
108+
},
98109
{
99110
"cell_type": "code",
100111
"execution_count": null,
101112
"metadata": {},
102113
"outputs": [],
103114
"source": [
104-
"try:\n",
105-
" import google.colab\n",
106-
" print(\"Running the notebook in Google Colab\")\n",
107-
" IN_COLAB = True\n",
108-
"except ImportError:\n",
109-
" print(\"Running the notebook on JupyterLab\")\n",
110-
" IN_COLAB = False"
115+
"if IN_COLAB:\n",
116+
" # Install uv, a fast Python package installer\n",
117+
" !pip install uv\n",
118+
" \n",
119+
" # Install MaxText and post-training dependencies\n",
120+
" !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
121+
" !python3 -m MaxText.install_maxtext_extra_deps"
111122
]
112123
},
113124
{
114125
"cell_type": "markdown",
115-
"metadata": {
116-
"id": "D9ms-jTSZUQL"
117-
},
126+
"metadata": {},
118127
"source": [
119-
"## Installation: MaxText and Post training Dependencies\n",
120-
"\n",
121-
"**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",
122-
"\n",
123-
"\n",
124-
"**Running the notebook on Google Colab:** Run the above-mentioned commands in the Colab and then restart session\n",
125-
"\n",
126128
"**Session restart Instructions for Colab:**\n",
127129
"1. Navigate to the menu at the top of the screen.\n",
128130
"2. Click on **Runtime**.\n",
129131
"3. Select **Restart session** from the dropdown menu.\n",
130132
"\n",
131-
"You will be asked to confirm the action in a pop-up dialog. Click on **Yes**.\n"
133+
"You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
132134
]
133135
},
134136
{
@@ -168,6 +170,7 @@
168170
"from datetime import datetime\n",
169171
"from flax import nnx\n",
170172
"from huggingface_hub import login\n",
173+
"from jax import numpy as jnp\n",
171174
"\n",
172175
"print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")"
173176
]
@@ -192,10 +195,9 @@
192195
},
193196
"outputs": [],
194197
"source": [
195-
"try:\n",
196-
" from google.colab import userdata\n",
198+
"if IN_COLAB:\n",
197199
" HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
198-
"except ImportError:\n",
200+
"else:\n",
199201
" HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
200202
"\n",
201203
"# If not found in the environment, prompt the user for input securely\n",
@@ -438,6 +440,7 @@
438440
" rollout_vllm_hbm_utilization=0.8,\n",
439441
" rollout_vllm_init_with_random_weights=True,\n",
440442
" rollout_vllm_tpu_backend_type=\"jax\",\n",
443+
" data_type=jnp.bfloat16,\n",
441444
" ),\n",
442445
")"
443446
]

0 commit comments

Comments
 (0)