Gestisci un LLM utilizzando TPU multi-host su GKE con Saxml

Standard

---

Questo tutorial mostra come gestire un modello linguistico di grandi dimensioni (LLM) utilizzando Tensor Processing Unit (TPU) su Google Kubernetes Engine (GKE) con Saxml.

Contesto

Saxml è un sistema sperimentale che gestisce i framework Paxml, JAX e PyTorch. Puoi usare le TPU per accelerare l'elaborazione dei dati con questi framework Per eseguire la demo del deployment delle TPU in GKE, questo tutorial serve il modello di test 175B LmCloudSpmd175B32Test. GKE esegue il deployment di questo modello di test su due pool di nodi TPU v5e con topologia 4x8 rispettivamente.

Per eseguire correttamente il deployment del modello di test, la topologia TPU è stata definita in base alle dimensioni del modello. Dato che il modello N miliardo a 16 bit richiede circa 2 volte (2xN) GB di memoria, il modello 175B LmCloudSpmd175B32Test richiede circa 350 GB di memoria. Il chip singolo TPU v5e ha 16 GB. Per supportare 350 GB, GKE ha bisogno di 21 chip v5e (350/16= 21). In base alla mappatura della configurazione TPU, la configurazione TPU corretta per questo tutorial è:

• Tipo di macchina: ct5lp-hightpu-4t
• Topologia: 4x8 (32 numero di chip TPU)

La selezione della topologia TPU giusta per la gestione di un modello è importante quando esegui il deployment di TPU in GKE. Per scoprire di più, consulta Pianificare la configurazione TPU.

Obiettivi

Questo tutorial è destinato agli ingegneri MLOps o DevOps oppure agli amministratori di piattaforma che vogliono utilizzare le funzionalità di orchestrazione di GKE per gestire modelli di dati.

Questo tutorial illustra i seguenti passaggi:

1. Prepara l'ambiente con un cluster GKE Standard. Il cluster ha due pool di nodi TPU v5e con topologia 4x8.
2. Esegui il deployment di Saxml. Saxml richiede un server di amministrazione, un gruppo di pod che fungono da server del modello, un server HTTP predefinito e un bilanciatore del carico.
3. Utilizza il Saxml per gestire l'LLM.

Il seguente diagramma mostra l'architettura implementata dal seguente tutorial:

Figura: architettura di esempio di una TPU multi-host su GKE.

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.

  Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

  Go to project selector

• Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.

• Attiva l'API richiesta.

  Abilita l'API

• In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

  Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

  Go to project selector

• Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.

• Attiva l'API richiesta.

  Abilita l'API

• Assicurati di disporre dei seguenti ruoli nel progetto: roles/container.admin, roles/iam.serviceAccountAdmin

  Verifica i ruoli

  1. Nella console Google Cloud, vai alla pagina IAM.

     Vai a IAM
  2. Seleziona il progetto.

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

  4. Nella colonna Ruolo per la riga contenente il tuo indirizzo email, controlla se l'elenco dei ruoli include quelli richiesti.

  Concedi i ruoli

  1. Nella console Google Cloud, vai alla pagina IAM.

     Vai a IAM
  2. Seleziona il progetto.
  3. Fai clic su Concedi l'accesso.
  4. Nel campo Nuove entità, inserisci il tuo indirizzo email.
  5. Nell'elenco Seleziona un ruolo, scegli un ruolo.
  6. Per concedere altri ruoli, fai clic su Aggiungi un altro ruolo e aggiungi ciascun ruolo aggiuntivo.
  7. Fai clic su Salva.

• Assicurati che il progetto abbia una quota sufficiente per Cloud TPU in GKE.

prepara l'ambiente

1. Nella console Google Cloud, avvia un'istanza di Cloud Shell:
   Apri Cloud Shell

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

1. 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 di GKE. GKE supporta TPU v5e nella versione 1.27.2-gke.2100 e successive. Per maggiori informazioni, consulta la pagina relativa alla disponibilità delle TPU in GKE.

   La creazione del cluster potrebbe richiedere diversi minuti.

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

3. 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 seguenti risorse:

• Un cluster standard con quattro nodi CPU.
• Due pool di nodi TPU v5e con topologia 4x8. Ogni pool di nodi rappresenta otto nodi TPU con 4 chip ciascuno.

Il modello 175B deve essere fornito su una sezione TPU v5e multi-host con almeno una sezione di topologia 4x8 (32 chip TPU v5e).

crea un bucket Cloud Storage

Creare un bucket Cloud Storage per archiviare le configurazioni dei server di amministratore Saxml. Un server di amministrazione in esecuzione salva periodicamente lo 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 di Workload Identity per GKE

Assegna un account di servizio Kubernetes all'applicazione e configuralo in modo che funzioni come account di servizio IAM.

1. Configura kubectl per comunicare con il cluster:

   gcloud container clusters get-credentials saxml --zone=${ZONE}
   

2. Crea un ServiceAccount Kubernetes per l'applicazione da utilizzare:

   kubectl create serviceaccount sax-sa --namespace default
   

3. Crea un account di servizio IAM per l'applicazione:

   gcloud iam service-accounts create sax-iam-sa
   

4. Aggiungi un'associazione dei criteri IAM per consentire al tuo account di servizio IAM di 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
   

5. Consentire all'Account di servizio Kubernetes di assumere l'identità dell'account di servizio IAM aggiungendo un'associazione di criteri IAM tra i due account di servizio. Questa associazione consente a Kubernetes ServiceAccount di agire come account di servizio IAM, in modo che l'account di servizio Kubernetes possa leggere e scrivere in 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]"
   

6. Annota l'account di servizio Kubernetes con l'indirizzo email dell'account di servizio IAM. In questo modo, l'app di esempio può sapere quale account di servizio utilizzare per accedere ai servizi Google Cloud. Pertanto, quando l'app utilizza una qualsiasi libreria client standard delle API di Google per accedere ai servizi Google Cloud, utilizza quell'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 dei modelli Saxml.

Deployment del server amministratore Saxml

1. Crea il seguente manifest sax-admin-server.yaml:

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: sax-admin-server
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: sax-admin-server
     template:
       metadata:
         labels:
           app: sax-admin-server
       spec:
         hostNetwork: false
         serviceAccountName: sax-sa
         containers:
         - name: sax-admin-server
           image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.1.0
           securityContext:
             privileged: true
           ports:
           - containerPort: 10000
           env:
           - name: GSBUCKET
             value: BUCKET_NAME

   Sostituisci BUCKET_NAME con il nome del tuo bucket Cloud Storage.

2. Applica il manifest:

   kubectl apply -f sax-admin-server.yaml
   

3. Verifica che il pod del server dell'amministratore sia 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
   

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, usa IndexedJob, StatefulSet con un servizio headless o JobSet che crea automaticamente un servizio headless per tutti i job appartenenti al JobSet. La seguente sezione mostra come gestire più gruppi di pod di server dei modelli con JobSet.

1. 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 JobSet. Ad esempio, v0.2.3.

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

3. Esegui il deployment di due server del modello in due pool di nodi TPU. Salva il seguente manifest di sax-model-server-set:

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: sax-model-server-set
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: sax-model-server
         replicas: 2
         template:
           spec:
             parallelism: 8
             completions: 8
             backoffLimit: 0
             template:
               spec:
                 serviceAccountName: sax-sa
                 hostNetwork: true
                 dnsPolicy: ClusterFirstWithHostNet
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 4x8
                 containers:
                 - name: sax-model-server
                   image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0
                   args: ["--port=10001","--sax_cell=/sax/test", "--platform_chip=tpuv5e"]
                   ports:
                   - containerPort: 10001
                   - containerPort: 8471
                   securityContext:
                     privileged: true
                   env:
                   - name: SAX_ROOT
                     value: "gs://BUCKET_NAME/sax-root"
                   - name: MEGASCALE_NUM_SLICES
                     value: ""
                   resources:
                     requests:
                       google.com/tpu: 4
                     limits:
                       google.com/tpu: 4

   Sostituisci BUCKET_NAME con il nome del tuo bucket Cloud Storage.

   In questo file manifest:

   • replicas: 2 è il numero di repliche del job. Ogni job rappresenta un server del modello. Pertanto, un gruppo di 8 pod.
   • parallelism: 8 e completions: 8 sono uguali al numero di nodi in ogni 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 la comunicazione delle VM TPU
   • name: MEGASCALE_NUM_SLICES annulla la configurazione della variabile di ambiente perché GKE non esegue l'addestramento multislice.

4. Applica il manifest:

   kubectl apply -f sax-model-server-set.yaml
   

5. Verifica lo stato dei pod di server di amministrazione e di server del modello Saxml:

   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 dei modelli: 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 dei modelli in ogni gruppo.

Il cluster Saxml ha due server dei modelli di cui è stato eseguito il deployment in due pool di nodi TPU v5e con topologia 4x8 rispettivamente.

esegui il deployment del server HTTP Saxml e del bilanciatore del carico

1. Utilizza la seguente immagine predefinita del server HTTP dell'immagine. Salva il seguente manifest di sax-http.yaml:

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: sax-http
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: sax-http
     template:
       metadata:
         labels:
           app: sax-http
       spec:
         hostNetwork: false
         serviceAccountName: sax-sa
         containers:
         - name: sax-http
           image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.0.0
           ports:
           - containerPort: 8888
           env:
           - name: SAX_ROOT
             value: "gs://BUCKET_NAME/sax-root"
   ---
   apiVersion: v1
   kind: Service
   metadata:
     name: sax-http-lb
   spec:
     selector:
       app: sax-http
     ports:
     - protocol: TCP
       port: 8888
       targetPort: 8888
     type: LoadBalancer

   Sostituisci BUCKET_NAME con il nome del tuo bucket Cloud Storage.

2. Applica il manifest sax-http.yaml:

   kubectl apply -f sax-http.yaml
   

   Nota: puoi utilizzare il tuo server HTTP creato per Saxml. Per scoprire di più, consulta Inferenze con Saxml e un server HTTP.

3. Attendi che il contenitore del server HTTP termini la creazione:

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

4. Attendi che al servizio sia 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, esegui il deployment e gestisci il modello sul Saxml nella sezione multihost v5e TPU:

Carica il modello

1. 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"
   

2. Carica il modello di test LmCloudSpmd175B in due pool di nodi 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 perfezionato; le ponderazioni vengono generate 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
   }
   

3. 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 è completamente caricato.

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

Pubblica 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 di risposta del modello. Questa risposta potrebbe non essere significativa perché il modello di test ha ponderazioni casuali.

Operazione riuscita: hai gestito un LLM utilizzando TPU multi-host su GKE.

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

1. Elimina il cluster che hai creato per questo tutorial:

   gcloud container clusters delete saxml --zone ${ZONE}
   

2. Elimina l'account di servizio:

   gcloud iam service-accounts delete sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com
   

3. Elimina il bucket Cloud Storage:

   gcloud storage rm -r gs://${GSBUCKET}
   

Passaggi successivi

• Scopri le attuali versioni di TPU con l'architettura di sistema di Cloud TPU.
• Scopri di più sulle TPU in GKE.