Menjalankan kode PyTorch pada slice Pod TPU

PyTorch/XLA mengharuskan semua VM TPU dapat mengakses kode dan data 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 di project untuk mengizinkan traffic masuk untuk port 8470 - 8479. Untuk mengetahui informasi selengkapnya tentang cara menambahkan aturan firewall, lihat Menggunakan aturan firewall

Menyiapkan lingkungan Anda

1. Di Cloud Shell, jalankan perintah berikut untuk memastikan Anda menjalankan gcloud versi saat ini:

   $ gcloud components update

   Jika Anda perlu menginstal gcloud, gunakan perintah berikut:

   $ sudo apt install -y google-cloud-sdk

2. 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

1. Tambahkan sertifikat SSH ke project Anda:

   ssh-add ~/.ssh/google_compute_engine

2. 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.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"

3. Meng-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.5 https://github.com/pytorch/xla.git"

4. Menjalankan skrip pelatihan di 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 berikut 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.

1. Putuskan koneksi dari instance Compute Engine, jika Anda belum melakukannya:

   (vm)$ exit

   Perintah Anda sekarang akan menjadi username@projectname, yang menunjukkan bahwa Anda berada di Cloud Shell.

2. Hapus resource Cloud TPU dan Compute Engine Anda.

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

3. Verifikasi bahwa resource telah dihapus dengan menjalankan gcloud compute tpus execution-groups list. Penghapusan mungkin memerlukan waktu beberapa menit. Output dari perintah berikut tidak boleh menyertakan resource apa pun yang dibuat dalam tutorial ini:

   $ gcloud compute tpus tpu-vm list --zone=${ZONE}