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

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

在 Docker 容器中訓練 PyTorch 模型

Cloud TPU 上的 PyTorch 說明 TPU 的 PyTorch 程式庫，並提供取得 Python 套件和 Docker 映像檔的說明文件連結。您可以直接使用 TorchTPU (請參閱「開始使用 GKE」，瞭解公開發布的容器映像檔)，或使用 TorchTitan 參考程式庫，建構及執行容器化訓練工作負載。

在 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 worker 上執行 JAX 程式碼。其中一個方法是使用 gcloud compute tpus tpu-vm ssh 指令，並搭配 --worker=all 和 --command 旗標。下列程序說明如何建立 Docker 映像檔，以便更輕鬆地設定各個 TPU worker。

  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 worker 上執行容器

    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 worker 上執行訓練指令碼

    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. 關閉所有 worker 上的容器

    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

後續步驟