Menjalankan kode PyTorch pada slice 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 TPU. Slice TPU adalah beberapa board TPU yang terhubung satu sama lain melalui koneksi jaringan khusus berkecepatan tinggi. Dokumen ini adalah pengantar untuk menjalankan kode PyTorch di slice TPU.

Membuat slice 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

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

Jika perintah berikut gagal karena error koneksi SSH, hal ini mungkin karena VM TPU tidak memiliki alamat IP eksternal. Untuk mengakses VM TPU tanpa alamat IP eksternal, ikuti petunjuk di Menghubungkan ke VM TPU tanpa alamat IP publik.

1. Instal 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. 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 TPU

Jalankan skrip pelatihan di semua pekerja. Skrip pelatihan menggunakan strategi sharding Single Program Multiple Data (SPMD). Untuk informasi selengkapnya tentang SPMD, lihat Panduan Pengguna SPMD PyTorch/XLA.

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"

Pelatihan ini memerlukan waktu sekitar 15 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 sambungan dari instance Cloud TPU, 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 Anda.

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

3. Verifikasi bahwa resource telah dihapus dengan menjalankan gcloud compute tpus tpu-vm 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}