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

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 XPK

Use Accelerated Processing Kit (XPK) to create GKE clusters for proof-of-concept and testing.

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

Before you begin

Before you start, complete the following steps:

Ensure you have a Google Cloud project with billing enabled.
Get access to TPU7x. For more information, contact your account team.
Ensure the account you're using with XPK has the roles listed in the XPK GitHub repository.

Install XPK and dependencies

Install XPK. Follow the instructions in the XPK GitHub repository.
Install Docker using instructions provided by your administrator or follow the official installation instructions. Once installed, run the following commands to configure Docker and test the installation:
```
gcloud auth configure-docker
sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
docker run hello-world # Test Docker
```
Set the following environment variables:
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
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_ACCELERATOR_TYPE: The TPU version and topology. For example, tpu7x-4x4x8. For a list of supported topologies, see Supported configurations.
- 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

Create a single-NIC, single slice cluster

Choose one of the following options to create your cluster. Using a custom network with 8,896 MTU is recommended for optimal performance.

Custom network

To create a custom network with 8,896 MTU and use it for your cluster, follow these steps:

Set environment variables for the network and firewall names:
```
export NETWORK_NAME=NETWORK_NAME
export NETWORK_FW_NAME=FIREWALL_NAME
```
Replace the following:
- NETWORK_NAME: A name for the network.
- FIREWALL_NAME: A name for the network firewall rule.

Create a custom network with an MTU of 8,896:

gcloud compute networks create ${NETWORK_NAME} \
    --mtu=8896 \
    --project=${PROJECT_ID} \
    --subnet-mode=auto \
    --bgp-routing-mode=regional

Create a firewall rule that allows TCP, ICMP, and UDP traffic on your network:

gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
    --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

Set an environment variable for the XPK cluster arguments to use the network you created:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
```

Create the XPK cluster. The following command provisions on-demand capacity:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

To use reserved capacity, replace --on-demand with --reservation=RESERVATION_NAME. To use TPU Spot VMs, replace