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.

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

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

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.

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.

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