在 Docker 容器中运行 Cloud TPU 应用

Docker 容器可将您的代码和所有所需依赖项组合到一个可分发软件包中，从而简化应用配置。您可以在 TPU 虚拟机中运行 Docker 容器，以简化 Cloud TPU 应用的配置和共享。本文档介绍了如何为 Cloud TPU 支持的每个机器学习框架设置 Docker 容器。

重要提示 ：您最多可以在一个 TPU 虚拟机上运行 256 个 Docker Pod。

在 Docker 容器中训练 TensorFlow 模型

TPU 设备

1. 在当前目录中创建一个名为 Dockerfile 的文件，并将以下文本粘贴到其中

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
   RUN git clone https://github.com/tensorflow/models.git
   WORKDIR ./models
   RUN pip install -r official/requirements.txt
   ENV PYTHONPATH=/models

2. 创建 Cloud Storage 存储桶

   gcloud storage buckets create gs://your-bucket-name --location=europe-west4

3. 创建 TPU 虚拟机

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-vm-tf-2.17.0-pjrt

4. 将 Dockerfile 复制到您的 TPU 虚拟机

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:

5. 通过 SSH 连接到 TPU 虚拟机

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a

6. 构建 Docker 映像

   sudo docker build -t your-image-name . 

7. 启动 Docker 容器

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash

8. 设置环境变量

   export STORAGE_BUCKET=gs://your-bucket-name
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x

9. 训练 ResNet

   python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,trainer.train_steps=100"

训练脚本运行完毕后，请务必清理资源。

1. 输入 exit 以退出 Docker 容器
2. 输入 exit 以退出 TPU 虚拟机
3. 删除 TPU 虚拟机

    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

TPU Pod

1. 在当前目录中创建一个名为 Dockerfile 的文件，并将以下文本粘贴到其中

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
   RUN git clone https://github.com/tensorflow/models.git
   WORKDIR ./models
   RUN pip install -r official/requirements.txt
   ENV PYTHONPATH=/models

2. 创建 TPU 虚拟机

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v3-32 \
   --version=tpu-vm-tf-2.17.0-pod-pjrt

3. 将 Dockerfile 复制到您的 TPU 虚拟机

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:

4. 通过 SSH 连接到 TPU 虚拟机

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a

5. 构建 Docker 映像

   sudo docker build -t your-image-name . 

6. 启动 Docker 容器

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash

7. 训练 ResNet

   python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

训练脚本运行完毕后，请务必清理资源。

1. 输入 exit 以退出 Docker 容器
2. 输入 exit 以退出 TPU 虚拟机
3. 删除 TPU 虚拟机

     $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

在 Docker 容器中训练 PyTorch 模型

TPU 设备

1. 创建 Cloud TPU 虚拟机

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

2. 通过 SSH 连接到 TPU 虚拟机

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a

3. 使用每夜 PyTorch/XLA 映像在 TPU 虚拟机中启动容器。

   sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm bash

   注意：运行此命令后，您的命令提示符将发生变化，表示终端已连接到正在运行的容器。

4. 配置 TPU 运行时

   有两种 PyTorch/XLA 运行时选项：PJRT 和 XRT。除非您有理由使用 XRT，否则我们建议您使用 PJRT。如需详细了解 不同的运行时配置，看您是否有理由使用 XRT。如需详细了解不同的运行时配置，请参阅 PJRT 运行时文档。

   PJRT

   export PJRT_DEVICE=TPU

   XRT

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

5. 克隆 PyTorch XLA 代码库

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

6. 训练 ResNet50

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

训练脚本运行完毕后，请务必清理资源。

1. 输入 exit 以退出 Docker 容器
2. 输入 exit 以退出 TPU 虚拟机
3. 删除 TPU 虚拟机

    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

TPU Pod

在 TPU Pod 上运行 PyTorch 代码时，您必须同时在所有 TPU 工作器上运行代码。为此，一种方法是使用带有 --worker=all 和 --command 标志的 gcloud compute tpus tpu-vm ssh 命令。以下过程展示如何创建 Docker 以便更轻松地设置每个 TPU 工作器。

1. 创建 TPU 虚拟机

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=us-central2-b \
   --accelerator-type=v4-32 \
   --version=tpu-ubuntu2204-base

2. 将当前用户添加到 docker 组

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=us-central2-b \
   --worker=all \
   --command="sudo usermod -a -G docker $USER"

3. 在所有 TPU 工作器上的容器中运行训练脚本。

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="docker run --rm --privileged --net=host  -e PJRT_DEVICE=TPU gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm python /pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"

   Docker 命令标志：

   • --rm 会在其进程终止后移除容器。
   • --privileged 会向容器公开 TPU 设备。
   • --net=host 会将容器的所有端口绑定到 TPU 虚拟机，以允许 Pod 中的主机之间进行通信。
   • -e 设置环境变量。

训练脚本运行完毕后，请务必清理资源。

使用以下命令删除 TPU 虚拟机：

$ gcloud compute tpus tpu-vm delete your-tpu-name \
  --zone=us-central2-b

在 Docker 容器中训练 JAX 模型

TPU 设备

1. 创建 TPU 虚拟机

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

2. 通过 SSH 连接到 TPU 虚拟机

   gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a

3. 在 TPU 虚拟机中启动 Docker 守护程序

   sudo systemctl start docker

4. 启动 Docker 容器

   sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash

5. 安装 JAX

   pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

6. 安装 FLAX

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

7. 运行 FLAX MNIST 训练脚本

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

训练脚本运行完毕后，请务必清理资源。

1. 输入 exit 以退出 Docker 容器
2. 输入 exit 以退出 TPU 虚拟机

3. 删除 TPU 虚拟机

   $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

TPU Pod

在 TPU Pod 上运行 JAX 代码时，您必须同时在所有 TPU 工作器上运行 JAX 代码。为此，一种方法是使用带有 --worker=all 和 --command 标志的 gcloud compute tpus tpu-vm ssh 命令。以下过程介绍了如何创建 Docker 映像，以便更轻松地设置每个 TPU 工作器。

1. 在当前目录中创建一个名为 Dockerfile 的文件，然后将 关注文本

   FROM python:3.8
   RUN pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   RUN pip install --upgrade clu
   RUN git clone https://github.com/google/flax.git
   RUN pip install --user -e flax
   WORKDIR ./flax/examples/mnist

2. 构建 Docker 映像

   docker build -t your-image-name .

3. 请先为 Docker 映像添加标记，然后再将该映像推送到 Artifact Registry。 如需详细了解如何使用 Artifact Registry，请参阅 使用容器映像。

   docker tag your-image-name europe-west-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

4. 将 Docker 映像推送到 Artifact Registry

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

5. 创建 TPU 虚拟机

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type==v2-8 \
   --version=tpu-ubuntu2204-base

6. 在所有 TPU 工作器上从 Artifact Registry 拉取 Docker 映像。

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="sudo usermod -a -G docker ${USER}"

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"

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

7. 在所有 TPU 工作器上运行容器。

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

8. 在所有 TPU 工作器上运行训练脚本：

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --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"

训练脚本运行完毕后，请务必清理资源。

1. 关停所有工作器上的容器：

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker kill your-container-name"

2. 使用以下命令删除 TPU 虚拟机：

   $ gcloud compute tpus tpu-vm delete your-tpu-name \
   --zone=europe-west4-a

后续步骤

• Cloud TPU 教程
• 管理 TPU
• Cloud TPU 系统架构
• 在 TPU Pod 切片上运行 TensorFlow 代码
• 在 TPU Pod 切片上运行 JAX 代码