Skip to content

Update docker image guide #3855

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
melissawm wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Update docker image guide #3855

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion docs/index.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,6 @@ maxdepth: 2
hidden:
---
install_maxtext
build_maxtext
tutorials
run_maxtext
guides
Expand Down
34 changes: 22 additions & 12 deletions docs/tutorials.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,44 +18,54 @@

Explore our tutorials to learn how to use MaxText, from your first run to advanced post-training techniques.

::::{grid} 1 2 2 2
:gutter: 2

:::{grid-item-card} 🚀 Getting Started
````{grid} 1 2 2 2
---
gutter: 2
---
```{grid-item-card} 🚀 Getting Started
:link: tutorials/first_run
:link-type: doc

Installation, prerequisites, verification, and your first training run.
:::
```

:::{grid-item-card} 📚 Pre-training
```{grid-item-card} 🛠️ Build and upload MaxText Docker Images
---
link: tutorials/build_maxtext
link-type: doc
---
Step-by-step guide for building MaxText docker images.
```

```{grid-item-card} 📚 Pre-training
:link: tutorials/pretraining
:link-type: doc

Step-by-step guides for pre-training with real datasets like C4 using HuggingFace, Grain, or TFDS.
:::
```

:::{grid-item-card} 🧩 Post-training
```{grid-item-card} 🧩 Post-training
:link: tutorials/post_training_index
:link-type: doc

Techniques for SFT, RL, and other post-training workflows on TPU.
:::
```

:::{grid-item-card} 📊 Inference
```{grid-item-card} 📊 Inference
:link: tutorials/inference
:link-type: doc

Step-by-step guides for running inference of MaxText models on vLLM.
:::
::::
```
````

```{toctree}
---
hidden:
maxdepth: 1
---
tutorials/first_run.md
tutorials/build_maxtext.md
tutorials/pretraining.md
tutorials/post_training_index.md
tutorials/inference.md
Expand Down
85 changes: 11 additions & 74 deletions docs/build_maxtext.md → docs/tutorials/build_maxtext.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
limitations under the License.
-->

# Build and Upload MaxText Docker Images
# Build and upload MaxText Docker images

This guide covers setting up a MaxText development environment and building container images for TPU and GPU workloads. These images can be used to run MaxText on GKE clusters with TPUs or GPUs, and are also required for running MaxText through XPK.

Expand Down Expand Up @@ -42,80 +42,15 @@ gcloud auth configure-docker
docker run hello-world
```

## Installation Modes
## Install MaxText

We recommend building MaxText inside a Python virtual environment using `uv` for speed and dependency management.
Follow the instructions in the [Installation guide](install_maxtext.md) to install the `maxtext[runner]` extra, which includes all necessary dependencies for building MaxText Docker images and running workloads through XPK. It automatically installs XPK, so you do not need to install it separately to manage your clusters and workloads.

### Option 1: From PyPI (Recommended)

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

```bash
# Install uv, a fast Python package installer
pip install uv
# Alternatively, if pip install fails:
# curl -LsSf https://astral.sh/uv/install.sh | sh

# Create virtual environment
export VENV_NAME=<VENV_NAME> # e.g., docker_venv
uv venv --python 3.12 --seed ${VENV_NAME?}
source ${VENV_NAME?}/bin/activate

# Install MaxText with the [runner] extra
# This enables Docker image building and workload scheduling via XPK.
# Once installed, you will have access to the `build_maxtext_docker_image`
# and `upload_maxtext_docker_image` commands.
uv pip install maxtext[runner]==0.2.1 --resolution=lowest
```

> **Note:** The `maxtext[runner]` extra includes all necessary dependencies for building MaxText Docker images and running workloads through XPK. It automatically installs XPK, so you do not need to install it separately to manage your clusters and workloads.

### Option 2: From Source

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

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

:::\{only} is_not_latest

By default, cloning the repository provides the latest version (**HEAD**).
If you wish to use the latest features, please follow the [latest guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html).
If you want to ensure compatibility with the specific version of the documentation
you are currently viewing, you must checkout the corresponding tag for that version
before proceeding with the installation.

```{eval-rst}
.. parsed-literal::

git checkout |version|
```

:::

```bash
# Create virtual environment
export VENV_NAME=<VENV_NAME> # e.g., docker_venv
uv venv --python 3.12 --seed ${VENV_NAME?}
source ${VENV_NAME?}/bin/activate

# Install MaxText with the [runner] extra in editable mode.
# This enables Docker image building and workload scheduling via XPK.
# Once installed, you will have access to the `build_maxtext_docker_image`
# and `upload_maxtext_docker_image` commands.
uv pip install -e .[runner] --resolution=lowest
```

> **Note:** The `maxtext[runner]` extra includes all necessary dependencies for building MaxText Docker images and running workloads through XPK. It automatically installs XPK, so you do not need to install it separately to manage your clusters and workloads.

## Build MaxText Docker Image
## Build MaxText Docker image

Select the appropriate build commands based on your hardware (`TPU` or `GPU`) and your specific workflow (`pre-training` or `post-training`). Each of these commands will generate a local Docker image named `maxtext_base_image`.

### TPU Pre-Training Docker Image
### TPU pre-training Docker image

```bash
# Option 1: Build with the stable versions of dependencies (default)
Expand All @@ -128,7 +63,7 @@ build_maxtext_docker_image MODE=nightly
build_maxtext_docker_image MODE=nightly JAX_VERSION=$JAX_VERSION
```

### GPU Pre-Training Docker Image
### GPU pre-training Docker image

```bash
# Option 1: Build with the stable versions of dependencies (default)
Expand All @@ -144,22 +79,24 @@ build_maxtext_docker_image DEVICE=gpu MODE=pinned
build_maxtext_docker_image DEVICE=gpu MODE=nightly JAX_VERSION=$JAX_VERSION
```

### TPU Post-Training Docker Image
### TPU post-training Docker image

```bash
# This build process takes approximately 10 to 15 minutes.
build_maxtext_docker_image WORKFLOW=post-training
```

## Upload MaxText Docker Image to Artifact Registry
## Upload MaxText Docker image to Artifact Registry

```bash
# Make sure to set `CLOUD_IMAGE_NAME` with your desired image name.
export CLOUD_IMAGE_NAME=<IMAGE_NAME>
upload_maxtext_docker_image CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME?}
```

> **Note:** You will need the [**Artifact Registry Writer**](https://docs.cloud.google.com/artifact-registry/docs/access-control#permissions) role to push Docker images to your project's Artifact Registry and to allow the cluster to pull them during workload execution. If you don't have this permission, contact your project administrator to grant you this role through "Google Cloud Console -> IAM -> Grant access".
```{note}
You will need the [**Artifact Registry Writer**](https://docs.cloud.google.com/artifact-registry/docs/access-control#permissions) role to push Docker images to your project's Artifact Registry and to allow the cluster to pull them during workload execution. If you don't have this permission, contact your project administrator to grant you this role through "Google Cloud Console -> IAM -> Grant access".
```

## Troubleshooting

Expand Down
Loading