在 Docker 容器中執行 TPU 工作負載

注意：本頁內容適用於 Cloud TPU API，該 API 不再處於積極開發階段。Cloud TPU API 只會收到錯誤修正和安全性更新。如要使用最新功能並支援最新 TPU 版本，建議您使用 Compute Engine 或 Google Kubernetes Engine 管理 TPU 資源。詳情請參閱「Compute Engine 中的 TPU 資源」和「GKE 中的 TPU 簡介」。TPU7x 和更新的 TPU 版本不支援 Cloud TPU API。

Docker 容器可將程式碼和所有必要依附元件合併在一個可發布的套件中，讓應用程式設定更輕鬆。您可以在 TPU VM 中執行 Docker 容器，簡化 Cloud TPU 應用程式的設定和共用程序。本文說明如何為 Cloud TPU 支援的每個機器學習框架設定 Docker 容器。

重要事項： 您最多可以在 TPU VM 上執行 256 個 Docker Pod。

在 Docker 容器中訓練 PyTorch 模型

執行下列指令前，請先建立下列環境變數：

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

環境變數說明

• PROJECT_ID： 您的 Google Cloud 專案 ID。使用現有專案或建立新專案。
• TPU_NAME：TPU 的名稱。
• ZONE： 要建立 TPU VM 的可用區。如要進一步瞭解支援的區域，請參閱 TPU 地區和區域。
• ACCELERATOR_TYPE： 加速器類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各個 TPU 版本支援的加速器類型，請參閱 TPU 版本。
• RUNTIME_VERSION：Cloud TPU 軟體版本。

TPU 裝置

1. 建立 Cloud TPU VM

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

2. 使用 SSH 連線至 TPU VM

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

3. 確認 Google Cloud 使用者已獲派 Artifact Registry Reader 角色。詳情請參閱「授予 Artifact Registry 角色」。

4. 使用每夜建構的 PyTorch/XLA 映像檔，在 TPU VM 中啟動容器

   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. 設定 TPU 執行階段

   PyTorch/XLA 執行階段有兩種選項：PJRT 和 XRT。除非您有使用 XRT 的理由，否則建議使用 PJRT。如要進一步瞭解不同的執行階段設定，請參閱 PJRT 執行階段說明文件。

   PJRT

   export PJRT_DEVICE=TPU
   

   XRT

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

6. 複製 PyTorch XLA 存放區

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

7. 訓練 ResNet50

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

訓練指令碼執行完畢後，請清理資源。

1. 輸入 exit 即可退出 Docker 容器
2. 輸入 exit 即可退出 TPU VM

3. 刪除 TPU VM

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

TPU 配量

在 TPU 配量上執行 PyTorch 程式碼時，必須同時在所有 TPU 工作站上執行程式碼。其中一種方式是使用 gcloud compute tpus tpu-vm ssh 指令，並加上 --worker=all 和 --command 旗標。以下程序說明如何建立 Docker 映像檔，方便設定每個 TPU 工作站。

1. 建立 TPU VM

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

2. 將目前使用者新增至 Docker 群組

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

3. 複製 PyTorch XLA 存放區

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

4. 在所有 TPU 工作站的容器中執行訓練指令碼

   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 指令旗標：

   • --rm 會在容器程序終止後移除容器。
   • --privileged 會將 TPU 裝置公開給容器。
   • --net=host 會將容器的所有通訊埠繫結至 TPU VM，允許 Pod 中的主機彼此通訊。
   • -e 會設定環境變數。

訓練指令碼執行完畢後，請清理資源。

使用下列指令刪除 TPU VM：

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

在 Docker 容器中訓練 JAX 模型

執行下列指令前，請先建立下列環境變數：

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

環境變數說明

• PROJECT_ID： 您的 Google Cloud 專案 ID。使用現有專案或建立新專案。
• TPU_NAME：TPU 的名稱。
• ZONE： 要建立 TPU VM 的可用區。如要進一步瞭解支援的區域，請參閱 TPU 地區和區域。
• ACCELERATOR_TYPE： 加速器類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各個 TPU 版本支援的加速器類型，請參閱 TPU 版本。
• RUNTIME_VERSION：Cloud TPU 軟體版本。

TPU 裝置

1. 建立 TPU VM

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

2. 使用 SSH 連線至 TPU VM

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

3. 在 TPU VM 中啟動 Docker Daemon

   sudo systemctl start docker
   

4. 啟動 Docker 容器

   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. 安裝 JAX

   pip install jax[tpu]
   

6. 安裝 FLAX

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

7. 安裝 tensorflow 和 tensorflow-datasets 套件

   pip install tensorflow
   pip install tensorflow-datasets
   

8. 執行 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 VM

3. 刪除 TPU VM

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

TPU 配量

在 TPU 配量上執行 JAX 程式碼時，您必須同時在所有 TPU 工作站上執行 JAX 程式碼。其中一種方式是使用 gcloud compute tpus tpu-vm ssh 指令，並加上 --worker=all 和 --command 旗標。以下程序說明如何建立 Docker 映像檔，方便設定各個 TPU 工作站。

1. 在目前目錄中建立名為 Dockerfile 的檔案，然後貼上以下文字

   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. 準備 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. 建構 Docker 映像檔

   docker build -t your-image-name .
   

4. 將 Docker 映像檔推送至 Artifact Registry 之前，請先為映像檔新增標記。 如要進一步瞭解如何使用 Artifact Registry，請參閱「使用容器映像檔」。

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

5. 將 Docker 映像檔推送至 Artifact Registry

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

6. 建立 TPU VM

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

7. 在所有 TPU 工作站上，從 Artifact Registry 提取 Docker 映像檔

   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. 在所有 TPU 工作站上執行容器

   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. 在所有 TPU 工作站上執行訓練指令碼

   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"
   

訓練指令碼執行完畢後，請清理資源。

1. 關閉所有工作人員的容器

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

2. 刪除 TPU VM

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

後續步驟

• Cloud TPU 教學課程
• 管理 TPU
• Cloud TPU 架構
• 在 TPU 配量上執行 JAX 程式碼