Menjalankan kode PyTorch pada slice Pod TPU

PyTorch/XLA mengharuskan semua VM TPU dapat mengakses data dan kode model. Anda dapat menggunakan skrip startup untuk mendownload software yang diperlukan untuk mendistribusikan data model ke semua VM TPU.

Jika menghubungkan VM TPU ke Virtual Private Cloud (VPC), Anda harus menambahkan aturan firewall dalam project agar traffic masuk ke port 8470 - 8479. Untuk mengetahui informasi lebih lanjut tentang cara menambahkan aturan firewall, lihat Menggunakan aturan firewall

Menyiapkan lingkungan Anda

Di Cloud Shell, jalankan perintah berikut untuk memastikan Anda menjalankan versi gcloud saat ini:
```
$ gcloud components update
```
Jika Anda perlu menginstal gcloud, gunakan perintah berikut:
```
$ sudo apt install -y google-cloud-sdk
```

Buat beberapa variabel lingkungan:

$ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=us-central2-b
$ export RUNTIME_VERSION=tpu-ubuntu2204-base
$ export ACCELERATOR_TYPE=v4-32

Membuat VM TPU

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

Mengonfigurasi dan menjalankan skrip pelatihan

Tambahkan sertifikat SSH ke project Anda:
```
ssh-add ~/.ssh/google_compute_engine
```

Menginstal PyTorch/XLA di semua pekerja VM TPU

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --worker=all --command="
  pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"

Clone XLA di semua pekerja VM TPU

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

Menjalankan skrip pelatihan pada semua pekerja

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

Pelatihan ini memerlukan waktu sekitar 5 menit. Setelah selesai, Anda akan melihat pesan yang mirip dengan ini:

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

Pembersihan

Setelah selesai menggunakan VM TPU, ikuti langkah-langkah berikut untuk membersihkan resource Anda.

Putuskan koneksi dari Compute Engine:
```
(vm)$ exit
```
Verifikasi bahwa resource telah dihapus dengan menjalankan perintah berikut. Pastikan TPU Anda tidak lagi tercantum. Proses penghapusan mungkin memerlukan waktu beberapa menit.
```
$ gcloud compute tpus tpu-vm list \
  --zone europe-west4-a
```