Menjalankan kode PyTorch di slice TPU

Perhatian: Halaman ini berlaku untuk Cloud TPU API, yang tidak lagi dalam pengembangan aktif. Cloud TPU API hanya akan menerima perbaikan bug dan update keamanan. Untuk mendapatkan fitur terbaru dan dukungan untuk versi TPU terbaru, sebaiknya gunakan Compute Engine atau Google Kubernetes Engine untuk mengelola resource TPU. Untuk mengetahui informasi selengkapnya, lihat Resource TPU di Compute Engine dan Tentang TPU di GKE. Cloud TPU API tidak didukung di TPU7x dan versi TPU yang lebih baru.

Tidak digunakan lagi: Halaman ini berlaku untuk PyTorch/XLA, yang tidak digunakan lagi dan digantikan oleh TorchTPU. Repositori PyTorch/XLA akan diarsipkan pada 15 November 2026. Mulai workload PyTorch baru dengan TorchTPU dan migrasikan workload PyTorch/XLA yang ada ke TorchTPU.

Sebelum menjalankan perintah dalam dokumen ini, ikuti petunjuk di bagian Menyiapkan akun dan project Cloud TPU.

Setelah kode PyTorch berjalan di satu VM TPU, Anda dapat meningkatkan skala 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 merupakan pengantar untuk menjalankan kode PyTorch di slice TPU.

Peran yang diperlukan

Untuk mendapatkan izin yang Anda perlukan untuk membuat TPU dan terhubung ke TPU menggunakan SSH, minta administrator untuk memberi Anda peran IAM berikut pada project Anda:

• TPU Admin (roles/tpu.admin)
• Service Account User (roles/iam.serviceAccountUser)
• Compute Viewer (roles/compute.viewer)

Untuk mengetahui informasi selengkapnya tentang pemberian peran, lihat Mengelola akses ke project, folder, dan organisasi.

Anda mungkin juga bisa mendapatkan izin yang diperlukan melalui peran khusus atau peran bawaan lainnya.

Membuat slice Cloud TPU

1. Tentukan variabel lingkungan untuk parameter perintah 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

   Deskripsi variabel lingkungan

   • PROJECT_ID: Project ID Google Cloud Anda. Gunakan project yang sudah ada atau buat project baru.
   • TPU_NAME: Nama TPU.
   • ZONE: Zona tempat membuat VM TPU. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
   • 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.
   • RUNTIME_VERSION: Versi software Cloud TPU.

2. Buat VM TPU Anda 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 Anda

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

1. Buat file bernama requirements.txt dengan konten berikut:

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

2. Salin requirements.txt ke setiap VM di slice Anda:

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

3. Instal dependensi di setiap VM di slice Anda:

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

4. Clone XLA di semua worker 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 Anda

Jalankan skrip pelatihan di semua worker. Skrip pelatihan menggunakan strategi sharding Single Program Multiple Data (SPMD). Untuk mengetahui 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 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 Anda.

1. Putuskan koneksi dari instance Cloud TPU, jika Anda belum melakukannya:

   exit
   

   Prompt 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 $TPU_NAME --zone=$ZONE
   

3. Pastikan 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