Run TPU workloads in a Docker container

Docker containers make configuring applications easier by combining your code and all needed dependencies in one distributable package. You can run Docker containers within TPU VMs to simplify configuring and sharing your Cloud TPU applications. This document describes how to set up a Docker container for each ML framework supported by Cloud TPU.

Train a PyTorch model in a Docker container

Before running the following commands, create the following environment variables:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Environment variable descriptions

  • PROJECT_ID: Your Google Cloud project ID. Use an existing project or create a new one.
  • TPU_NAME: The name of the TPU.
  • ZONE: The zone in which to create the TPU VM. For more information about supported zones, see TPU regions and zones.
  • ACCELERATOR_TYPE: The accelerator type specifies the version and size of the Cloud TPU you want to create. For more information about supported accelerator types for each TPU version, see TPU versions.
  • RUNTIME_VERSION: The Cloud TPU software version.

TPU device

  1. Create Cloud TPU VM

    gcloud compute tpus tpu-vm create $TPU_NAME \
--zone=$ZONE \
--accelerator-type=$ACCELERATOR_TYPE \
--version=$RUNTIME_VERSION

  2. Connect to the TPU VM using SSH

    gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

  3. Ensure your Google Cloud user has been granted the Artifact Registry Reader role. For more information, see Granting Artifact Registry roles.

  4. Start a container in the TPU VM using the nightly PyTorch/XLA image

    sudo docker run --net=host -ti --rm --name your-container-name --privileged \
us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 \
bash

  5. Configure TPU runtime

    There are two PyTorch/XLA runtime options: PJRT and XRT. We recommend you use PJRT unless you have a reason to use XRT. To learn more about the different runtime configurations, see the PJRT runtime documentation.

    PJRT

    export PJRT_DEVICE=TPU

    XRT

    export XRT_TPU_CONFIG="localservice;0;localhost:51011"

  6. Clone the PyTorch XLA repository

    git clone --recursive https://github.com/pytorch/xla.git

  7. Train ResNet50

    python3 xla/test/test_train_mp_imagenet.py \
--fake_data \
--model=resnet50 \
--num_epochs=1

When the training script completes, clean up the resources.

  1. Type exit to exit from the Docker container
  2. Type exit to exit from the TPU VM

  3. Delete the TPU VM

    gcloud compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE

TPU slice

When you run PyTorch code on a TPU slice, you must run your code on all TPU workers at the same time. One way to do this is to use the gcloud compute tpus tpu-vm ssh command with the --worker=all and --command flags. The following procedure shows you how to create a Docker image to make setting up each TPU worker easier.

  1. Create a TPU VM

    gcloud compute tpus tpu-vm create $TPU_NAME \
--zone=$ZONE \
--accelerator-type=$ACCELERATOR_TYPE \
--version=$RUNTIME_VERSION

  2. Add the current user to the Docker group

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
--zone=$ZONE \
--worker=all \
--command='sudo usermod -a -G docker $USER'

  3. Clone the PyTorch XLA repository

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
--zone=$ZONE \
--command="git clone --recursive https://github.com/pytorch/xla.git"

  4. Run the training script in a container on all TPU workers

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
--zone=$ZONE \
--command="docker run --rm --privileged --net=host  -v ~/xla:/xla -e PJRT_DEVICE=TPU us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 python /xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"

    Docker command flags:

    • --rm removes the container after its process terminates.
    • --privileged exposes the TPU device to the container.
    • --net=host binds all of the container's ports to the TPU VM to allow communication between the hosts in the pod.
    • -e sets environment variables.

When the training script completes, clean up the resources.

Delete the TPU VM using the following command:

gcloud compute tpus tpu-vm delete $TPU_NAME \
--zone=$ZONE

Train a JAX model in a Docker container

Before running the following commands, create the following environment variables:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Environment variable descriptions

  • PROJECT_ID: Your Google Cloud project ID. Use an existing project or create a new one.
  • TPU_NAME: The name of the TPU.
  • ZONE: The zone in which to create the TPU VM. For more information about supported zones, see TPU regions and zones.
  • ACCELERATOR_TYPE: The accelerator type specifies the version and size of the Cloud TPU you want to create. For more information about supported accelerator types for each TPU version, see TPU versions.
  • RUNTIME_VERSION: The Cloud TPU software version.

TPU device

  1. Create the TPU VM

    gcloud compute tpus tpu-vm create $TPU_NAME \
--zone=$ZONE \
--accelerator-type=$ACCELERATOR_TYPE \
--version=$RUNTIME_VERSION

  2. Connect to the TPU VM using SSH

    gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

  3. Start Docker daemon in TPU VM

    sudo systemctl start docker

  4. Start Docker container

    sudo docker run --net=host -ti --rm --name your-container-name \
--privileged us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 \
bash

  5. Install JAX

    pip install jax[tpu]

  6. Install FLAX

    pip install --upgrade clu
git clone https://github.com/google/flax.git
pip install --user -e flax

  7. Install tensorflow and tensorflow-datasets packages

    pip install tensorflow
pip install tensorflow-datasets

  8. Run the FLAX MNIST training script

    cd flax/examples/mnist
python3 main.py --workdir=/tmp/mnist \
--config=configs/default.py \
--config.learning_rate=0.05 \
--config.num_epochs=5

When the training script completes, clean up the resources.

  1. Type exit to exit from the Docker container
  2. Type exit to exit from the TPU VM

  3. Delete the TPU VM

    gcloud compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE

TPU slice

When you run JAX code on a TPU slice, you must run your JAX code on all TPU workers at the same time. One way to do this is to use the gcloud compute tpus tpu-vm ssh command with the --worker=all and --command flags. The following procedure shows you how to create a Docker image to make setting up each TPU worker easier.

  1. Create a file named Dockerfile in your current directory and paste the following text

    FROM python:3.10
RUN pip install jax[tpu]
RUN pip install --upgrade clu
RUN git clone https://github.com/google/flax.git
RUN pip install --user -e flax
RUN pip install tensorflow
RUN pip install tensorflow-datasets
WORKDIR ./flax/examples/mnist

  2. Prepare an Artifact Registry

    gcloud artifacts repositories create your-repo \
--repository-format=docker \
--location=europe-west4 --description="Docker repository" \
--project=$PROJECT_ID

gcloud artifacts repositories list \
--project=$PROJECT_ID

gcloud auth configure-docker europe-west4-docker.pkg.dev

  3. Build the Docker image

    docker build -t your-image-name .

  4. Add a tag to your Docker image before pushing it to the Artifact Registry. For more information on working with Artifact Registry, see Work with container images.

    docker tag your-image-name europe-west4-docker.pkg.dev/$PROJECT_ID/your-repo/your-image-name:your-tag

  5. Push your Docker image to the Artifact Registry

    docker push europe-west4-docker.pkg.dev/$PROJECT_ID/your-repo/your-image-name:your-tag

  6. Create a TPU VM

    gcloud compute tpus tpu-vm create $TPU_NAME \
--zone=$ZONE \
--accelerator-type=$ACCELERATOR_TYPE \
--version=$RUNTIME_VERSION

  7. Pull the Docker image from the Artifact Registry on all TPU workers

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
--zone=$ZONE \
--command='sudo usermod -a -G docker ${USER}'

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
--zone=$ZONE \
--command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
--zone=$ZONE \
--command="docker pull europe-west4-docker.pkg.dev/$PROJECT_ID/your-repo/your-image-name:your-tag"

  8. Run the container on all TPU workers

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
--zone=$ZONE \
--command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/$PROJECT_ID/your-repo/your-image-name:your-tag bash"

  9. Run the training script on all TPU workers

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
--zone=$ZONE \
--command="docker exec --privileged your-container-name python3 main.py --workdir=/tmp/mnist \
--config=configs/default.py \
--config.learning_rate=0.05 \
--config.num_epochs=5"

When the training script completes, clean up the resources.

  1. Shut down the container on all workers

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
--zone=$ZONE \
--command="docker kill your-container-name"

  2. Delete the TPU VM

    gcloud compute tpus tpu-vm delete $TPU_NAME \
--zone=$ZONE

What's next