Inferenza JetStream MaxText su VM TPU v6e

Questo tutorial mostra come utilizzare JetStream per pubblicare modelli MaxText su TPU v6e. JetStream è un motore ottimizzato per la velocità effettiva e la memoria per l'inferenza dei modelli linguistici di grandi dimensioni (LLM) sui dispositivi XLA (TPU). In questo tutorial esegui il benchmark di inferenza per il modello Llama2-7B.

Prima di iniziare

Preparati a eseguire il provisioning di una TPU v6e con 4 chip:

1. Segui la guida Configurare l'ambiente Cloud TPU per assicurarti di disporre dell'accesso appropriato per utilizzare le Cloud TPU.

2. Crea un'identità di servizio per la VM TPU.

   gcloud alpha compute tpus tpu-vm service-identity create --zone=zone

3. Crea un account di servizio TPU e concedi accesso ai servizi Google Cloud .

   Gli account di servizio consentono al Google Cloud servizio TPU di accedere ad altri Google Cloud servizi. È consigliato un account di servizio gestito dall'utente. Puoi creare un account di servizio dalla console Google Cloud o tramite il comandogcloud.

   Crea un account di servizio utilizzando lo strumento a riga di comando gcloud:

   gcloud iam service-accounts create your-service-account-name \
   --description="your-sa-description" \
   --display-name="your-sa-display-name"
   export SERVICE_ACCOUNT_NAME=your-service-account-name

   Crea un account di servizio dalla console Google Cloud:

   1. Vai alla pagina Account di servizio nella console Google Cloud.
   2. Fai clic su Crea account di servizio.
   3. Inserisci il nome dell'account di servizio.
   4. (Facoltativo) Inserisci una descrizione per l'account di servizio.
   5. Fai clic su Crea e continua.
   6. Scegli i ruoli da concedere all'account di servizio.
   7. Fai clic su Continua.
   8. (Facoltativo) Specifica gli utenti o i gruppi che possono gestire l'account di servizio.
   9. Fai clic su Fine per completare la creazione del account di servizio.

   Dopo aver creato l'account di servizio, segui questi passaggi per concedere i ruoli dell'account di servizio.

   Sono necessari i seguenti ruoli:

   • TPU Admin: necessario per creare una TPU
   • Amministratore Storage: necessario per accedere a Cloud Storage
   • Writer log
   • Monitoring Metric Writer: necessario per scrivere le metriche in Cloud Monitoring

   L'amministratore deve concederti il ruolo roles/resourcemanager.projectIamAdmin per consentirti di assegnare i ruoli IAM agli utenti. Anche un utente con il ruolo Amministratore IAM progetto roles/resourcemanager.projectIamAdmin può assegnare questo ruolo.

   Utilizza i seguenti comandi gcloud per aggiungere i ruoli dell'account di servizio:

   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/tpu.admin
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/storage.admin
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/logging.logWriter
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
      --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
      --role roles/monitoring.metricWriter

   Puoi assegnare i ruoli anche utilizzando la console Google Cloud.

   Nella console Google Cloud, seleziona i seguenti ruoli:

   1. Seleziona il tuo account di servizio e fai clic su Aggiungi principale.
   2. Nel campo Nuove entità, inserisci l'indirizzo email del tuo account di servizio.
   3. Nel menu a discesa Seleziona un ruolo, cerca il ruolo (ad esempio Amministratore Storage) e selezionalo.
   4. Fai clic su Salva.

4. Esegui l'autenticazione con Google Cloud e configura il progetto e la zona predefiniti per Google Cloud CLI.

   gcloud auth login
   gcloud config set project PROJECT_ID
   gcloud config set compute/zone ZONE

Capacità sicura

Quando è tutto pronto per assicurarti la capacità TPU, consulta la pagina delle quote per saperne di più sul sistema delle quote Cloud. Se hai altre domande su come assicurarti la capacità, contatta il team di vendita o dell'account Cloud TPU.

Esegui il provisioning dell'ambiente Cloud TPU

Puoi eseguire il provisioning delle VM TPU con GKE, con GKE e XPK, o come risorse in coda.

Prerequisiti

• Questo tutorial è stato testato con Python 3.10 o versioni successive.
• Verifica che il tuo progetto disponga di una quota TPUS_PER_TPU_FAMILY sufficiente, che specifica il numero massimo di chip a cui puoi accedere all'interno del progettoGoogle Cloud .
• Verifica che il tuo progetto disponga di una quota TPU sufficiente per:
  • Quota VM TPU
  • Quota di indirizzi IP
  • Quota Hyperdisk bilanciata
• Autorizzazioni del progetto per gli utenti
  • Se utilizzi GKE con XPK, consulta Autorizzazioni di Cloud Console per l'account utente o di servizio per conoscere le autorizzazioni necessarie per eseguire XPK.

Creare variabili di ambiente

In Cloud Shell, crea le seguenti variabili di ambiente:

export NODE_ID=TPU_NODE_ID # TPU name
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=v6e-4
export ZONE=us-central2-b
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=YOUR_SERVICE_ACCOUNT
export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID
export VALID_DURATION=VALID_DURATION

# Additional environment variable needed for Multislice:
export NUM_SLICES=NUM_SLICES

# Use a custom network for better performance as well as to avoid having the
# default network becoming overloaded.
export NETWORK_NAME=${PROJECT_ID}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

Esegui il provisioning di una TPU v6e

gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
    --node-id TPU_NAME \
    --project PROJECT_ID \
    --zone ZONE \
    --accelerator-type v6e-4 \
    --runtime-version v2-alpha-tpuv6e \
    --service-account SERVICE_ACCOUNT

Utilizza i comandi list o describe per eseguire query sullo stato della risorsa in coda.

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project ${PROJECT_ID} --zone ${ZONE}

Per un elenco completo degli stati delle richieste di risorse in coda, consulta la documentazione relativa alle risorse in coda.

Connettiti alla TPU tramite SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

Dopo aver eseguito la connessione alla TPU, puoi eseguire il benchmark di inferenza.

Esegui il benchmark di inferenza Llama2-7B

Per configurare JetStream e MaxText, convertire i checkpoint del modello ed eseguire il benchmark di inferenza, segui le istruzioni nel repository GitHub.

Al termine del benchmark di inferenza, assicurati di ripulire le risorse TPU.

Esegui la pulizia

Elimina la TPU:

gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \
    --project PROJECT_ID \
    --zone ZONE \
    --force \
    --async