<!--
 Copyright 2023-2025 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]==0.2.1 --resolution=lowest
     ```

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

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

     ```bash
     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](https://maxtext.readthedocs.io/en/latest/build_maxtext.html)
     guide.

     ```bash
     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](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
   ```

:::\{only} is_not_latest

By default, cloning the repository provides the latest version (**HEAD**).
If you wish to use the latest features, please follow the [latest guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html).
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.

```{eval-rst}
.. parsed-literal::

  git checkout |version|
```

:::

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

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