Menjalankan kode PyTorch di slice TPU

Sebelum menjalankan perintah dalam dokumen ini, pastikan Anda telah mengikuti petunjuk di Menyiapkan akun dan project Cloud TPU.

Setelah kode PyTorch Anda berjalan di satu VM TPU, Anda dapat menskalakan kode dengan menjalankannya di slice TPU. Slice TPU adalah beberapa papan 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 guna membuat TPU dan terhubung ke TPU menggunakan SSH, minta administrator Anda untuk memberi Anda peran IAM berikut di project Anda:

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 beberapa variabel lingkungan agar perintah lebih mudah digunakan.

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=europe-west4-b
export ACCELERATOR_TYPE=v5p-32
export RUNTIME_VERSION=v2-alpha-tpuv5

    Deskripsi variabel lingkungan

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

Jika perintah berikut gagal karena error koneksi SSH, hal ini mungkin disebabkan 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 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 memerlukan waktu sekitar 15 menit. Setelah selesai, Anda akan melihat pesan yang mirip 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:

    (vm)$ exit

    Kini perintah Anda akan menjadi username@projectname, yang menunjukkan Anda berada dalam 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 yang dibuat dalam tutorial ini:

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