Train a model using TPU7x (Ironwood)

This document describes how to provision TPU7x resources and gives an example of deploying a training workload using MaxText and Cluster Toolkit.

TPU7x is the first release within the Ironwood family, Google Cloud's seventh generation TPU. The Ironwood generation is designed for large-scale AI training and inference. For more information, see TPU7x.

For more examples optimized for TPU7x, see Training Recipes for Ironwood TPU on GitHub.

Deploy a training workload with MaxText and Cluster Toolkit

Use Cluster Toolkit to create production-ready GKE clusters and run training workloads.

The following sections show how to deploy a training workload using MaxText and Cluster Toolkit.

Before you begin

Before you start, complete the following steps:

Ensure you have a Google Cloud project with billing enabled.
Install and initialize the Google Cloud CLI.
Install Cluster Toolkit.

Set the following environment variables:
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"
```
Replace the following:
- YOUR_PROJECT_ID: Your Google Cloud project ID.
- YOUR_ZONE: The zone in which to create the cluster.
- YOUR_CLUSTER_NAME: The name of the new cluster.
- YOUR_BUCKET_NAME: The name of your Cloud Storage bucket, which will be the output directory for model training.

If you don't have an existing Cloud Storage bucket, create one using the following command:

gcloud storage buckets create ${BASE_OUTPUT_DIR} \
    --project=${PROJECT_ID} \
    --location=US \
    --default-storage-class=STANDARD \
    --uniform-bucket-level-access

Deploy the TPU7x (Ironwood) cluster

Deploy a GKE TPU7x cluster by following the instructions in Deploy a GKE TPU 7x cluster. Ensure that your cluster node pool topology matches your training workload requirements (such as 4x4x8 for the sample MaxText workload below).

Build and upload the MaxText Docker image

Build and push the MaxText Docker image:

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Build the Docker image
bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.8.2

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

Define the MaxText training command

Prepare the command to run your training script within the Docker container.

The MaxText 1B model is a configuration within the MaxText framework designed for training a language model with approximately 1 billion parameters. Use this model to experiment with small chip scales. Performance is not optimized.

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_training \
        steps=30"

Deploy the training workload

Submit your training job using gcluster job submit. For details on job submission options, see the Cluster Toolkit Job Guide:

gcluster job submit \
    --name="maxtext-1b-$(date +%H%M)" \
    --cluster=${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --compute-type=tpu7x-standard-4t \
    --topology=4x4x8 \
    --image="gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}" \
    --command="${MAXTEXT_COMMAND}"

Workload names must be unique within the cluster. In this example, $(date +%H%M) is appended to the workload name to ensure uniqueness.