Menjalankan penghitungan di VM Cloud TPU menggunakan PyTorch

Menjalankan penghitungan di VM Cloud TPU menggunakan PyTorch

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.

Dokumen ini memberikan pengantar singkat tentang cara menggunakan PyTorch dan Cloud TPU.

Catatan: Contoh ini menunjukkan cara menjalankan kode di TPU v5litepod-8 (v5e), yang merupakan TPU host tunggal. TPU host tunggal hanya memiliki satu VM TPU. Untuk menjalankan kode di TPU dengan lebih dari satu VM TPU (misalnya, v5litepod-32 atau yang lebih besar), lihat Menjalankan kode PyTorch di slice Cloud TPU.

Sebelum memulai

Sebelum menjalankan perintah dalam dokumen ini, Anda harus membuat akun Google Cloud , menginstal Google Cloud CLI, dan mengonfigurasi perintah gcloud. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan lingkungan Cloud 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 Cloud TPU menggunakan gcloud

1. Tentukan beberapa variabel lingkungan untuk mempermudah penggunaan perintah.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-east5-a
   export ACCELERATOR_TYPE=v5litepod-8
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite

   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 \
        --project=$PROJECT_ID \
        --zone=$ZONE \
        --accelerator-type=$ACCELERATOR_TYPE \
        --version=$RUNTIME_VERSION
   

Menghubungkan ke VM Cloud TPU Anda

Hubungkan ke VM TPU Anda melalui SSH menggunakan perintah berikut:

gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

Jika Anda tidak dapat terhubung ke VM TPU menggunakan 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.

Menginstal PyTorch/XLA di VM TPU Anda

1. Simpan kode berikut ke file bernama requirements.txt:

   --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
   ray[default]==2.40.0
   

2. Instal dependensinya:

   pip install -r requirements.txt
   

Memastikan PyTorch dapat mengakses TPU

Gunakan perintah berikut untuk memverifikasi bahwa PyTorch dapat mengakses TPU Anda:

PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

Output dari perintah akan terlihat seperti berikut:

['xla:0', 'xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

Melakukan penghitungan dasar

1. Buat file bernama tpu-test.py di direktori saat ini dan tambahkan skrip berikut ke dalamnya:

   import torch
   import torch_xla.core.xla_model as xm
   
   dev = xm.xla_device()
   t1 = torch.randn(3,3,device=dev)
   t2 = torch.randn(3,3,device=dev)
   print(t1 + t2)
   

2. Jalankan skrip:

   PJRT_DEVICE=TPU python3 tpu-test.py
   

   Output dari skrip menunjukkan hasil komputasi:

   tensor([[-0.2121,  1.5589, -0.6951],
            [-0.7886, -0.2022,  0.9242],
            [ 0.8555, -1.8698,  1.4333]], device='xla:1')
   

Pembersihan

Agar akun Google Cloud Anda tidak dikenai biaya untuk resource yang digunakan pada halaman ini, ikuti langkah-langkah berikut.

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 Cloud TPU Anda.

   gcloud compute tpus tpu-vm delete $TPU_NAME \
        --project=$PROJECT_ID \
        --zone=$ZONE
   

3. Pastikan resource telah dihapus dengan menjalankan perintah berikut. Pastikan TPU Anda tidak lagi tercantum. Penghapusan mungkin memerlukan waktu beberapa menit.

   gcloud compute tpus tpu-vm list \
        --zone=$ZONE
   

Langkah berikutnya

Baca selengkapnya tentang VM Cloud TPU:

• Menjalankan kode PyTorch di slice TPU
• Mengelola TPU
• Arsitektur sistem Cloud TPU
• Dokumentasi PyTorch/XLA