Skip to content

install_maxtext.md

Latest commit

 

History

History
174 lines (127 loc) · 5.28 KB

install_maxtext.md

File metadata and controls

174 lines (127 loc) · 5.28 KB
Raw

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.

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:

    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.

    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.

      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.

      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.

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

      uv pip install maxtext[runner]=={{version}} --resolution=lowest
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.

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:

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

  2. Create virtual environment:

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

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

    • Option 2: Install .[cuda12]

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

    • Option 3: Install .[tpu-post-train]

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

    • Option 4: Install .[runner]

      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:

    uv pip check
    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:

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