Crea uno slice TPU multi-host

Crea uno slice TPU multi-host

Scopri come creare uno slice TPU multi-host utilizzando un gruppo di istanze gestite (MIG), connetterti allo slice ed eseguire un calcolo. Questa guida rapida utilizza l'opzione di consumo on demand. Esegui i comandi di questa guida rapida nel terminale locale o in Cloud Shell.

Prima di iniziare

Accedi al tuo account Google Cloud . Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti senza costi per l'esecuzione, il test e il deployment dei carichi di lavoro.

Installa Google Cloud CLI.

Nota: se hai già installato gcloud CLI, assicurati di avere l'ultima versione eseguendo gcloud components update.

Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

Per inizializzare gcloud CLI, esegui questo comando:

gcloud init

Crea o seleziona un Google Cloud progetto.

Ruoli richiesti per selezionare o creare un progetto

• Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
• Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

Nota: se non prevedi di conservare le risorse che crei in questa procedura, crea un progetto invece di selezionarne uno già esistente. Una volta completata questa procedura, puoi eliminare il progetto e tutte le relative risorse.

• Creare un progetto Google Cloud :

  gcloud projects create PROJECT_ID

  Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.

• Seleziona il progetto Google Cloud che hai creato:

  gcloud config set project PROJECT_ID

  Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

Se utilizzi un progetto esistente per questa guida, verifica di disporre delle autorizzazioni necessarie per completare la guida. Se hai creato un nuovo progetto, disponi già delle autorizzazioni necessarie.

Verifica che la fatturazione sia attivata per il tuo progetto Google Cloud .

Abilita l'API Compute Engine:

Ruoli richiesti per abilitare le API

Per abilitare le API, devi disporre dell'autorizzazione serviceusage.services.enable. Se hai creato il progetto, probabilmente disponi già di questa autorizzazione tramite il ruolo Proprietario (roles/owner). In caso contrario, puoi ottenere questa autorizzazione tramite il ruolo Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin). Scopri come concedere i ruoli.

gcloud services enable compute.googleapis.com

Installa Google Cloud CLI.

Nota: se hai già installato gcloud CLI, assicurati di avere l'ultima versione eseguendo gcloud components update.

Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

Per inizializzare gcloud CLI, esegui questo comando:

gcloud init

Crea o seleziona un Google Cloud progetto.

Ruoli richiesti per selezionare o creare un progetto

• Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
• Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

Nota: se non prevedi di conservare le risorse che crei in questa procedura, crea un progetto invece di selezionarne uno già esistente. Una volta completata questa procedura, puoi eliminare il progetto e tutte le relative risorse.

• Creare un progetto Google Cloud :

  gcloud projects create PROJECT_ID

  Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.

• Seleziona il progetto Google Cloud che hai creato:

  gcloud config set project PROJECT_ID

  Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

Se utilizzi un progetto esistente per questa guida, verifica di disporre delle autorizzazioni necessarie per completare la guida. Se hai creato un nuovo progetto, disponi già delle autorizzazioni necessarie.

Verifica che la fatturazione sia attivata per il tuo progetto Google Cloud .

Abilita l'API Compute Engine:

Ruoli richiesti per abilitare le API

Per abilitare le API, devi disporre dell'autorizzazione serviceusage.services.enable. Se hai creato il progetto, probabilmente disponi già di questa autorizzazione tramite il ruolo Proprietario (roles/owner). In caso contrario, puoi ottenere questa autorizzazione tramite il ruolo Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin). Scopri come concedere i ruoli.

gcloud services enable compute.googleapis.com

Ruoli obbligatori

Per ottenere le autorizzazioni necessarie per creare un gruppo di istanze gestite che forma uno slice TPU multihost, connetterti a ogni VM nel gruppo di istanze gestite utilizzando SSH ed eseguire comandi, chiedi all'amministratore di concederti i seguenti ruoli IAM nel progetto:

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

Per saperne di più sulla concessione dei ruoli, consulta Gestisci l'accesso a progetti, cartelle e organizzazioni.

Potresti anche riuscire a ottenere le autorizzazioni richieste tramite i ruoli personalizzati o altri ruoli predefiniti.

Crea un template di istanza

Per creare un modello di istanza per le VM TPU v6e, utilizza il comando gcloud compute instance-templates create:

Nota: se utilizzi un tipo di macchina diverso, assicurati di utilizzare un'immagine del sistema operativo supportata. Per un elenco delle immagini sistema operativo gestite da Google per ogni versione di TPU supportata, consulta Immagini sistema operativo 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

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.

Nota: potresti riscontrare un errore di disponibilità temporanea delle risorse mentre il MIG mette in coda la capacità. Se i tempi di attesa sono prolungati, puoi provare con un'altra regione. Per un elenco delle località supportate, consulta Regioni e zone TPU. Se scegli una regione diversa, devi prima creare una policy del workload in quella regione ripetendo il passaggio Crea una policy del workload. Assicurati di aggiornare la regione in tutti i comandi successivi.

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

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

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questa pagina, elimina il progetto Google Cloud con le risorse.

In alternativa, se vuoi conservare il progetto, puoi eliminare solo il gruppo di istanze gestite e tutte le VM del gruppo utilizzando il comando gcloud compute instance-groups managed delete:

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

Passaggi successivi

• Scopri di più sulle risorse TPU in Compute Engine
• Crea uno slice TPU multi-host utilizzando altre opzioni di consumo
• Scopri di più sulle istanze TPU nei MIG