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

Prepara l'ambiente

1. Nella console Trusted Cloud , avvia un'istanza 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 ZONE=COMPUTE_ZONE
     export BUCKET_NAME=PROJECT_ID-gke-bucket
   

   Sostituisci i seguenti valori:

   • PROJECT_ID: il tuo Trusted Cloud ID progetto.
   • COMPUTE_ZONE: la zona in cui è disponibile ct5lp-hightpu-4t.

   In questo comando, BUCKET_NAME specifica il nome del Trusted Cloud bucket di archiviazione in cui memorizzare le configurazioni del server amministratore Saxml.

Crea un cluster GKE Standard

Utilizza Cloud Shell per svolgere le seguenti operazioni:

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 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} \
       --machine-type=ct5lp-hightpu-4t \
       --tpu-topology=4x8 \
       --num-nodes=8 \
       --cluster=saxml
   

   Il valore del flag --num-nodes viene calcolato dividendo la topologia TPU per il numero di chip TPU per sezione TPU. In questo caso: (4 * 8) / 4.

3. Crea il secondo pool di nodi denominato tpu2:

   gcloud container node-pools create tpu2 \
       --zone=${ZONE} \
       --machine-type=ct5lp-hightpu-4t \
       --tpu-topology=4x8 \
       --num-nodes=8 \
       --cluster=saxml
   

   Il valore del flag --num-nodes viene calcolato dividendo la topologia TPU per il numero di chip TPU per sezione TPU. In questo caso: (4 * 8) / 4.

Hai creato le seguenti risorse:

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

Il modello 175B deve essere pubblicato su uno slice TPU v5e multi-host con slice di topologia 4x8 (32 chip TPU v5e) come minimo.

Crea un bucket Cloud Storage

Crea un bucket Cloud Storage per 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://${BUCKET_NAME}

Configura l'accesso ai tuoi workload utilizzando Workload Identity Federation for GKE

Assegna un service account Kubernetes all'applicazione e configura questo service account Kubernetes in modo che funga da 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 da utilizzare per la tua applicazione:

   kubectl create serviceaccount sax-sa --namespace default
   

3. Crea un account di servizio IAM per la tua applicazione:

   gcloud iam service-accounts create sax-iam-sa
   

4. Aggiungi un binding della policy IAM per il tuoaccount di serviziot IAM per leggere e scrivere in Cloud Storage:

   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
     --member "serviceAccount:sax-iam-sa@${PROJECT_ID}.s3ns-system.iam.gserviceaccount.com" \
     --role roles/storage.admin
   

5. Consenti al service account Kubernetes di rappresentare l'account di servizio IAM aggiungendo un'associazione dei criteri IAM tra i due service account. Questa associazione consente a Kubernetes ServiceAccount di agire come account di servizio IAM, in modo che possa leggere e scrivere in Cloud Storage.

   gcloud iam service-accounts add-iam-policy-binding sax-iam-sa@${PROJECT_ID}.s3ns-system.iam.gserviceaccount.com \
     --role roles/iam.workloadIdentityUser \
     --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/sax-sa]"
   

6. Annota il account di servizio Kubernetes con l'indirizzo email del account di servizio IAM. In questo modo, l'app di esempio sa quale account di servizio utilizzare per accedere ai servizi Trusted Cloud . Pertanto, quando l'app utilizza librerie client API di Google standard per accedere ai servizi Trusted Cloud , utilizza questo account di servizio IAM.

   kubectl annotate serviceaccount sax-sa \
     iam.gke.io/gcp-service-account=sax-iam-sa@${PROJECT_ID}.s3ns-system.iam.gserviceaccount.com
   

Esegui il deployment di Saxml

In questa sezione, esegui il deployment del server amministratore Saxml e del server del modello Saxml.

Esegui il 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

2. Sostituisci BUCKET_NAME con il bucket Cloud Storage creato in precedenza:

   perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-admin-server.yaml
   

3. Applica il manifest:

   kubectl apply -f sax-admin-server.yaml
   

4. Verifica che il pod del server 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
   

Esegui il deployment del server di modelli Saxml

I carichi di lavoro in esecuzione in slice TPU multi-host richiedono un identificatore di rete stabile per ogni pod per rilevare i peer nello stesso slice 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. Un Jobset è un'API per i carichi di lavoro che ti consente di gestire un gruppo di job Kubernetes come un'unica unità. Il caso d'uso più comune per un JobSet è l'addestramento distribuito, ma puoi utilizzarlo anche per eseguire carichi di lavoro batch.

La seguente sezione mostra come gestire più gruppi di pod del server di 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 del JobSet. Ad esempio, v0.2.3.

2. Convalida che il controller JobSet sia 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 di modelli in due node pool di sezioni TPU. Salva il seguente manifest 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

4. Sostituisci BUCKET_NAME con il bucket Cloud Storage creato in precedenza:

   perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-model-server-set.yaml
   

   In questo manifest:

   • replicas: 2 è il numero di repliche del job. Ogni job rappresenta un server di modelli. 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 zero per contrassegnare il job come non riuscito se un pod non va a buon fine.
   • ports.containerPort: 8471 è la porta predefinita per la comunicazione delle VM
   • name: MEGASCALE_NUM_SLICES annulla l'impostazione della variabile di ambiente perché GKE non esegue l'addestramento multislice.

5. Applica il manifest:

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

6. Verifica lo stato dei pod Saxml Admin Server e Model 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-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, ci sono 16 contenitori del server di 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 di modelli principali in ogni gruppo.

Il tuo cluster Saxml ha due server di modelli di cui è stato eseguito il deployment su due node pool di sezioni TPU v5e con topologia 4x8 rispettivamente.

Deployment del server HTTP Saxml e del bilanciatore del carico

1. Utilizza la seguente immagine del server HTTP predefinita. Salva il seguente manifest 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

2. Sostituisci BUCKET_NAME con il bucket Cloud Storage creato in precedenza:

   perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-http.yaml
   

3. Applica il manifest sax-http.yaml:

   kubectl apply -f sax-http.yaml
   

4. Attendi il completamento della creazione del contenitore del server HTTP:

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

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

Utilizzare Saxml

Carica, esegui il deployment e gestisci il modello su Saxml nello slice multihost TPU v5e:

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 node pool di slice 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
   }
   

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

Invia 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 pesi casuali.

Annullare 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"
}