Questo tutorial mostra come eseguire il deployment e gestire un modello linguistico di grandi dimensioni (LLM) utilizzando un pool di nodi di sezioni TPU multi-host su Google Kubernetes Engine (GKE) con Saxml per un'architettura scalabile ed efficiente.
Contesto
Saxml è un sistema sperimentale che gestisce i framework
Paxml,
JAX e
PyTorch. Puoi usare le TPU per
accelerare l'elaborazione dati con questi framework. Per dimostrare il deployment delle TPU in GKE, questo tutorial utilizza il modello di test 175B LmCloudSpmd175B32Test. GKE esegue il deployment di questo modello di test, rispettivamente, su due pool di nodi della sezione TPU v5e con topologia 4x8
.
Per eseguire correttamente il deployment del modello di test, la topologia TPU è stata definita in base alle dimensioni del modello. Dato che il modello da N miliardi a 16 bit richiede circa 2 volte (2xN) GB di memoria, il modello LmCloudSpmd175B32Test da 175B richiede circa 350 GB di memoria. Il chip TPU v5e singolo ha 16 GB. Per supportare 350 GB, GKE ha bisogno di 21 chip TPU v5e (350/16= 21). In base alla mappatura della configurazione delle TPU, la configurazione TPU adeguata per questo tutorial è:
- Tipo di macchina:
ct5lp-hightpu-4t
- Topologia:
4x8
(32 numero di chip TPU)
La scelta della topologia TPU giusta per la gestione di un modello è importante durante il deployment di TPU in GKE. Per scoprire di più, consulta Pianificare la configurazione di TPU.
Obiettivi
Questo tutorial è destinato agli ingegneri MLOps o DevOps o agli amministratori di piattaforma che vogliono utilizzare funzionalità di orchestrazione GKE per gestire i modelli di dati.
Questo tutorial illustra i seguenti passaggi:
- Prepara l'ambiente con un cluster GKE Standard. Il cluster ha due pool di nodi delle sezioni TPU v5e con topologia
4x8
. - Esegui il deployment di Saxml. Saxml ha bisogno di un server amministratore, un gruppo di pod che funziona come server del modello, un server HTTP predefinito e un bilanciatore del carico.
- Utilizza il Saxml per gestire l'LLM.
Il seguente diagramma mostra l'architettura implementata dal seguente tutorial:
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 gratuiti per l'esecuzione, il test e il deployment dei carichi di lavoro.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.
-
Attiva l'API richiesta.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.
-
Attiva l'API richiesta.
-
Assicurati di disporre dei seguenti ruoli nel progetto: roles/container.admin, roles/iam.serviceAccountAdmin
Verifica i ruoli
-
Nella console Google Cloud, vai alla pagina IAM.
Vai a IAM - Seleziona il progetto.
-
Nella colonna Entità, individua la riga contenente il tuo indirizzo email.
Se il tuo indirizzo email non è in questa colonna, significa che non disponi di alcun ruolo.
- Nella colonna Ruolo per la riga contenente il tuo indirizzo email, controlla se l'elenco dei ruoli include quelli richiesti.
Concedi i ruoli
-
Nella console Google Cloud, vai alla pagina IAM.
Vai a IAM - Seleziona il progetto.
- Fai clic su Concedi l'accesso.
- Nel campo Nuove entità, inserisci il tuo indirizzo email.
- Nell'elenco Seleziona un ruolo, scegli un ruolo.
- Per concedere altri ruoli, fai clic su Aggiungi un altro ruolo e aggiungi ciascun ruolo aggiuntivo.
- Fai clic su Salva.
-
- Assicurati che il progetto abbia una quota sufficiente per Cloud TPU in GKE.
prepara l'ambiente
Nella console Google Cloud, avvia un'istanza di Cloud Shell:
Apri Cloud ShellImposta le variabili di ambiente predefinite:
gcloud config set project PROJECT_ID export PROJECT_ID=$(gcloud config get project) export REGION=COMPUTE_REGION export ZONE=COMPUTE_ZONE export GSBUCKET=PROJECT_ID-gke-bucket
Sostituisci i seguenti valori:
- PROJECT_ID: il tuo ID progetto Google Cloud.
- COMPUTE_REGION: la regione di Compute Engine.
- COMPUTE_ZONE: la zona in cui è disponibile
ct5lp-hightpu-4t
.
Crea un cluster GKE Standard
Utilizza Cloud Shell per:
Crea un cluster Standard che utilizza la Federazione delle identità per i carichi di lavoro per GKE:
gcloud container clusters create saxml \ --zone=${ZONE} \ --workload-pool=${PROJECT_ID}.svc.id.goog \ --cluster-version=VERSION \ --num-nodes=4
Sostituisci
VERSION
con il numero di versione GKE. GKE supporta TPU v5e nella versione 1.27.2-gke.2100 e successive. Per maggiori informazioni, consulta Disponibilità delle TPU in GKE.La creazione del cluster potrebbe richiedere diversi minuti.
Crea il primo pool di nodi denominato
tpu1
:gcloud container node-pools create tpu1 \ --zone=${ZONE} \ --num-nodes=8 \ --machine-type=ct5lp-hightpu-4t \ --tpu-topology=4x8 \ --cluster=saxml
Crea il secondo pool di nodi denominato
tpu2
:gcloud container node-pools create tpu2 \ --zone=${ZONE} \ --num-nodes=8 \ --machine-type=ct5lp-hightpu-4t \ --tpu-topology=4x8 \ --cluster=saxml
Hai creato le risorse seguenti:
- Un cluster Standard con quattro nodi CPU.
- Due pool di nodi di sezioni TPU v5e con topologia
4x8
. Ogni pool di nodi rappresenta otto nodi di sezione TPU con 4 chip TPU ciascuno.
Il modello 175B deve essere pubblicato su una sezione TPU v5e multi-host con una sezione di topologia 4x8
(32 chip TPU v5e) come minimo.
Crea un bucket Cloud Storage
Creare un bucket Cloud Storage in cui archiviare le configurazioni del server amministratore Saxml. Un server amministratore in esecuzione salva periodicamente il proprio stato e i dettagli dei modelli pubblicati.
In Cloud Shell, esegui questo comando:
gcloud storage buckets create gs://${GSBUCKET}
Configura l'accesso ai carichi di lavoro utilizzando la federazione delle identità per i carichi di lavoro per GKE
Assegnare un ServiceAccount Kubernetes all'applicazione e configurarlo in modo che agisca come account di servizio IAM.
Configura
kubectl
per comunicare con il tuo cluster:gcloud container clusters get-credentials saxml --zone=${ZONE}
Crea un ServiceAccount Kubernetes da utilizzare per la tua applicazione:
kubectl create serviceaccount sax-sa --namespace default
Crea un account di servizio IAM per la tua applicazione:
gcloud iam service-accounts create sax-iam-sa
Aggiungi un'associazione dei criteri IAM per il tuo account di servizio IAM in modo che possa leggere e scrivere in Cloud Storage:
gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member "serviceAccount:sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com" \ --role roles/storage.admin
Consenti all'account di servizio Kubernetes di impersonare l'account di servizio IAM aggiungendo un'associazione di criteri IAM tra i due account di servizio. Questa associazione consente all'account di servizio Kubernetes di agire come account di servizio IAM, in modo che l'account di servizio Kubernetes possa leggere e scrivere su Cloud Storage.
gcloud iam service-accounts add-iam-policy-binding sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/iam.workloadIdentityUser \ --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/sax-sa]"
Annota l'account di servizio Kubernetes con l'indirizzo email dell'account di servizio IAM. Ciò consente all'app di esempio di sapere quale account di servizio utilizzare per accedere ai servizi Google Cloud. Perciò, quando l'app utilizza librerie client delle API di Google standard per accedere ai servizi Google Cloud, utilizza l'account di servizio IAM.
kubectl annotate serviceaccount sax-sa \ iam.gke.io/gcp-service-account=sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com
Deployment di Saxml
In questa sezione eseguirai il deployment del server amministratore Saxml e del server del modello Saxml.
Distribuisci il server di amministrazione Saxml
Crea il seguente manifest
sax-admin-server.yaml
:Sostituisci
BUCKET_NAME
con il nome del tuo bucket Cloud Storage.Applica il manifest:
kubectl apply -f sax-admin-server.yaml
Verifica che il pod del server dell'amministratore sia attivo e in esecuzione:
kubectl get deployment
L'output è simile al seguente:
NAME READY UP-TO-DATE AVAILABLE AGE sax-admin-server 1/1 1 1 52s
Esegui il deployment del server del modello Saxml
I carichi di lavoro in esecuzione in sezioni TPU multi-host richiedono un identificatore di rete stabile per ogni pod per rilevare i peer nella stessa sezione TPU. Per definire questi identificatori, utilizza IndexedJob, StatefulSet con un servizio headless o JobSet, che crea automaticamente un servizio headless per tutti i job che appartengono al JobSet. La seguente sezione mostra come gestire più gruppi di pod del server modello con JobSet.
Installa JobSet v0.2.3 o versioni successive.
kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml
Sostituisci
JOBSET_VERSION
con la versione del set di job. Ad esempio,v0.2.3
.Convalida il controller JobSet in esecuzione nello spazio dei nomi
jobset-system
:kubectl get pod -n jobset-system
L'output è simile al seguente:
NAME READY STATUS RESTARTS AGE jobset-controller-manager-69449d86bc-hp5r6 2/2 Running 0 2m15s
Esegui il deployment di due server del modello in due pool di nodi delle sezioni TPU. Salva il seguente manifest di
sax-model-server-set
:Sostituisci
BUCKET_NAME
con il nome del tuo bucket Cloud Storage.In questo manifest:
replicas: 2
è il numero di repliche del job. Ogni job rappresenta un server del modello. Di conseguenza, un gruppo di 8 pod.parallelism: 8
ecompletions: 8
sono uguali al numero di nodi in ciascun pool di nodi.backoffLimit: 0
deve essere pari a zero per contrassegnare il job come non riuscito in caso di errore di un pod.ports.containerPort: 8471
è la porta predefinita per le comunicazioni delle VMname: MEGASCALE_NUM_SLICES
annulla la configurazione della variabile di ambiente perché GKE non sta eseguendo l'addestramento multisezione.
Applica il manifest:
kubectl apply -f sax-model-server-set.yaml
Verifica lo stato dei pod del server di amministrazione Saxml e del server del modello:
kubectl get pods
L'output è simile al seguente:
NAME READY STATUS RESTARTS AGE sax-admin-server-557c85f488-lnd5d 1/1 Running 0 35h sax-model-server-set-sax-model-server-0-0-nj4sm 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-1-sl8w4 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-2-hb4rk 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-3-qv67g 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-4-pzqz6 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-5-nm7mz 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-6-7br2x 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-7-4pw6z 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-0-8mlf5 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-1-h6z6w 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-2-jggtv 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-3-9v8kj 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-4-6vlb2 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-5-h689p 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-6-bgv5k 1/1 Running 0 24m sax-model-server-set-sax-model-server-1-7-cd6gv 1/1 Running 0 24m
In questo esempio sono presenti 16 container di server del modello: sax-model-server-set-sax-model-server-0-0-nj4sm
e sax-model-server-set-sax-model-server-1-0-8mlf5
sono i due server principali del modello in ogni gruppo.
Il cluster Saxml ha due server del modello di cui è stato eseguito il deployment su due pool di nodi di sezioni TPU v5e con
topologia 4x8
rispettivamente.
Esegui il deployment del server HTTP Saxml e del bilanciatore del carico
Utilizza la seguente immagine del server HTTP con immagine predefinita. Salva il seguente manifest di
sax-http.yaml
:Sostituisci
BUCKET_NAME
con il nome del tuo bucket Cloud Storage.Applica il manifest
sax-http.yaml
:kubectl apply -f sax-http.yaml
Attendi il completamento della creazione del container HTTP Server:
kubectl get pods
L'output è simile al seguente:
NAME READY STATUS RESTARTS AGE sax-admin-server-557c85f488-lnd5d 1/1 Running 0 35h sax-http-65d478d987-6q7zd 1/1 Running 0 24m sax-model-server-set-sax-model-server-0-0-nj4sm 1/1 Running 0 24m ...
Attendi che al servizio venga assegnato un indirizzo IP esterno:
kubectl get svc
L'output è simile al seguente:
NAME TYPE CLUSTER-IP EXTERNAL-IP PORT(S) AGE sax-http-lb LoadBalancer 10.48.11.80 10.182.0.87 8888:32674/TCP 7m36s
Utilizza Saxml
Carica, distribuisci e gestisci il modello su Saxml nella sezione multihost TPU v5e:
Carica il modello
Recupera l'indirizzo IP del bilanciatore del carico per Saxml.
LB_IP=$(kubectl get svc sax-http-lb -o jsonpath='{.status.loadBalancer.ingress[*].ip}') PORT="8888"
Carica il modello di test
LmCloudSpmd175B
in due pool di nodi delle sezioni TPU v5e:curl --request POST \ --header "Content-type: application/json" \ -s ${LB_IP}:${PORT}/publish --data \ '{ "model": "/sax/test/spmd", "model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test", "checkpoint": "None", "replicas": 2 }'
Il modello di test non ha un checkpoint ottimizzato, i pesi vengono generati in modo casuale. Il caricamento del modello potrebbe richiedere fino a 10 minuti.
L'output è simile al seguente:
{ "model": "/sax/test/spmd", "path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test", "checkpoint": "None", "replicas": 2 }
Controlla l'idoneità del modello:
kubectl logs sax-model-server-set-sax-model-server-0-0-nj4sm
L'output è simile al seguente:
... loading completed. Successfully loaded model for key: /sax/test/spmd
Il modello è stato caricato completamente.
Ottieni informazioni sul modello:
curl --request GET \ --header "Content-type: application/json" \ -s ${LB_IP}:${PORT}/listcell --data \ '{ "model": "/sax/test/spmd" }'
L'output è simile al seguente:
{ "model": "/sax/test/spmd", "model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test", "checkpoint": "None", "max_replicas": 2, "active_replicas": 2 }
Pubblica il modello
Elabora una richiesta di prompt:
curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/generate --data \
'{
"model": "/sax/test/spmd",
"query": "How many days are in a week?"
}'
L'output mostra un esempio della risposta del modello. Questa risposta potrebbe non essere significativa perché il modello di test ha ponderazioni casuali.
Annulla la pubblicazione del modello
Esegui questo comando per annullare la pubblicazione del modello:
curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/unpublish --data \
'{
"model": "/sax/test/spmd"
}'
L'output è simile al seguente:
{
"model": "/sax/test/spmd"
}
Esegui la pulizia
Per evitare che al tuo Account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.
Elimina le risorse di cui è stato eseguito il deployment
Elimina il cluster che hai creato per questo tutorial:
gcloud container clusters delete saxml --zone ${ZONE}
Elimina l'account di servizio:
gcloud iam service-accounts delete sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com
Elimina il bucket Cloud Storage:
gcloud storage rm -r gs://${GSBUCKET}
Passaggi successivi
- Scopri di più sulle versioni attuali di TPU con l'architettura di sistema Cloud TPU.
- Scopri di più sulle TPU in GKE.