Crea una slice TPU multi-host

Scopri come creare una slice TPU multi-host utilizzando un gruppo di istanze gestite (MIG), connetterti alla 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

  1. Accedi al tuo Google Cloud account. Se non hai mai utilizzato Google Cloud, crea un account per valutare il rendimento 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.

  2. Installa Google Cloud CLI.

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

  4. Per inizializzare gcloud CLI, esegui questo comando: 

    gcloud init

  5. 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 su cui ti è stato concesso un ruolo.
    • Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l' resourcemanager.projects.create autorizzazione. Scopri come concedere i ruoli.

    • Crea un Google Cloud progetto:

      gcloud projects create PROJECT_ID

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

    • Seleziona il Google Cloud progetto che hai creato:

      gcloud config set project PROJECT_ID

      Sostituisci PROJECT_ID con il nome del Google Cloud progetto.

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

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

  8. 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 hai già questa autorizzazione tramite il ruolo Proprietario (roles/owner). In caso contrario, puoi ottenere questa autorizzazione tramite il ruolo Amministratore Service Usage (roles/serviceusage.serviceUsageAdmin). Scopri come concedere i ruoli. 

    gcloud services enable compute.googleapis.com

  9. Installa Google Cloud CLI.

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

  11. Per inizializzare gcloud CLI, esegui questo comando: 

    gcloud init

  12. 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 su cui ti è stato concesso un ruolo.
    • Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l' resourcemanager.projects.create autorizzazione. Scopri come concedere i ruoli.

    • Crea un Google Cloud progetto:

      gcloud projects create PROJECT_ID

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

    • Seleziona il Google Cloud progetto che hai creato:

      gcloud config set project PROJECT_ID

      Sostituisci PROJECT_ID con il nome del Google Cloud progetto.

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

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

  15. 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 hai già questa autorizzazione tramite il ruolo Proprietario (roles/owner). In caso contrario, puoi ottenere questa autorizzazione tramite il ruolo Amministratore Service Usage (roles/serviceusage.serviceUsageAdmin). Scopri come concedere i ruoli. 

    gcloud services enable compute.googleapis.com
    16.

Ruoli obbligatori

Per ottenere le autorizzazioni che ti servono per creare un MIG che forma una slice TPU multi-host, connetterti a ogni VM nel MIG utilizzando SSH ed eseguire comandi, chiedi all'amministratore di concederti i seguenti ruoli IAM sul tuo progetto:

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 template di istanza per le VM TPU v6e, utilizza il gcloud compute instance-templates create comando:

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 le slice TPU multi-host interconnesse.

Per creare una policy del workload per una slice TPU multi-host, utilizza il gcloud compute resource-policies create workload-policy comando 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 i seguenti comandi per creare un MIG che forma una slice TPU multi-host.

  1. Per creare un MIG che forma una slice TPU multi-host, utilizza il gcloud compute instance-groups managed create comando:

    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. Verifica che le istanze gestite siano in esecuzione utilizzando i seguenti comandi:

Installa JAX

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

  1. Controlla la 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 slice

Per eseguire il codice JAX su una slice TPU, devi eseguire il codice su ogni host nella slice TPU. La chiamata di funzione jax.device_count() smette di rispondere finché non viene chiamata su ogni host nella slice. Il seguente esempio mostra come eseguire un calcolo JAX su una slice 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 slice

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

Libera spazio

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

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

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

Passaggi successivi