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.3 --resolution=lowest
install_tpu_pre_train_extra_deps
```
- Option 2: Install maxtext[cuda12], used for pre-training and decoding on GPUs.
```
uv pip install maxtext[cuda12]==0.2.3 --resolution=lowest
install_cuda12_pre_train_extra_deps
```
- 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.3 --resolution=lowest
install_tpu_post_train_extra_deps
```
- Option 4: Install maxtext[runner], used for building MaxText’s Docker images and scheduling workloads through XPK. Once installed, you will have access to the build_maxtext_docker_image, upload_maxtext_docker_image, and xpk commands. For more details on building and uploading Docker images, see the Build MaxText Docker Image guide.
```
uv pip install maxtext[runner]==0.2.3 --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

Option 2: Install .[cuda12]

uv pip install -e .[cuda12] --resolution=lowest
install_cuda12_pre_train_extra_deps

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