<!--
 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.
MaxText offers following installation modes:

1. maxtext[tpu]. Used for pre-training and decode on TPUs.
2. maxtext[cuda12]. Used for pre-training and decode on GPUs.
3. maxtext[tpu-post-train]. Used for post-training on TPUs. Currently, this option should also be used for running vllm_decode on TPUs.
4. maxtext[runner]. Used for building MaxText's Docker images and scheduling workloads through XPK.

## From PyPI (Recommended)

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

```bash
# 1. Install uv, a fast Python package installer
pip install uv

# 2. Create virtual environment
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

# 3. Install MaxText and its dependencies. Choose a single
# installation option from this list to fit your use case.

# Option 1: Installing maxtext[tpu]
uv pip install maxtext[tpu] --resolution=lowest
install_maxtext_tpu_github_deps

# Option 2: Installing maxtext[cuda12]
uv pip install maxtext[cuda12] --resolution=lowest
install_maxtext_cuda12_github_dep

# Option 3: Installing maxtext[tpu-post-train]
uv pip install maxtext[tpu-post-train] --resolution=lowest
install_maxtext_tpu_post_train_extra_deps

# Option 4: Installing maxtext[runner]
uv pip install maxtext[runner] --resolution=lowest
```

> **Note:** The `install_maxtext_tpu_github_deps`, `install_maxtext_cuda12_github_dep`, and
> `install_maxtext_tpu_post_train_extra_deps` commands are temporarily required to install dependencies directly from GitHub
> that are not yet available on PyPI. As shown above, choose the one that corresponds to your use case.

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

## From Source

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

```bash
# 1. Clone the repository
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext

# 2. Create virtual environment
pip install uv
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

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

# Option 1: Installing .[tpu]
uv pip install -e .[tpu] --resolution=lowest
install_maxtext_tpu_github_deps

# Option 2: Installing .[cuda12]
uv pip install -e .[cuda12] --resolution=lowest
install_maxtext_cuda12_github_dep

# Option 3: Installing .[tpu-post-train]
uv pip install -e .[tpu-post-train] --resolution=lowest
install_maxtext_tpu_post_train_extra_deps

# Option 4: Installing maxtext[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 ...`.

# Update MaxText dependencies

## Introduction

This document provides a guide to updating dependencies in MaxText using the `seed-env` tool. `seed-env` helps generate deterministic and reproducible Python environments by creating fully-pinned `requirements.txt` files from a base set of requirements.

Please keep dependencies updated throughout development. This will allow each commit to work properly from both a feature and dependency perspective. We will periodically upload commits to PyPI for stable releases. But it is also critical to keep dependencies in sync for users installing MaxText from source.

## Overview of the Process

To update dependencies, you will follow these general steps:

1. **Modify Base Requirements**: Update the desired dependencies in `base_requirements/requirements.txt` or the hardware-specific files (`base_requirements/tpu-base-requirements.txt`, `base_requirements/gpu-base-requirements.txt`).
2. **Generate New Files**: Run the `seed-env` CLI tool to generate new, fully-pinned requirements files based on your changes.
3. **Update Project Files**: Copy the newly generated files into the `generated_requirements/` directory.
4. **Handle GitHub Dependencies**: Move any dependencies that are installed directly from GitHub from the generated files to `src/dependencies/github_deps/pre_train_deps.txt`.
5. **Verify**: Test the new dependencies to ensure the project installs and runs correctly.

The following sections provide detailed instructions for each step.

## Step 1: Install seed-env

First, you need to install the `seed-env` command-line tool by running `pip install seed-env uv`. Or follow the instructions in the
[seed-env repository](https://github.com/google-ml-infra/actions/tree/main/python_seed_env#install-the-seed-env-tool) if you want to build `seed-env` from source.

## Step 2: Find the JAX Build Commit Hash

The dependency generation process is pinned to a specific nightly build of JAX. You need to find the commit hash for the desired JAX build.

You can find the latest commit hashes in the [JAX `build/` folder](https://github.com/jax-ml/jax/commits/main/build). Choose a recent, successful build and copy its full commit hash.

## Step 3: Generate the Requirements Files

Next, run the `seed-env` CLI to generate the new requirements files. You will need to do this separately for the TPU and GPU environments. The generated files will be placed in a directory specified by `--output-dir`.

### For TPU

Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step.

```bash
seed-env \
  --local-requirements=src/dependencies/requirements/base_requirements/tpu-base-requirements.txt \
  --host-name=MaxText \
  --seed-commit=<jax-build-commit-hash> \
  --python-version=3.12 \
  --requirements-txt=tpu-requirements.txt \
  --output-dir=generated_tpu_artifacts
```

### For GPU

Similarly, run the command for the GPU requirements.

```bash
seed-env \
  --local-requirements=src/dependencies/requirements/base_requirements/cuda12-base-requirements.txt \
  --host-name=MaxText \
  --seed-commit=<jax-build-commit-hash> \
  --python-version=3.12 \
  --requirements-txt=cuda12-requirements.txt \
  --hardware=cuda12 \
  --output-dir=generated_gpu_artifacts
```

## Step 4: Update Project Files

After generating the new requirements, you need to update the files in the MaxText repository.

1. **Copy the generated files:**

   - Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`.
   - Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`.

2. **Update `pre_train_deps.txt` (if necessary):**
   Currently, MaxText uses a few dependencies, such as `mlperf-logging` and `google-jetstream`, that are installed directly from GitHub source. These are defined in `base_requirements/requirements.txt`, and the `seed-env` tool will carry them over to the generated requirements files.

## Step 5: Verify the New Dependencies

Finally, test that the new dependencies install correctly and that MaxText runs as expected.

1. **Install MaxText and dependencies**: For instructions on installing MaxText on your VM, please refer to the [official documentation](https://maxtext.readthedocs.io/en/maxtext-v0.2.0/install_maxtext.html#from-source).

2. **Verify the installation**: Run MaxText tests to ensure everything is working as expected with the newly installed dependencies and there are no regressions.