Install MaxText

Install MaxText#

This document discusses how to install MaxText. We recommend installing MaxText inside a Python virtual environment.

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
uv pip install maxtext --resolution=lowest
install_maxtext_github_deps

Note: The install_maxtext_github_deps command is temporarily required to install dependencies directly from GitHub that are not yet available on PyPI.

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
# install the tpu package
uv pip install -e .[tpu] --resolution=lowest
# or install the gpu package by running the following line
# uv pip install -e .[cuda12] --resolution=lowest
install_maxtext_github_deps

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.

Overview of the Process#

To update dependencies, you will follow these general steps:

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).
Generate New Files: Run the seed-env CLI tool to generate new, fully-pinned requirements files based on your changes.
Update Project Files: Copy the newly generated files into the generated_requirements/ directory.
Handle GitHub Dependencies: Move any dependencies that are installed directly from GitHub from the generated files to src/install_maxtext_extra_deps/extra_deps_from_github.txt.
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=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=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.

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.
Update extra_deps_from_github.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.

Create a clean environment: It’s best to start with a fresh Python virtual environment.

uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

Run the setup script: Execute bash setup.sh to install the new dependencies.

pip install uv
# install the tpu package
uv pip install -e .[tpu] --resolution=lowest
# or install the gpu package by running the following line:
# uv pip install -e .[cuda12] --resolution=lowest
install_maxtext_github_deps

Run tests: Run MaxText tests to ensure there are no regressions.