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:
- Accedi al tuo Account Google. Se non l'hai ancora fatto, registrati per creare un nuovo account.
- Nella console Google Cloud, seleziona o crea un progetto Google Cloud dalla pagina di selezione del progetto.
- Abilita la fatturazione per il tuo progetto Google Cloud . La fatturazione è obbligatoria per tutto l'utilizzo di Google Cloud .
- Installa i componenti gcloud alpha.
Esegui il seguente comando per installare la versione più recente dei componenti
gcloud
.gcloud components update
Abilita l'API TPU tramite il seguente comando
gcloud
utilizzando Cloud Shell. Puoi anche attivarlo dalla console Google Cloud.gcloud services enable tpu.googleapis.com
Crea un'identità di servizio per la VM TPU.
gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE
Crea un account di servizio TPU e concedi l'accesso ai servizi Google Cloud .
I service account consentono al servizio TPU di Google Cloud di accedere ad altri servizi di Google Cloud. È consigliato un account di servizio gestito dall'utente. Segui queste guide per creare e concedere i ruoli. Sono necessari i seguenti ruoli:
- Amministratore TPU: necessario per creare una TPU
- Amministratore dello spazio di archiviazione: necessario per accedere a Cloud Storage
- Scrittore di log: necessario per scrivere i log con l'API Logging
- Monitoring Metric Writer: necessario per scrivere le metriche in Cloud Monitoring
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à
Contatta il team di vendita o dedicato all'account Cloud TPU per richiedere una quota TPU e per eventuali domande sulla capacità.
Esegui il provisioning dell'ambiente Cloud TPU
Puoi eseguire il provisioning delle TPU v6e 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 nel tuo progetto Google Cloud . - Questo tutorial è stato testato con la seguente configurazione:
- Python
3.10 or later
- Versioni software Nightly:
- a notte JAX
0.4.32.dev20240912
- nightly LibTPU
0.1.dev20240912+nightly
- a notte JAX
- Versioni software stabili:
- JAX + libreria JAX di
v0.4.35
- JAX + libreria JAX di
- Python
- 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
Descrizioni dei flag dei comandi
Variabile | Descrizione |
NODE_ID | L'ID assegnato dall'utente della TPU creata quando viene allocata la richiesta di risorsa in coda. |
PROJECT_ID | Nome del progettoGoogle Cloud . Utilizza un progetto esistente o creane uno nuovo. |
ZONA | Consulta il documento Regioni e zone TPU per le zone supportate. |
ACCELERATOR_TYPE | Consulta la documentazione relativa ai tipi di acceleratore per conoscere i tipi di acceleratore supportati. |
RUNTIME_VERSION | v2-alpha-tpuv6e
|
SERVICE_ACCOUNT | Si tratta dell'indirizzo email del tuo account di servizio che puoi trovare in
console Google Cloud -> IAM -> Account di servizio
Ad esempio: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com |
NUM_SLICES | Il numero di sezioni da creare (necessario solo per Multislice) |
QUEUED_RESOURCE_ID | L'ID testo assegnato dall'utente della richiesta di risorsa in coda. |
VALID_DURATION | La durata di validità della richiesta di risorse in coda. |
NETWORK_NAME | Il nome di una rete secondaria da utilizzare. |
NETWORK_FW_NAME | Il nome di un firewall di rete secondario da utilizzare. |
Provisiona 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
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