Membuat slice TPU multi-host

Membuat slice TPU multi-host

Pelajari cara membuat slice TPU multi-host menggunakan grup instance terkelola (MIG), menghubungkan ke slice, dan menjalankan perhitungan. Panduan memulai ini menggunakan opsi konsumsi on-demand. Jalankan perintah dalam panduan memulai ini di terminal lokal atau Cloud Shell.

Sebelum memulai

Login ke akun Google Cloud Anda. Jika Anda baru menggunakan Google Cloud, buat akun untuk mengevaluasi performa produk kami dalam skenario dunia nyata. Pelanggan baru juga mendapatkan kredit gratis senilai $300 untuk menjalankan, menguji, dan men-deploy workload.

Instal Google Cloud CLI.

Catatan: Jika Anda telah menginstal gcloud CLI sebelumnya, pastikan Anda memiliki versi terbaru dengan menjalankan gcloud components update.

Jika Anda menggunakan penyedia identitas (IdP) eksternal, Anda harus login ke gcloud CLI dengan identitas gabungan Anda terlebih dahulu.

Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut:

gcloud init

Buat atau pilih Google Cloud project.

Peran yang diperlukan untuk memilih atau membuat project

• Pilih project: Memilih project tidak memerlukan peran IAM tertentu—Anda dapat memilih project mana pun yang telah diberi peran.
• Membuat project: Untuk membuat project, Anda memerlukan peran Project Creator (roles/resourcemanager.projectCreator), yang berisi izin resourcemanager.projects.create. Pelajari cara memberikan peran.

Catatan: Jika Anda tidak berencana untuk menyimpan resource yang Anda buat dalam prosedur ini, buatlah project, bukan memilih project yang ada. Setelah menyelesaikan langkah ini, Anda dapat menghapus project, yang menghapus semua resource yang terkait dengan project ini.

• Buat Google Cloud project:

  gcloud projects create PROJECT_ID

  Ganti PROJECT_ID dengan nama untuk Google Cloud project yang Anda buat.

• Pilih project Google Cloud yang Anda buat:

  gcloud config set project PROJECT_ID

  Ganti PROJECT_ID dengan nama project Google Cloud Anda.

Jika Anda menggunakan project yang sudah ada untuk panduan ini, pastikan Anda memiliki izin yang diperlukan untuk menyelesaikan panduan ini. Jika berhasil membuat project baru, berarti Anda sudah memiliki izin yang diperlukan.

Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

Aktifkan Compute Engine API:

Peran yang diperlukan untuk mengaktifkan API

Untuk mengaktifkan API, Anda memerlukan izin serviceusage.services.enable. Jika Anda membuat project, kemungkinan Anda sudah memiliki izin ini melalui peran Pemilik (roles/owner). Jika tidak, Anda bisa mendapatkan izin ini melalui peran Admin Penggunaan Layanan (roles/serviceusage.serviceUsageAdmin). Pelajari cara memberikan peran.

gcloud services enable compute.googleapis.com

Instal Google Cloud CLI.

Catatan: Jika Anda telah menginstal gcloud CLI sebelumnya, pastikan Anda memiliki versi terbaru dengan menjalankan gcloud components update.

Jika Anda menggunakan penyedia identitas (IdP) eksternal, Anda harus login ke gcloud CLI dengan identitas gabungan Anda terlebih dahulu.

Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut:

gcloud init

Buat atau pilih Google Cloud project.

Peran yang diperlukan untuk memilih atau membuat project

• Pilih project: Memilih project tidak memerlukan peran IAM tertentu—Anda dapat memilih project mana pun yang telah diberi peran.
• Membuat project: Untuk membuat project, Anda memerlukan peran Project Creator (roles/resourcemanager.projectCreator), yang berisi izin resourcemanager.projects.create. Pelajari cara memberikan peran.

Catatan: Jika Anda tidak berencana untuk menyimpan resource yang Anda buat dalam prosedur ini, buatlah project, bukan memilih project yang ada. Setelah menyelesaikan langkah ini, Anda dapat menghapus project, yang menghapus semua resource yang terkait dengan project ini.

• Buat Google Cloud project:

  gcloud projects create PROJECT_ID

  Ganti PROJECT_ID dengan nama untuk Google Cloud project yang Anda buat.

• Pilih project Google Cloud yang Anda buat:

  gcloud config set project PROJECT_ID

  Ganti PROJECT_ID dengan nama project Google Cloud Anda.

Jika Anda menggunakan project yang sudah ada untuk panduan ini, pastikan Anda memiliki izin yang diperlukan untuk menyelesaikan panduan ini. Jika berhasil membuat project baru, berarti Anda sudah memiliki izin yang diperlukan.

Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

Aktifkan Compute Engine API:

Peran yang diperlukan untuk mengaktifkan API

Untuk mengaktifkan API, Anda memerlukan izin serviceusage.services.enable. Jika Anda membuat project, kemungkinan Anda sudah memiliki izin ini melalui peran Pemilik (roles/owner). Jika tidak, Anda bisa mendapatkan izin ini melalui peran Admin Penggunaan Layanan (roles/serviceusage.serviceUsageAdmin). Pelajari cara memberikan peran.

gcloud services enable compute.googleapis.com

Peran yang diperlukan

Agar mendapatkan izin yang diperlukan untuk membuat MIG yang membentuk slice TPU multi-host, hubungkan ke setiap VM di MIG menggunakan SSH, dan jalankan perintah, minta administrator untuk memberi Anda peran IAM berikut di project Anda:

• Compute Instance Admin (v1) (roles/compute.instanceAdmin.v1)
• Service Account User (roles/iam.serviceAccountUser)
• Service Usage Admin (roles/serviceusage.serviceUsageAdmin)

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

Untuk membuat template instance bagi VM TPU v6e, gunakan perintah gcloud compute instance-templates create:

Catatan: Jika Anda menggunakan jenis mesin yang berbeda, pastikan Anda menggunakan image OS yang didukung. Untuk mengetahui daftar image OS yang dikelola Google untuk setiap versi TPU yang didukung, lihat Image OS TPU.

gcloud compute instance-templates create quickstart-tpu-instance-template \
    --machine-type=ct6e-standard-4t \
    --maintenance-policy=TERMINATE \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --image-project=ubuntu-os-accelerator-images \
    --region=us-east5

Membuat kebijakan beban kerja

Kebijakan workload menentukan properti fisik instance komputasi Anda. Dalam slice TPU, topologi akselerator menentukan susunan fisik chip TPU. Topologi akselerator harus ditentukan untuk slice TPU multi-host yang saling terhubung.

Untuk membuat kebijakan beban kerja untuk slice TPU multi-host, gunakan perintah gcloud compute resource-policies create workload-policy dengan flag --accelerator-topology. Perintah berikut akan membuat kebijakan beban kerja dengan topologi 2x4:

gcloud compute resource-policies create workload-policy quickstart-tpu-workload-policy \
    --type=high-throughput \
    --accelerator-topology=2x4 \
    --region=us-east5

Membuat MIG

Jalankan perintah berikut untuk membuat MIG yang membentuk slice TPU multi-host.

Catatan: Anda mungkin mengalami error ketersediaan resource sementara saat MIG mengantre untuk mendapatkan kapasitas. Jika mengalami waktu tunggu yang lama, Anda dapat mencoba wilayah lain. Untuk mengetahui daftar lokasi yang didukung, lihat region dan zona TPU. Jika memilih region lain, Anda harus membuat kebijakan workload di region tersebut terlebih dahulu dengan mengulangi langkah Buat kebijakan workload. Pastikan untuk memperbarui region di semua perintah berikutnya.

1. Untuk membuat MIG yang membentuk slice TPU multi-host, gunakan perintah gcloud compute instance-groups managed create:

   gcloud compute instance-groups managed create quickstart-tpu-mig \
       --size=2 \
       --target-size-policy-mode=bulk \
       --template=quickstart-tpu-instance-template \
       --region=us-east5 \
       --target-distribution-shape=any-single-zone \
       --instance-redistribution-type=none \
       --default-action-on-vm-failure=do-nothing \
       --workload-policy=projects/PROJECT_ID/regions/us-east5/resourcePolicies/quickstart-tpu-workload-policy
   

   Ganti PROJECT_ID dengan project ID Google Cloud Anda.

2. (Opsional) Verifikasi bahwa instance terkelola berjalan menggunakan perintah berikut:

   • Untuk melihat status MIG secara keseluruhan, gunakan perintah gcloud compute instance-groups managed describe:

     gcloud compute instance-groups managed describe quickstart-tpu-mig \
         --region=us-east5
     

   • Untuk melihat status setiap instance, gunakan perintah gcloud compute instance-groups managed list-instances:

     gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
         --region=us-east5
     

Menginstal JAX

Instal dependensi dan framework JAX di lingkungan virtual pada semua instance VM TPU di MIG. Jika VM TPU Anda memiliki Python versi yang lebih lama dari 3.11, Anda harus menginstal Python 3.11 untuk menjalankan JAX versi terbaru.

1. Periksa versi Python yang berjalan di VM TPU Anda:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='python3 --version'
   

   Jika versinya lebih lama dari Python 3.11, instal Python 3.11:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt update && \
       sudo apt install -y software-properties-common && \
       sudo add-apt-repository -y ppa:deadsnakes/ppa && \
       sudo apt update && \
       sudo apt install -y python3.11 python3.11-dev'
   

2. Buat lingkungan virtual:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt install -y python3.11-venv && \
       python3.11 -m venv ~/jax_venv'
   

3. Instal JAX di lingkungan virtual:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='source ~/jax_venv/bin/activate && \
       pip install --upgrade pip -q && \
       pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q'
   

Menjalankan kode JAX di slice

Untuk menjalankan kode JAX pada slice TPU, Anda harus menjalankan kode di setiap host dalam slice TPU. Panggilan fungsi jax.device_count() berhenti merespons hingga dipanggil di setiap host dalam slice. Contoh berikut menunjukkan cara menjalankan penghitungan JAX pada slice TPU.

Menyiapkan kode

Buat file bernama example.py di setiap instance:

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command="cat << 'EOF' > ~/example.py
import jax

# Initialize the slice
jax.distributed.initialize()

# The total number of TPU cores in the slice
device_count = jax.device_count()

# The number of TPU cores attached to this host
local_device_count = jax.local_device_count()

# The psum is performed over all mapped devices across the slice
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)

# Print from a single host to avoid duplicated output
if jax.process_index() == 0:
    print('global device count:', jax.device_count())
    print('local device count:', jax.local_device_count())
    print('pmap result:', r)
EOF"

Jalankan kode pada slice

Jalankan program example.py di setiap VM TPU dalam slice:

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command='source ~/jax_venv/bin/activate && python3 ~/example.py'

Outputnya akan mirip dengan berikut ini:

global device count: 8
local device count: 4
pmap result: [8. 8. 8. 8.]

Pembersihan

Agar akun Google Cloud Anda tidak dikenai biaya untuk resource yang digunakan pada halaman ini, hapus project Google Cloud yang berisi resource tersebut.

Atau, jika ingin mempertahankan project, Anda dapat menghapus MIG dan semua VM dalam grup saja menggunakan perintah gcloud compute instance-groups managed delete:

gcloud compute instance-groups managed delete quickstart-tpu-mig --region=us-east5

Langkah berikutnya

• Pelajari resource TPU di Compute Engine
• Membuat slice TPU multi-host menggunakan opsi pemakaian lainnya
• Pelajari instance TPU di MIG