Membuat slice TPU multi-host

Membuat template instance

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

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.

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 Cloud de Confiance by S3NS 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.]