You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Copy file name to clipboardExpand all lines: docs/guides/run_python_notebook.md
+2-2Lines changed: 2 additions & 2 deletions
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -103,7 +103,7 @@ To install, click the `Extensions` icon on the left sidebar (or press `Ctrl+Shif
103
103
104
104
### Step 4: Install MaxText and Dependencies
105
105
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-source)and specifically follow `Option 3: Installing [tpu-post-train]`. This will ensure all post-training dependencies are installed inside your virtual environment.
107
107
108
108
### Step 5: Install the necessary library for Jupyter
109
109
@@ -162,7 +162,7 @@ pip3 install jupyterlab
162
162
163
163
### Step 4: Install MaxText and Dependencies
164
164
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-source)and specifically follow `Option 3: Installing [tpu-post-train]`. This will ensure all post-training dependencies are installed inside your virtual environment.
166
166
167
167
### Step 5: Register virtual environment as a Jupyter Kernel
> **Note:** The `install_maxtext_tpu_github_deps`, `install_maxtext_cuda12_github_dep`, and
53
-
`install_maxtext_tpu_post_train_extra_deps` commands are temporarily required to install dependencies directly from GitHub
54
-
that are not yet available on PyPI. As shown above, choose the one that corresponds to your use case.
55
+
> `install_maxtext_tpu_post_train_extra_deps` commands are temporarily required to install dependencies directly from GitHub
56
+
> that are not yet available on PyPI. As shown above, choose the one that corresponds to your use case.
55
57
56
58
> **Note:** The maxtext package contains a comprehensive list of all direct and transitive dependencies, with lower bounds, generated by [seed-env](https://github.com/google-ml-infra/actions/tree/main/python_seed_env). We highly recommend the `--resolution=lowest` flag. It instructs `uv` to install the specific, tested versions of dependencies defined by MaxText, rather than the latest available ones. This ensures a consistent and reproducible environment, which is critical for stable performance and for running benchmarks.
57
59
58
60
## From Source
61
+
59
62
If you plan to contribute to MaxText or need the latest unreleased features, install from source.
60
63
61
64
```bash
@@ -98,11 +101,11 @@ Please keep dependencies updated throughout development. This will allow each co
98
101
99
102
To update dependencies, you will follow these general steps:
100
103
101
-
1.**Modify Base Requirements**: Update the desired dependencies in `base_requirements/requirements.txt` or the hardware-specific files (`base_requirements/tpu-base-requirements.txt`, `base_requirements/gpu-base-requirements.txt`).
102
-
2.**Generate New Files**: Run the `seed-env` CLI tool to generate new, fully-pinned requirements files based on your changes.
103
-
3.**Update Project Files**: Copy the newly generated files into the `generated_requirements/` directory.
104
-
4.**Handle GitHub Dependencies**: Move any dependencies that are installed directly from GitHub from the generated files to `src/install_maxtext_extra_deps/extra_deps_from_github.txt`.
105
-
5.**Verify**: Test the new dependencies to ensure the project installs and runs correctly.
104
+
1.**Modify Base Requirements**: Update the desired dependencies in `base_requirements/requirements.txt` or the hardware-specific files (`base_requirements/tpu-base-requirements.txt`, `base_requirements/gpu-base-requirements.txt`).
105
+
2.**Generate New Files**: Run the `seed-env` CLI tool to generate new, fully-pinned requirements files based on your changes.
106
+
3.**Update Project Files**: Copy the newly generated files into the `generated_requirements/` directory.
107
+
4.**Handle GitHub Dependencies**: Move any dependencies that are installed directly from GitHub from the generated files to `src/install_maxtext_extra_deps/extra_deps_from_github.txt`.
108
+
5.**Verify**: Test the new dependencies to ensure the project installs and runs correctly.
106
109
107
110
The following sections provide detailed instructions for each step.
108
111
@@ -154,25 +157,26 @@ seed-env \
154
157
155
158
After generating the new requirements, you need to update the files in the MaxText repository.
156
159
157
-
1.**Copy the generated files:**
158
-
- Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`.
159
-
- Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`.
160
+
1.**Copy the generated files:**
161
+
162
+
- Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`.
163
+
- Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`.
Currently, MaxText uses a few dependencies, such as `mlperf-logging` and `google-jetstream`, that are installed directly from GitHub source. These are defined in `base_requirements/requirements.txt`, and the `seed-env` tool will carry them over to the generated requirements files.
Currently, MaxText uses a few dependencies, such as `mlperf-logging` and `google-jetstream`, that are installed directly from GitHub source. These are defined in `base_requirements/requirements.txt`, and the `seed-env` tool will carry them over to the generated requirements files.
163
167
164
168
## Step 5: Verify the New Dependencies
165
169
166
170
Finally, test that the new dependencies install correctly and that MaxText runs as expected.
167
171
168
-
1.**Create a clean environment:** It's best to start with a fresh Python virtual environment.
172
+
1.**Create a clean environment:** It's best to start with a fresh Python virtual environment.
169
173
170
174
```bash
171
175
uv venv --python 3.12 --seed maxtext_venv
172
176
source maxtext_venv/bin/activate
173
177
```
174
178
175
-
2.**Run the setup script:** Execute `bash setup.sh` to install the new dependencies.
179
+
2.**Run the setup script:** Execute `bash setup.sh` to install the new dependencies.
Copy file name to clipboardExpand all lines: src/maxtext/examples/rl_llama3_demo.ipynb
+32-11Lines changed: 32 additions & 11 deletions
Original file line number
Diff line number
Diff line change
@@ -89,17 +89,38 @@
89
89
"source": [
90
90
"## Installation: MaxText and Post training Dependencies\n",
91
91
"\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",
92
+
"**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it."
Copy file name to clipboardExpand all lines: src/maxtext/examples/sft_llama3_demo_tpu.ipynb
+30-9Lines changed: 30 additions & 9 deletions
Original file line number
Diff line number
Diff line change
@@ -101,17 +101,38 @@
101
101
"source": [
102
102
"## Installation: MaxText and Post training Dependencies\n",
103
103
"\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:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it."
Copy file name to clipboardExpand all lines: src/maxtext/examples/sft_qwen3_demo.ipynb
+24-19Lines changed: 24 additions & 19 deletions
Original file line number
Diff line number
Diff line change
@@ -95,40 +95,46 @@
95
95
" IN_COLAB = False"
96
96
]
97
97
},
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:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it."
107
+
]
108
+
},
98
109
{
99
110
"cell_type": "code",
100
111
"execution_count": null,
101
112
"metadata": {},
102
113
"outputs": [],
103
114
"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",
"## 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",
126
132
"**Session restart Instructions for Colab:**\n",
127
133
"1. Navigate to the menu at the top of the screen.\n",
128
134
"2. Click on **Runtime**.\n",
129
135
"3. Select **Restart session** from the dropdown menu.\n",
130
136
"\n",
131
-
"You will be asked to confirm the action in a pop-up dialog. Click on **Yes**.\n"
137
+
"You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
0 commit comments