Skip to content

install_maxtext.md

Latest commit

 

History

History
214 lines (149 loc) · 10.3 KB

install_maxtext.md

File metadata and controls

214 lines (149 loc) · 10.3 KB

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.
# 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: Installing maxtext[tpu]
uv pip install maxtext[tpu]==0.2.1 --resolution=lowest
install_tpu_pre_train_extra_deps

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

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

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

Note: The maxtext[runner] extra is used for building MaxText 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.

Note: The install_tpu_pre_train_extra_deps, install_cuda12_pre_train_extra_deps, and install_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.

Note: MaxText is only tested on Linux during releases.

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.
# 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: Installing .[tpu]
uv pip install -e .[tpu] --resolution=lowest
install_tpu_pre_train_extra_deps

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

# Option 3: Installing .[tpu-post-train]
uv pip install -e .[tpu-post-train] --resolution=lowest
install_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 ....

UV Project management

For simplicity this guide uses the traditional uv pip install syntax. If you are using the uv project management features (with a pyproject.toml and uv.lock in your own project), you would need to run your commands differently. You would need to use uv add instead of uv pip install and uv run as a prefix to all other commands. For example, to install and run MaxText in a uv-managed project:

# 1. Initialize your uv project
mkdir my-maxtext-project && cd my-maxtext-project
uv init

# 2. Add MaxText as a dependency
uv add maxtext[tpu]==0.2.1 --resolution=lowest

# 3. Install MaxText's extra GitHub dependencies. These will be automatically added to your pyproject.toml
uv run install_tpu_pre_train_extra_deps # This will be added to your pyproject.toml

# 4. Run MaxText training for a few steps
uv run python3 -m maxtext.trainers.pre_train.train run_name=${RUN_NAME?} base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
    dataset_type=synthetic steps=5 per_device_batch_size=1 model_name=llama2-7b

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 src/dependencies/requirements/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 src/dependencies/requirements/generated_requirements/tpu-requirements.txt.
    • Move generated_gpu_artifacts/cuda12-requirements.txt to src/dependencies/requirements/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.