Install MaxText

Install MaxText#

This document discusses how to install MaxText.

We recommend installing MaxText inside a Python virtual environment and using the uv package manager following uv’s official installation instructions.

Note

MaxText is only tested on Linux during releases.

From PyPI (recommended)#

This is the easiest way to get started with the latest stable version.

Create a virtual environment:

uv venv --python 3.12 --seed <VENV_NAME>
source <VENV_NAME>/bin/activate

Install MaxText and its dependencies.

Choose a single installation option from this list to fit your use case.

Important

If you want to switch to a different installation option (e.g., from [tpu] to [tpu-post-train]), we strongly recommend starting with a fresh virtual environment to avoid dependency conflicts.
- Option 1: Install maxtext[tpu], used for pre-training and decoding on TPUs.
```
uv pip install maxtext[tpu]==0.2.5 --resolution=lowest
install_tpu_pre_train_extra_deps
# Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
# install_tpu_pre_train_extra_deps --with-tf
```
- Option 2: Install maxtext[cuda12], used for pre-training and decoding on GPUs.
```
uv pip install maxtext[cuda12]==0.2.5 --resolution=lowest
install_cuda12_pre_train_extra_deps
# Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
# install_cuda12_pre_train_extra_deps --with-tf
```
- Option 3: Install maxtext[tpu-post-train], used for post-training on TPUs. Currently, this option should also be used for running vllm_decode on TPUs.
```
UV_TORCH_BACKEND=cpu uv pip install maxtext[tpu-post-train]==0.2.5 --resolution=lowest
install_tpu_post_train_extra_deps
```
- Option 4: Install maxtext[runner], used for building and uploading MaxText’s Docker images. Once installed, you will have access to the build_maxtext_docker_image and upload_maxtext_docker_image commands. For GKE job submission, install the gcluster CLI separately as described in the Cluster Toolkit guide. For more details on building and uploading Docker images, see the Build MaxText Docker Image guide.
```
uv pip install maxtext[runner]==0.2.5 --resolution=lowest
```

Note

The maxtext package contains a comprehensive list of all direct and transitive dependencies, with lower bounds, generated by 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.

From source#

If you plan to contribute to MaxText or need the latest unreleased features, install from source.

Important

If you want to switch to a different installation option (e.g., from [tpu] to [tpu-post-train]), we strongly recommend starting with a fresh virtual environment to avoid dependency conflicts.

Clone the repository:

git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext

Create virtual environment:

uv venv --python 3.12 --seed <VENV_NAME>
source <VENV_NAME>/bin/activate

Install dependencies in editable mode. Choose a single installation option from this list to fit your use case.

Option 1: Install .[tpu]:

uv pip install -e .[tpu] --resolution=lowest
install_tpu_pre_train_extra_deps
# Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
# install_tpu_pre_train_extra_deps --with-tf

Option 2: Install .[cuda12]

uv pip install -e .[cuda12] --resolution=lowest
install_cuda12_pre_train_extra_deps
# Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
# install_cuda12_pre_train_extra_deps --with-tf

Option 3: Install .[tpu-post-train]

UV_TORCH_BACKEND=cpu uv pip install -e .[tpu-post-train] --resolution=lowest
install_tpu_post_train_extra_deps

Option 4: Install .[runner]

uv pip install -e .[runner] --resolution=lowest

Verify installation#

After installing MaxText (either from PyPI or from source), verify that your virtual environment is healthy and free of package version conflicts:

Verify the environment has no dependency conflicts:
```
uv pip check
```
Note

If uv pip check reports any package version conflicts, they can usually be resolved by starting with a fresh virtual environment (see above) and reinstalling using the platform-specific target and the --resolution=lowest flag to ensure all dependencies resolve to their verified versions.

Verify MaxText is importable and runnable:

python3 -c "import maxtext"
python3 -m maxtext.trainers.pre_train.train --help