Menjalankan kode PyTorch pada slice Pod TPU

Sebelum menjalankan perintah dalam dokumen ini, pastikan Anda telah mengikuti petunjuk di Menyiapkan akun dan project Cloud TPU.

Setelah kode PyTorch berjalan di satu VM TPU, Anda dapat menskalakan kode dengan menjalankannya di slice Pod TPU. Slice Pod TPU adalah beberapa board TPU yang terhubung satu sama lain melalui koneksi jaringan khusus berkecepatan tinggi. Dokumen ini adalah pengantar untuk menjalankan kode PyTorch pada slice Pod TPU.

Membuat slice Pod Cloud TPU

1. Tentukan beberapa variabel lingkungan agar perintah lebih mudah digunakan.

   export PROJECT_ID=your-project
   export ACCELERATOR_TYPE=v5p-32
   export ZONE=europe-west4-b
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export TPU_NAME=your-tpu-name

   Deskripsi variabel lingkungan

   PROJECT_ID

   Project ID Google Cloud Anda.

   ACCELERATOR_TYPE

   Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.

   ZONE

   Zona tempat Anda berencana membuat Cloud TPU.

   RUNTIME_VERSION

   Versi software Cloud TPU.

   TPU_NAME

   Nama yang ditetapkan pengguna untuk Cloud TPU Anda.

2. Buat VM TPU dengan menjalankan perintah berikut:

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

Menginstal PyTorch/XLA di slice Pod

Setelah membuat slice Pod TPU, Anda harus menginstal PyTorch di semua host dalam slice Pod TPU. Anda dapat melakukannya menggunakan perintah gcloud compute tpus tpu-vm ssh menggunakan parameter --worker=all dan --commamnd.

1. 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"

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

Menjalankan skrip pelatihan di slice Pod TPU

Jalankan 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 serupa seperti berikut:

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}