使用 TPU v5e 訓練模型

每個 Pod 的晶片數量較少 (256 個)，因此 TPU v5e 經過最佳化調整後，成為高價值的產品，適合用於 Transformer、文字轉圖像和卷積類神經網路 (CNN) 的訓練、微調和服務。如要進一步瞭解如何使用 Cloud TPU v5e 進行服務，請參閱「使用 v5e 執行推論作業」。

如要進一步瞭解 Cloud TPU v5e TPU 硬體和設定，請參閱「TPU v5e」。

開始使用

以下各節說明如何開始使用 TPU v5e。

要求配額

如要使用 TPU v5e 進行訓練，您需要配額。隨選 TPU、預留 TPU 和 TPU Spot VM 的配額類型不同。如果使用 TPU v5e 進行推論，則需要不同的配額。如要進一步瞭解配額，請參閱「配額」。如要申請 TPU v5e 配額，請與 Cloud 銷售團隊聯絡。

建立 Google Cloud 帳戶和專案

如要使用 Cloud TPU，您需要 Google Cloud 帳戶和專案。詳情請參閱「設定 Cloud TPU 環境」。

建立 Cloud TPU

最佳做法是使用 queued-resource create 指令，將 Cloud TPU v5e 佈建為已加入佇列的資源。詳情請參閱「管理佇列資源」。

您也可以使用 Create Node API (gcloud compute tpus tpu-vm create) 佈建 Cloud TPU v5e。詳情請參閱「管理 TPU 資源」。

如要進一步瞭解可用的 v5e 訓練設定，請參閱「Cloud TPU v5e 訓練類型」。

設定架構

本節說明使用 JAX 或 PyTorch 和 TPU v5e 訓練自訂模型的一般設定程序。

如需推論設定操作說明，請參閱 v5e 推論簡介。

定義一些環境變數：

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

設定 JAX

如果切片形狀大於 8 個晶片，則一個切片中會有數個 VM。在這種情況下，您需要使用 --worker=all 旗標，在單一步驟中於所有 TPU VM 上執行安裝作業，不必使用 SSH 分別登入每個 VM：

gcloud compute tpus tpu-vm ssh $TPU_NAME  \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

指令旗標說明

  • TPU_NAME：TPU 的使用者指派文字 ID，在分配排入佇列的資源要求時建立。
  • PROJECT_ID： Google Cloud 專案名稱。使用現有專案或建立新專案，請參閱「 設定專案」一文。 Google Cloud
  • ZONE：如要瞭解支援的區域，請參閱 TPU 區域和可用區文件。
  • worker：可存取基礎 TPU 的 TPU VM。

您可以執行下列指令，檢查裝置數量 (這裡顯示的輸出內容是使用 v5litepod-16 切片產生)。這段程式碼會檢查 JAX 是否能看到 Cloud TPU TensorCore，並執行基本作業，藉此測試所有項目是否已正確安裝：

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

畫面中會顯示如下的輸出結果：

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() 會顯示指定切片中的晶片總數。 jax.local_device_count() 表示這個切片中單一 VM 可存取的晶片數量。

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

畫面中會顯示如下的輸出結果：

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

請參閱本文中的 JAX 教學課程，開始使用 JAX 進行 v5e 訓練。

PyTorch 設定

請注意，v5e 僅支援 PJRT 執行階段，且 PyTorch 2.1 以上版本會將 PJRT 設為所有 TPU 版本的預設執行階段。

本節說明如何開始在 v5e 上使用 PJRT，搭配 PyTorch/XLA，並為所有工作人員提供指令。

安裝依附元件

gcloud compute tpus tpu-vm ssh $TPU_NAME  \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip install mkl mkl-include
      pip install tf-nightly tb-nightly tbp-nightly
      pip install numpy
      sudo apt-get install libopenblas-dev -y
      pip install torch~=PYTORCH_VERSION torchvision torch_xla[tpu]~=PYTORCH_VERSION -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html'

將 PYTORCH_VERSION 替換為您要使用的 PyTorch 版本。PYTORCH_VERSION 用於指定 PyTorch/XLA 的相同版本。建議使用 2.6.0 版。

如要進一步瞭解 PyTorch 和 PyTorch/XLA 版本，請參閱「PyTorch - Get Started」和「PyTorch/XLA releases」。

如要進一步瞭解如何安裝 PyTorch/XLA，請參閱 PyTorch/XLA 安裝。

如果安裝 torch、torch_xla 或 torchvision 的 Wheel 時發生錯誤 (例如 pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222)，請使用下列指令降級版本：

pip3 install setuptools==62.1.0

使用 PJRT 執行指令碼

unset LD_PRELOAD

以下範例使用 Python 指令碼，在 v5e VM 上執行計算：

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

這麼做會產生類似以下內容的輸出結果：

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

請參閱本文中的 PyTorch 教學課程，開始使用 PyTorch 訓練 v5e。

在工作階段結束時刪除 TPU 和排入佇列的資源。如要刪除排入佇列的資源，請分 2 個步驟刪除切片和排入佇列的資源：

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

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=${ZONE} \
   --quiet

這兩個步驟也可用於移除處於 FAILED 狀態的已排隊資源要求。

JAX/FLAX 範例

以下各節提供範例，說明如何在 TPU v5e 上訓練 JAX 和 FLAX 模型。

在 v5e 上訓練 ImageNet

本教學課程說明如何使用偽輸入資料，在 v5e 上訓練 ImageNet。如要使用真實資料，請參閱 GitHub 上的 README 檔案。

設定

  1. 建立環境變數：

    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
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

    環境變數說明

    • PROJECT_ID： 您的 Google Cloud 專案 ID。使用現有專案或建立新專案。
    • TPU_NAME：TPU 的名稱。
    • ZONE： 要建立 TPU VM 的可用區。如要進一步瞭解支援的區域，請參閱 TPU 地區和區域。
    • ACCELERATOR_TYPE： 加速器類型會指定您要建立的 Cloud TPU 版本和大小。如要進一步瞭解各個 TPU 版本支援的加速器類型，請參閱 TPU 版本。
    • RUNTIME_VERSION：Cloud TPU 軟體版本。
    • SERVICE_ACCOUNT： 服務帳戶的電子郵件地址。前往 Google Cloud 控制台的「Service Accounts」(服務帳戶) 頁面，即可找到服務帳戶。

      例如： tpu-service-account@PROJECT_ID.iam.gserviceaccount.com

    • QUEUED_RESOURCE_ID：使用者指派的已加入佇列資源要求文字 ID。

  2. 建立 TPU 資源：

    gcloud compute tpus queued-resources create $QUEUED_RESOURCE_ID \
   --node-id=$TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --runtime-version=$RUNTIME_VERSION \
   --service-account=$SERVICE_ACCOUNT

    排入佇列的資源處於 ACTIVE 狀態時，您就能透過 SSH 連線至 TPU VM：

    gcloud compute tpus queued-resources describe $QUEUED_RESOURCE_ID  \
   --project=$PROJECT_ID \
   --zone=$ZONE

    QueuedResource 處於 ACTIVE 狀態時，輸出內容會類似於下列內容：

     state: ACTIVE

  3. 安裝最新版 JAX 和 jaxlib：

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. 複製 ImageNet 模型，並安裝對應需求：

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command="git clone https://github.com/coolkp/flax.git && cd flax && git checkout pmap-orbax-conversion && git pull"

  5. 如要產生虛假資料，模型需要資料集維度的相關資訊。這項資訊可從 ImageNet 資料集的中繼資料收集而來：

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command="cd flax/examples/imagenet && pip install -r requirements-cloud-tpu.txt"

訓練模型

完成所有前述步驟後，即可訓練模型。

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command="cd flax/examples/imagenet && bash ../../tests/download_dataset_metadata.sh && JAX_PLATFORMS=tpu python imagenet_fake_data_benchmark.py"

刪除 TPU 和排入佇列的資源

在工作階段結束時，刪除 TPU 和排入佇列的資源。

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

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

Hugging Face FLAX 模型

以 FLAX 實作的 Hugging Face 模型可在 Cloud TPU v5e 上直接運作。本節提供執行熱門模型的操作說明。

在 Imagenette 上訓練 ViT

本教學課程說明如何使用 Fast AI Imagenette 資料集，在 Cloud TPU v5e 上訓練 HuggingFace 的 Vision Transformer (ViT) 模型。

ViT 模型是第一個成功在 ImageNet 上訓練 Transformer 編碼器，且與卷積網路相比，結果相當出色的模型。詳情請參閱 ViT 總覽。

設定

  1. 建立環境變數：

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-16
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

    環境變數說明

    • PROJECT_ID： 您的 Google Cloud 專案 ID。使用現有專案或建立新專案。
    • TPU_NAME：TPU 的名稱。
    • ZONE： 要建立 TPU VM 的可用區。如要進一步瞭解支援的區域，請參閱 TPU 地區和區域。
    • ACCELERATOR_TYPE： 加速器類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各個 TPU 版本支援的加速器類型，請參閱 TPU 版本。
    • RUNTIME_VERSION：Cloud TPU 軟體版本。
    • SERVICE_ACCOUNT： 服務帳戶的電子郵件地址。前往 Google Cloud 控制台的「Service Accounts」(服務帳戶) 頁面，即可找到服務帳戶。

      例如： tpu-service-account@