<!--
 Copyright 2023-2026 Google LLC

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

      https://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 -->

# 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](https://docs.astral.sh/uv/getting-started/installation/).

```{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:**

   ```bash
   uv venv --python 3.12 --seed <VENV_NAME>
   source <VENV_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.

     ```bash
     uv pip install maxtext[tpu]=={{version}} --resolution=lowest
     install_tpu_pre_train_extra_deps
     ```

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

     ```bash
     uv pip install maxtext[cuda12]=={{version}} --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.

     ```bash
     UV_TORCH_BACKEND=cpu uv pip install maxtext[tpu-post-train]=={{version}} --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](build_maxtext.md)
     guide.

     ```bash
     uv pip install maxtext[runner]=={{version}} --resolution=lowest
     ```

```{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.
```

(install-from-source)=

## 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:

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

2. Create virtual environment:

   ```bash
   uv venv --python 3.12 --seed <VENV_NAME>
   source <VENV_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]`:

     ```bash
     uv pip install -e .[tpu] --resolution=lowest
     install_tpu_pre_train_extra_deps
     ```

   - **Option 2:** Install `.[cuda12]`

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

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

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

   - **Option 4:** Install `.[runner]`

     ```bash
     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:

1. **Verify the environment has no dependency conflicts:**

   ```bash
   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.
   ```

2. **Verify MaxText is importable and runnable:**

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