Crea uno slice TPU multi-host

Crea un template di istanza

Per creare un modello di istanza per le VM TPU v6e, utilizza il comando 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

Crea una policy del workload

Una policy del workload definisce le proprietà fisiche delle istanze di Compute. Nelle slice TPU, la topologia dell'acceleratore definisce la disposizione fisica dei chip TPU. La specifica di una topologia dell'acceleratore è obbligatoria per gli slice TPU multi-host interconnessi.

Per creare una policy del workload per uno slice TPU multi-host, utilizza il comando gcloud compute resource-policies create workload-policy con il flag --accelerator-topology. Il seguente comando crea una policy del workload con una topologia 2x4:

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

Crea un MIG

Esegui questi comandi per creare un MIG che forma uno slice TPU multihost.

1. Per creare un MIG che forma uno slice TPU multi-host, utilizza il comando 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
   

   Sostituisci PROJECT_ID con l'ID progetto Cloud de Confiance by S3NS .

2. (Facoltativo) Verifica che le istanze gestite siano in esecuzione utilizzando i seguenti comandi:

   • Per visualizzare lo stato generale del MIG, utilizza il comando gcloud compute instance-groups managed describe:

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

   • Per visualizzare lo stato di ogni istanza, utilizza il comando gcloud compute instance-groups managed list-instances:

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

Installare JAX

Installa le dipendenze e il framework JAX in un ambiente virtuale su tutte le istanze VM TPU nel MIG. Se sulle tue VM TPU è installata una versione di Python precedente alla 3.11, devi installare Python 3.11 per eseguire l'ultima versione di JAX.

1. Controlla quale versione di Python è in esecuzione sulle VM TPU:

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

   Se la versione è precedente a Python 3.11, installa 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. Crea un ambiente virtuale:

   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. Installa JAX nell'ambiente virtuale:

   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'
   

Esegui il codice JAX sulla sezione

Per eseguire il codice JAX su una slice TPU, devi eseguire il codice su ogni host della slice TPU. La chiamata di funzione jax.device_count() smette di rispondere finché non viene chiamato ogni host nella sezione. L'esempio seguente mostra come eseguire un calcolo JAX su una sezione TPU.

Prepara il codice

Crea un file denominato example.py su ogni istanza:

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"

Esegui il codice sulla sezione

Esegui il programma example.py su ogni VM TPU nello 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'

L'output dovrebbe essere simile al seguente:

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