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


Questo tutorial mostra come eseguire il deployment e gestire un modello linguistico di grandi dimensioni (LLM) utilizzando un pool di nodi della sezione TPU multi-host Google Kubernetes Engine (GKE) con Saxml per un'architettura efficiente e scalabile.

Sfondo

Saxml è un sistema sperimentale che gestisce Paxml JAX e framework PyTorch. Puoi utilizzare le TPU per accelerare l'elaborazione dei dati con questi framework. Per dimostrare il 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 delle sezioni TPU v5e rispettivamente con 4x8 topologia.

Per eseguire correttamente il deployment del modello di test, la topologia TPU è stata definita in base alle dimensioni del modello. Dato che il modello di N-miliardi a 16 bit approssimativamente richiede circa il 2 volte (2xN) GB di memoria, LmCloudSpmd175B32Test richiede circa 350 GB di memoria. Il singolo chip TPU v5e ha 16 GB. Per supportare 350 GB, GKE ha bisogno di 21 chip TPU 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 corretta per l'erogazione di un modello è importante quando si eseguono il deployment delle TPU in GKE. Per saperne di più, vedi Pianifica la configurazione di TPU.

Obiettivi

Questo tutorial è destinato agli MLOps o DevOps Engineer o agli amministratori di piattaforma che vogliono utilizzare funzionalità di orchestrazione GKE per gestire i dati di grandi dimensioni.

Questo tutorial illustra i seguenti passaggi:

  1. Prepara l'ambiente con un cluster GKE Standard. La ha due pool di nodi delle sezioni TPU v5e con topologia 4x8.
  2. Esegui il deployment di Saxml. Saxml ha bisogno di un server amministratore, un gruppo di pod che funzionano come server del modello, un server HTTP predefinito e un bilanciatore del carico.
  3. Utilizza Saxml per pubblicare l'LLM.

Il seguente diagramma mostra l'architettura implementata nel tutorial che segue:

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

Prima di iniziare

  • Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

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

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

  • Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin

    Check for the roles

    1. In the Google Cloud console, go to the IAM page.

      Go to IAM
    2. Select the project.
    3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

    4. For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.

    Grant the roles

    1. In the Google Cloud console, go to the IAM page.

      Vai a IAM
    2. Seleziona il progetto.
    3. Fai clic su Concedi l'accesso.
    4. Nel campo Nuove entità, inserisci il tuo identificatore utente. In genere si tratta dell'indirizzo email di un Account Google.

    5. Nell'elenco Seleziona un ruolo, seleziona un ruolo.
    6. Per concedere altri ruoli, fai clic su Aggiungi un altro ruolo e aggiungi ogni altro ruolo.
    7. Fai clic su Salva.

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:

Crea un cluster GKE Standard

Utilizza Cloud Shell per:

  1. Crea un cluster Standard che utilizza Federazione delle identità dei 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 della versione GKE. GKE supporta TPU v5e nella versione 1.27.2-gke.2100 e versioni successive. Per ulteriori informazioni, vedi 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 risorse seguenti:

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

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

Crea un bucket Cloud Storage

Crea un bucket Cloud Storage per archiviare le configurazioni del server di amministrazione Saxml. Un server amministratore in esecuzione salva periodicamente il proprio stato i dettagli dei modelli pubblicati.

In Cloud Shell, esegui quanto segue:

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 all'applicazione un ServiceAccount Kubernetes e configurarlo Kubernetes ServiceAccount per agire come account di servizio IAM.

  1. Configura kubectl in modo che comunichi con il tuo cluster:

    gcloud container clusters get-credentials saxml --zone=${ZONE}
    
  2. Crea un account di servizio 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'associazione dei criteri IAM per il tuo account di servizio IAM lettura e scrittura 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. Consenti all'account di servizio Kubernetes di impersonare l'account di servizio IAM aggiungendo un account IAM l'associazione di criteri tra i due account di servizio. Questa associazione consente a Kubernetes ServiceAccount affinché agisca come account di servizio IAM, in modo che Kubernetes ServiceAccount può 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]"
    
  6. Annota l'account di servizio Kubernetes con l'indirizzo email l'account di servizio. In questo modo, l'app di esempio saprà quale account di servizio utilizzare per accedere ai servizi Google Cloud. Pertanto, quando l'app utilizza qualsiasi libreria client API di Google standard per accedere ai servizi Google Cloud, utilizza questo 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 di amministrazione Saxml e del modello Saxml server web.

Esegui il deployment del server di amministrazione 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 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 scoprire i peer nella stessa sezione TPU. Per definire questi identificatori, utilizza IndexedJob, StatefulSet con un servizio headless o un JobSet che crea automaticamente un servizio headless tutti i job appartenenti a un oggetto JobSet. La sezione seguente mostra come gestire più gruppi di pod di 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 set di job. Ad esempio, v0.2.3.

  2. Verifica 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 pool di nodi di sezioni TPU. Salva quanto segue 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 manifest:

    • replicas: 2 è il numero di repliche del job. Ogni job rappresenta un server di modelli. Di conseguenza, un gruppo di 8 pod.
    • parallelism: 8 e completions: 8 sono uguali al numero di nodi in ciascuno 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 VM
    • name: MEGASCALE_NUM_SLICES azzera la 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 del server di amministrazione Saxml e del server di modelli:

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

Il cluster Saxml ha due server del modello di cui è stato eseguito il deployment su due pool di nodi sezione TPU v5e con 4x8 topologia.

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

  1. Utilizza l'immagine del server HTTP predefinita riportata di seguito. Salva il seguente manifestsax-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 nome del bucket Cloud Storage.

  2. Applica il manifest sax-http.yaml:

    kubectl apply -f sax-http.yaml
    
  3. 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
    ...
    
  4. 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 distribuisci il modello su Saxml nella sezione 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 Pool di nodi della sezione 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, i pesi vengono generati. 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 è stato caricato completamente.

  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

Presenta 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 il seguente 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