Google 會運用 AI 技術將內容翻譯成你偏好的語言，但可能會出錯。

在 TPU slice 上執行 PyTorch 程式碼

執行本文件中的指令前，請先按照「設定帳戶和 Cloud TPU 專案」中的操作說明操作。

在單一 TPU VM 上執行 PyTorch 程式碼後，您可以在 TPU 切片上執行程式碼，藉此擴充程式碼。TPU 配量是指多個 TPU 板，透過專屬高速網路連線相互連結。本文將介紹如何在 TPU slice 上執行 PyTorch 程式碼。

必要的角色

如要取得建立 TPU 並使用 SSH 連線所需的權限，請要求管理員在專案中授予您下列 IAM 角色：

TPU 管理員 (roles/tpu.admin)
服務帳戶使用者 (roles/iam.serviceAccountUser)
Compute 檢視者 (roles/compute.viewer)

如要進一步瞭解如何授予角色，請參閱「管理專案、資料夾和組織的存取權」。

您或許也能透過自訂角色或其他預先定義的角色，取得必要權限。

建立 Cloud TPU 切片

為 Google Cloud CLI 指令參數定義環境變數。
```
export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-central1-b
export ACCELERATOR_TYPE=v6e-8
export RUNTIME_VERSION=v2-alpha-tpuv6e
```
環境變數說明
- PROJECT_ID：您的 Google Cloud 專案 ID。使用現有專案或建立新專案。
- TPU_NAME：TPU 的名稱。
- ZONE：建立 TPU VM 的可用區。如要進一步瞭解支援的可用區，請參閱「TPU 區域和可用區」。
- ACCELERATOR_TYPE：加速器類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。
- RUNTIME_VERSION：Cloud TPU 軟體版本。

執行下列指令，建立 TPU VM：

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

在 slice 上安裝 PyTorch/XLA

建立 TPU 切片後，您必須在 TPU 切片中的所有主機上安裝 PyTorch。您可以使用 gcloud compute tpus tpu-vm ssh 指令，並搭配 --worker=all 和 --command 參數來執行這項操作。

建立名為 requirements.txt 的檔案，並加入下列內容：

--find-links https://storage.googleapis.com/libtpu-releases/index.html
--find-links https://storage.googleapis.com/libtpu-wheels/index.html
torch~=2.6.0
torch_xla[tpu]~=2.6.0
torchvision
ray[default]==2.40.0

將 requirements.txt 複製到切片中的每個 VM：

gcloud compute tpus tpu-vm tpu-vm scp ./requirements.txt \
     $TPU_NAME:~/ \
     --zone=$ZONE \
     --project=$PROJECT_ID \
     --worker=all

在配量中的每個 VM 上安裝依附元件：

gcloud compute tpus tpu-vm ssh $TPU_NAME \
     --zone=$ZONE \
     --project=$PROJECT_ID \
     --worker=all \
     --command="pip3 install -r requirements.txt"

在所有 TPU VM worker 上複製 XLA：

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

在 TPU slice 上執行訓練指令碼

在所有 worker 上執行訓練指令碼。訓練指令碼採用單一程式多重資料 (SPMD) 分區策略。如要進一步瞭解 SPMD，請參閱 PyTorch/XLA SPMD 使用手冊。

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --zone=$ZONE \
   --project=$PROJECT_ID \
   --worker=all \
   --command="PJRT_DEVICE=TPU python3 ~/xla/test/spmd/test_train_spmd_imagenet.py  \
   --fake_data \
   --model=resnet50  \
   --num_epochs=1 2>&1 | tee ~/logs.txt"

訓練作業約需 15 分鐘。完成後，您應該會看到類似下方的訊息：

   Epoch 1 test end 23:49:15, Accuracy=100.00
   10.164.0.11 [0] Max Accuracy: 100.00%

清除所用資源

完成 TPU VM 後，請按照下列步驟清除資源。

如果尚未中斷與 Cloud TPU 執行個體的連線，請執行下列指令：
```
exit
```
提示詞現在應該是 username@projectname，表示您位於 Cloud Shell 中。

刪除 Cloud TPU 資源。

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

執行 gcloud compute tpus tpu-vm list，確認資源已刪除。刪除作業可能需要幾分鐘才能完成。下列指令的輸出內容不應包含本教學課程中建立的任何資源：
```
gcloud compute tpus tpu-vm list --zone=$ZONE
```