在 Docker 容器中运行 Cloud TPU 应用

Docker 容器 您可以结合自己的代码和所有需要的应用 一个可分发软件包中的依赖项。你可以通过容器 TPU 虚拟机，用于简化 Cloud TPU 应用的配置和共享。这个 文档介绍了如何为每个机器学习框架设置 Docker 容器 Cloud TPU 支持的存储类型。

在 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。我们建议您 请使用 PJRT，除非您有使用 XRT 的理由。如需详细了解 不同的运行时配置，看您是否有理由使用 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。一种方法是使用 gcloud compute tpus tpu-vm ssh 命令及 --worker=all 和 --command 标志。以下过程展示如何创建 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。一种方法是使用 gcloud compute tpus tpu-vm ssh 带有 --worker=all 和 --command 标志的命令。以下 过程展示了如何创建 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

后续步骤