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.

1. Create a virtual environment:

   uv venv --python 3.12 --seed <virtual env name>
   source <virtual env name>/bin/activate
   

2. 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.1 --resolution=lowest
     

   • Option 2: Install maxtext[cuda12], used for pre-training and decoding on GPUs.

     uv pip install maxtext[cuda12]==0.2.1 --resolution=lowest
     

   • 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 pip install maxtext[tpu-post-train]==0.2.1 --resolution=lowest
     

   • 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.1 --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.

1. Clone the repository:

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

By default, cloning the repository provides the latest version (HEAD). If you wish to use the latest features, please follow the latest guide. If you want to ensure compatibility with the specific version of the documentation you are currently viewing, you must checkout the corresponding tag for that version before proceeding with the installation.

git checkout latest

2. Create virtual environment:

   uv venv --python 3.12 --seed <virtual env name>
   source <virtual env name>/bin/activate
   

3. 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 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
     

After installation, you can verify the package is available with python3 -c "import maxtext" and run training jobs with python3 -m maxtext.trainers.pre_train.train ....