Membuat slice TPU multi-host

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

Sebelum memulai

  1. Login keakun Anda. Google Cloud 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.

  2. Instal Google Cloud CLI.

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

  4. Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut: 

    gcloud init

  5. Buat atau pilih Google Cloud project.

    Peran yang diperlukan untuk memilih atau membuat project

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

    • Buat Google Cloud project:

      gcloud projects create PROJECT_ID

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

    • Pilih Google Cloud project yang telah Anda buat:

      gcloud config set project PROJECT_ID

      Ganti PROJECT_ID dengan nama projectAnda. Google Cloud

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

  7. Pastikan penagihan diaktifkan untuk Google Cloud project Anda.

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

  9. Instal Google Cloud CLI.

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

  11. Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut: 

    gcloud init

  12. Buat atau pilih Google Cloud project.

    Peran yang diperlukan untuk memilih atau membuat project

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

    • Buat Google Cloud project:

      gcloud projects create PROJECT_ID

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

    • Pilih Google Cloud project yang telah Anda buat:

      gcloud config set project PROJECT_ID

      Ganti PROJECT_ID dengan nama projectAnda. Google Cloud

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

  14. Pastikan penagihan diaktifkan untuk Google Cloud project Anda.

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

Peran yang diperlukan

Untuk mendapatkan izin yang Anda perlukan guna membuat MIG yang membentuk slice TPU multi-host, terhubung ke setiap VM di MIG menggunakan SSH, dan menjalankan perintah, 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 template instance

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

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 workload

Kebijakan workload menentukan properti fisik instance komputasi Anda. Dalam slice TPU, topologi akselerator menentukan pengaturan fisik chip TPU. Menentukan topologi akselerator diperlukan untuk slice TPU multi-host yang saling terhubung.

Untuk membuat kebijakan workload bagi slice TPU multi-host, gunakan gcloud compute resource-policies create workload-policy perintah dengan flag --accelerator-topology. Perintah berikut akan membuat kebijakan workload 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.

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

    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 Google Cloud project IDAnda.

  2. Pastikan instance terkelola berjalan menggunakan perintah berikut:

Menginstal JAX

Instal dependensi dan framework JAX di lingkungan virtual pada semua instance VM TPU di MIG. Jika VM TPU Anda memiliki versi Python 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 di slice TPU, Anda harus menjalankan kode di setiap host dalam slice TPU. Panggilan fungsi jax.device_count() akan berhenti merespons hingga dipanggil di setiap host dalam slice. Contoh berikut menunjukkan cara menjalankan perhitungan JAX di 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"

Menjalankan kode di 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 Anda tidak dikenai biaya untuk resource yang digunakan pada halaman ini, hapus project yang berisi resource tersebut. Google Cloud Google Cloud

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

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

Langkah berikutnya