Create a multi-host TPU slice

Create an instance template

To create an instance template for TPU v6e VMs, use the gcloud compute instance-templates create command:

gcloud compute instance-templates create quickstart-tpu-instance-template \
    --machine-type=ct6e-standard-4t \
    --maintenance-policy=TERMINATE \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --image-project=ubuntu-os-accelerator-images \
    --region=us-east5

Create a workload policy

A workload policy defines physical properties of your compute instances. In TPU slices, the accelerator topology defines the physical arrangement of TPU chips. Specifying an accelerator topology is required for interconnected multi-host TPU slices.

To create a workload policy for a multi-host TPU slice, use the gcloud compute resource-policies create workload-policy command with the --accelerator-topology flag. The following command creates a workload policy with a 2x4 topology:

gcloud compute resource-policies create workload-policy quickstart-tpu-workload-policy \
    --type=high-throughput \
    --accelerator-topology=2x4 \
    --region=us-east5

Create a MIG

Run the following commands to create a MIG that forms a multi-host TPU slice.

1. To create a MIG that forms a multi-host TPU slice, use the gcloud compute instance-groups managed create command:

   gcloud compute instance-groups managed create quickstart-tpu-mig \
       --size=2 \
       --target-size-policy-mode=bulk \
       --template=quickstart-tpu-instance-template \
       --region=us-east5 \
       --target-distribution-shape=any-single-zone \
       --instance-redistribution-type=none \
       --default-action-on-vm-failure=do-nothing \
       --workload-policy=projects/PROJECT_ID/regions/us-east5/resourcePolicies/quickstart-tpu-workload-policy
   

   Replace PROJECT_ID with your Cloud de Confiance by S3NS project ID.

2. Optionally, verify that the managed instances are running by using the following commands:

   • To view the overall MIG status, use the gcloud compute instance-groups managed describe command:

     gcloud compute instance-groups managed describe quickstart-tpu-mig \
         --region=us-east5
     

   • To view the status of each instance, use the gcloud compute instance-groups managed list-instances command:

     gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
         --region=us-east5
     

Install JAX

Install dependencies and the JAX framework in a virtual environment on all TPU VM instances in the MIG. If your TPU VMs have a Python version earlier than 3.11 installed, you must install Python 3.11 to run the latest version of JAX.

1. Check which Python version is running on your TPU VMs:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='python3 --version'
   

   If the version is earlier than Python 3.11, install Python 3.11:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt update && \
       sudo apt install -y software-properties-common && \
       sudo add-apt-repository -y ppa:deadsnakes/ppa && \
       sudo apt update && \
       sudo apt install -y python3.11 python3.11-dev'
   

2. Create a virtual environment:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt install -y python3.11-venv && \
       python3.11 -m venv ~/jax_venv'
   

3. Install JAX in the virtual environment:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='source ~/jax_venv/bin/activate && \
       pip install --upgrade pip -q && \
       pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q'
   

Run JAX code on the slice

To run JAX code on a TPU slice, you must run the code on each host in the TPU slice. The jax.device_count() function call stops responding until it's called on each host in the slice. The following example shows how to run a JAX calculation on a TPU slice.

Prepare the code

Create a file called example.py on each instance:

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command="cat << 'EOF' > ~/example.py
import jax

# Initialize the slice
jax.distributed.initialize()

# The total number of TPU cores in the slice
device_count = jax.device_count()

# The number of TPU cores attached to this host
local_device_count = jax.local_device_count()

# The psum is performed over all mapped devices across the slice
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)

# Print from a single host to avoid duplicated output
if jax.process_index() == 0:
    print('global device count:', jax.device_count())
    print('local device count:', jax.local_device_count())
    print('pmap result:', r)
EOF"

Run the code on the slice

Run the example.py program on each TPU VM in the slice:

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command='source ~/jax_venv/bin/activate && python3 ~/example.py'

The output should be similar to the following:

global device count: 8
local device count: 4
pmap result: [8. 8. 8. 8.]