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

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 guna membuat TPU dan terhubung ke TPU menggunakan SSH, minta administrator Anda untuk memberi Anda peran IAM berikut di project Anda:

• Admin TPU (roles/tpu.admin)
• Pengguna Akun Layanan (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 agar perintah lebih mudah digunakan.

   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 Anda. Google Cloud Gunakan project yang ada atau buat project baru.
   • TPU_NAME: Nama TPU.
   • ZONE: Zona tempat VM TPU akan dibuat. 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

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

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
   

Memverifikasi bahwa 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 tersebut 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 tersebut 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 akunAnda tidak dikenai biaya untuk resource yang digunakan pada halaman ini, ikuti langkah-langkah berikut. Google Cloud

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

   exit
   

   Kini perintah Anda akan menjadi username@projectname, yang menunjukkan Anda berada dalam Cloud Shell.

2. Hapus Cloud TPU Anda.

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

3. Verifikasi bahwa 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 lebih lanjut tentang VM Cloud TPU:

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