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

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

Sfondo

Saxml è un sistema sperimentale che gestisce i framework Paxml, JAX e PyTorch. Puoi utilizzare le TPU per accelerare l'elaborazione dei dati con questi framework. Per dimostrare il deployment delle TPU in GKE, questo tutorial utilizza il modello di test LmCloudSpmd175B32Test da 175 miliardi di parametri. GKE esegue il deployment di questo modello di test su due pool di nodi di sezioni TPU v5e con topologia 4x8 rispettivamente.

Per eseguire il deployment corretto del modello di test, la topologia TPU è stata definita in base alle dimensioni del modello. Dato che il modello da N miliardi di bit richiede circa 2 volte (2xN) GB di memoria, il modello LmCloudSpmd175B32Test da 175 miliardi di bit 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 esegui il deployment delle TPU in GKE. Per saperne di più, consulta la sezione Pianificare la configurazione della TPU.

Obiettivi

Questo tutorial è rivolto a ingegneri MLOps o DevOps o amministratori di piattaforme che vogliono utilizzare le funzionalità di orchestrazione di GKE per pubblicare modelli di dati.

Questo tutorial illustra i seguenti passaggi:

  1. Prepara l'ambiente con un cluster standard. Il cluster ha due node pool di 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 gestire l'LLM.

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

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.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

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

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

    •

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

    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 column 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 dell'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 ruolo successivo.
    7. Fai clic su Salva.

      8. Prepara l'ambiente

      1. Nella console Google 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 CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
  export BUCKET_NAME=PROJECT_ID-gke-bucket

        Sostituisci i seguenti valori:

        • PROJECT_ID: il tuo Google Cloud ID progetto.
        • CONTROL_PLANE_LOCATION: la zona di Compute Engine del control plane del cluster. Seleziona la zona in cui è disponibile ct5lp-hightpu-4t.

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

      Crea un cluster

      Utilizza Cloud Shell per svolgere le seguenti operazioni:

      1. Crea un cluster Standard che utilizza Workload Identity Federation for GKE:

        gcloud container clusters create saxml \
    --location=${CONTROL_PLANE_LOCATION} \
    --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 \
    --location=${CONTROL_PLANE_LOCATION} \
    --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 \
    --location=${CONTROL_PLANE_LOCATION} \
    --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 4 chip TPU ciascuno.

      Il modello da 175 miliardi deve essere pubblicato su uno slice TPU v5e multi-host con uno 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 --location=${CONTROL_PLANE_LOCATION}

      2. Crea un service account 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 tuo account di servizio IAM per 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. Consenti al service account Kubernetes di impersonare il service account IAM aggiungendo un'associazione dei criteri IAM tra i due service account. Questa associazione consente al service account Kubernetes 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}.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 Google Cloud . Pertanto, quando l'app utilizza librerie 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

      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 eseguiti 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. Un Jobset è un'API per i carichi di lavoro che ti consente di gestire un gruppo di job Kubernetes come unità. Il caso d'uso più comune per un JobSet è l'addestramento distribuito, ma puoi utilizzarlo anche per eseguire workload 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 di JobSet. 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 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 nodil.
        • 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, sono presenti 16 container 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 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 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"
}

      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 --location ${CONTROL_PLANE_LOCATION}

      2. Elimina il 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://${BUCKET_NAME}

      Passaggi successivi