Menjalankan penghitungan di VM Cloud TPU menggunakan PyTorch

Panduan memulai ini menunjukkan cara membuat Cloud TPU, menginstal PyTorch, dan menjalankan penghitungan sederhana di Cloud TPU. Untuk tutorial yang lebih mendalam yang menunjukkan cara melatih model di Cloud TPU, lihat salah satu Tutorial PyTorch Cloud TPU.

Sebelum memulai

Sebelum mengikuti panduan memulai ini, Anda harus membuat akun Google Cloud Platform, menginstal Google Cloud CLI, dan mengonfigurasi perintah gcloud. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan akun dan project Cloud TPU.

Membuat Cloud TPU dengan gcloud

Untuk membuat VM TPU di project, jaringan, dan zona pengguna default, jalankan:

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central1-a \
   --accelerator-type=v3-8 \
   --version=tpu-ubuntu2204-base

Deskripsi flag perintah

zone
Zona tempat Anda berencana membuat Cloud 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.
version
Versi software Cloud TPU.

Saat membuat TPU, Anda dapat meneruskan flag --network dan --subnetwork tambahan jika ingin menentukan jaringan dan subnetwork default. Jika tidak ingin menggunakan jaringan default, Anda harus meneruskan flag --network. Flag --subnetwork bersifat opsional dan dapat digunakan untuk menentukan subnetwork default untuk jaringan apa pun yang Anda gunakan (default atau ditentukan pengguna). Lihat halaman referensi API gcloud untuk mengetahui detail tentang flag ini.

Menghubungkan ke VM Cloud TPU

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-a

Menginstal PyTorch/XLA di VM TPU

   (vm)$ pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html

Menetapkan konfigurasi runtime TPU

Pastikan runtime PyTorch/XLA menggunakan TPU.

   (vm) $ export PJRT_DEVICE=TPU

Memverifikasi bahwa PyTorch dapat mengakses TPU

  1. Buat file bernama tpu-count.py di direktori saat ini, lalu salin dan tempelkan skrip berikut ke dalamnya.

    import torch
import torch_xla
import torch_xla.core.xla_model as xm
print(f'PyTorch can access {len(torch_xla.devices())} TPU cores')

  2. Jalankan skrip:

    (vm)$ python3 tpu-count.py

    Output dari skrip menunjukkan hasil komputasi:

    PyTorch can access 8 TPU cores

Melakukan penghitungan dasar

  1. Buat file bernama tpu-test.py di direktori saat ini, lalu salin dan tempel 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:

      (vm)$ 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 tidak menimbulkan biaya pada akun Google Cloud Anda untuk resource yang digunakan pada halaman ini, ikuti langkah-langkah berikut.

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

    $ gcloud compute tpus tpu-vm delete tpu-name \
  --zone=us-central1-a

Output perintah ini akan mengonfirmasi bahwa TPU Anda telah dihapus.

Langkah selanjutnya

Baca selengkapnya tentang VM Cloud TPU: