Inferenza JetStream PyTorch su VM TPU v6e

Questo tutorial mostra come utilizzare JetStream per eseguire il servizio di modelli PyTorch su TPU v6e. JetStream è un motore ottimizzato per la velocità effettiva e la memoria per l'inferenza dei modelli linguistici di grandi dimensioni (LLM) su 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 configurare un Google Cloud progetto, configurare Google Cloud CLI, attivare l'API Cloud TPU e assicurarti di avere accesso all'utilizzo delle Cloud TPU.

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

Garantire la capacità

Quando è tutto pronto per garantire la capacità TPU, consulta Quote Cloud TPU per saperne di più sulle quote Cloud TPU. Se hai altre domande sulla prenotazione della capacità, contatta il team di vendita o di gestione 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

  • 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 bilanciato
  • Autorizzazioni del progetto per gli utenti

Creare variabili di ambiente

In Cloud Shell, crea le seguenti variabili di ambiente:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-central2-b
export ACCELERATOR_TYPE=v6e-4
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

Descrizioni delle variabili di ambiente

Variabile Descrizione
PROJECT_ID Il tuo Google Cloud ID progetto. Utilizza un progetto esistente o creane uno nuovo.
TPU_NAME Il nome della TPU.
ZONE La zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
ACCELERATOR_TYPE Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni TPU.
RUNTIME_VERSION La versione software di Cloud TPU.
SERVICE_ACCOUNT L'indirizzo email del tuo account di servizio. Puoi trovarlo nella pagina Account di servizio della console Google Cloud .

Ad esempio: tpu-service-account@PROJECT_ID.iam.gserviceaccount.com

QUEUED_RESOURCE_ID L'ID di testo assegnato dall'utente della richiesta di risorsa in coda.

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 ${ACCELERATOR_TYPE} \
        --runtime-version ${RUNTIME_VERSION} \
        --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}

Esegui il benchmark JetStream PyTorch Llama2-7B

Per configurare JetStream-PyTorch, 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