LLM mit Multi-Host-TPUs in GKE mit Saxml bereitstellen

In dieser Anleitung erfahren Sie, wie Sie mithilfe von Saxml in Google Kubernetes Engine (GKE) ein großes Sprachmodell (LLM) mit Tensor Processing Units (TPUs) bereitstellen.

Hintergrund

Saxml ist ein experimentelles System, das Paxml-, JAX- und PyTorch-Frameworks bereitstellt. Sie können TPUs verwenden, um die Datenverarbeitung mit diesen Frameworks zu beschleunigen. In dieser Anleitung wird das 175B-Testmodell LmCloudSpmd175B32Test gezeigt, um die Bereitstellung von TPUs in GKE zu demonstrieren. GKE stellt dieses Testmodell in zwei v5e TPU-Knotenpools mit der 4x8-Topologie bereit.

Um das Testmodell ordnungsgemäß bereitzustellen, wurde die TPU-Topologie anhand der Größe des Modells definiert. Da das N Milliarden 16-Bit-Modell ungefähr zweimal (2xN) GB Arbeitsspeicher erfordert, benötigt das 175B-LmCloudSpmd175B32Test-Modell etwa 350 GB Arbeitsspeicher. Der einzelne Chip TPU v5e hat 16 GB. Zur Unterstützung von 350 GB benötigt GKE 21 v5e-Chips (350/16= 21). Basierend auf der Zuordnung der TPU-Konfiguration lautet die TPU-Konfiguration für diese Anleitung ordnungsgemäß:

  • Maschinentyp: ct5lp-hightpu-4t
  • Topologie: 4x8 (32 TPU-Chips)

Die Wahl der richtigen TPU-Topologie zum Bereitstellen eines Modells ist wichtig, wenn TPUs in GKE bereitgestellt werden. Weitere Informationen finden Sie unter TPU-Konfiguration planen.

Lernziele

Diese Anleitung richtet sich an MLOps- oder DevOps-Entwickler oder Plattformadministratoren, die GKE-Orchestrierungsfunktionen zum Bereitstellen von Datenmodellen verwenden möchten.

Diese Anleitung umfasst die folgenden Schritte:

  1. Umgebung mit einem GKE-Standardcluster vorbereiten Der Cluster hat zwei v5e-TPU-Knotenpools mit 4x8-Topologie.
  2. Saxml bereitstellen Saxml benötigt einen Administratorserver, eine Gruppe von Pods, die als Modellserver fungieren, einen vordefinierten HTTP-Server und einen Load-Balancer.
  3. Verwenden Sie das Saxml, um das LLM bereitzustellen.

Das folgende Diagramm zeigt die Architektur, die in der folgenden Anleitung implementiert wird:

Architektur einer TPU mit mehreren Hosts in GKE
Abbildung: Beispielarchitektur einer TPU mit mehreren Hosts in GKE.

Hinweise

  • Melden Sie sich bei Ihrem Google Cloud-Konto an. Wenn Sie mit Google Cloud noch nicht vertraut sind, erstellen Sie ein Konto, um die Leistungsfähigkeit unserer Produkte in der Praxis sehen und bewerten zu können. Neukunden erhalten außerdem ein Guthaben von 300 $, um Arbeitslasten auszuführen, zu testen und bereitzustellen.

  • Wählen Sie in der Google Cloud Console auf der Seite der Projektauswahl ein Google Cloud-Projekt aus oder erstellen Sie eines.

    Zur Projektauswahl

  • Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein.

  • Aktivieren Sie die erforderliche API.

    Aktivieren Sie die API

    •

  • Wählen Sie in der Google Cloud Console auf der Seite der Projektauswahl ein Google Cloud-Projekt aus oder erstellen Sie eines.

    Zur Projektauswahl

  • Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein.

  • Aktivieren Sie die erforderliche API.

    Aktivieren Sie die API

    •

  • Prüfen Sie, ob Sie die folgenden Rollen für das Projekt haben: roles/container.admin, roles/iam.serviceAccountAdmin

    Auf Rollen prüfen

    1. Öffnen Sie in der Google Cloud Console die Seite IAM.

      IAM aufrufen
    2. Wählen Sie das Projekt aus.

    3. Suchen Sie in der Spalte Hauptkonto die Zeile mit Ihrer E-Mail-Adresse.

      Ist Ihre E-Mail-Adresse nicht in dieser Spalte enthalten, haben Sie keine Rollen.

    4. Prüfen Sie in der Spalte Rolle der Zeile mit Ihrer E-Mail-Adresse, ob die Liste der Rollen die erforderlichen Rollen enthält.

    Rollen zuweisen

    1. Öffnen Sie in der Google Cloud Console die Seite IAM.

      IAM aufrufen
    2. Wählen Sie das Projekt aus.
    3. Klicken Sie auf Zugriff erlauben.
    4. Geben Sie in das Feld Neue Hauptkonten Ihre E-Mail-Adresse ein.
    5. Wählen Sie in der Liste Rolle auswählen eine Rolle aus.
    6. Wenn Sie weitere Rollen hinzufügen möchten, klicken Sie auf Weitere Rolle hinzufügen und fügen Sie weitere Rollen hinzu.
    7. Klicken Sie auf Speichern.

Umgebung vorbereiten

  1. Starten Sie in der Google Cloud Console eine Cloud Shell-Instanz:
    Cloud Shell öffnen

  2. Legen Sie die Standardumgebungsvariablen fest:

      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

    Ersetzen Sie die folgenden Werte:

GKE-Standardcluster erstellen

Erledigen Sie mit Cloud Shell Folgendes:

  1. Erstellen Sie einen Standard-Cluster, der die Identitätsföderation von Arbeitslasten für GKE verwendet:

    gcloud container clusters create saxml \
    --zone=${ZONE} \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --cluster-version=VERSION \
    --num-nodes=4

    Ersetzen Sie VERSION durch die GKE-Versionsnummer. GKE unterstützt TPU v5e in Version 1.27.2-gke.2100 und höher. Weitere Informationen finden Sie unter TPU-Verfügbarkeit in GKE.

    Die Erstellung eines Clusters kann einige Minuten dauern.

  2. Erstellen Sie den ersten Knotenpool mit dem Namen tpu1:

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

  3. Erstellen Sie den zweiten Knotenpool mit dem Namen tpu2:

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

Sie haben die folgenden Ressourcen erstellt:

  • Ein Standardcluster mit vier CPU-Knoten.
  • Zwei v5e-TPU-Knotenpools mit 4x8-Topologie. Jeder Knotenpool stellt acht TPU-Knoten mit jeweils 4 Chips dar.

Das 175B-Modell muss mindestens auf einem 5-TPU-Slice mit mehreren Hosts mit 4x8-Topologiesegment (32 v5e TPU-Chips) bereitgestellt werden.

Cloud Storage-Bucket erstellen

Cloud Storage-Bucket zum Speichern der Saxml-Administratorserverkonfigurationen erstellen. Ein laufender Administratorserver speichert regelmäßig seinen Status und die Details der veröffentlichten Modelle.

Führen Sie in Cloud Shell folgenden Befehl aus:

gcloud storage buckets create gs://${GSBUCKET}

Arbeitslastzugriff mit Identitätsföderation von Arbeitslasten für GKE konfigurieren

Weisen Sie der Anwendung ein Kubernetes-ServiceAccount zu und konfigurieren Sie dieses Kubernetes-Dienstkonto als IAM-Dienstkonto.

  1. Konfigurieren Sie kubectl für die Kommunikation mit Ihrem Cluster:

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

  2. Erstellen Sie ein Kubernetes-ServiceAccount für die Anwendung:

    kubectl create serviceaccount sax-sa --namespace default

  3. Erstellen Sie ein IAM-Dienstkonto für Ihre Anwendung:

    gcloud iam service-accounts create sax-iam-sa

  4. Fügen Sie eine IAM-Richtlinienbindung für Ihr IAM-Dienstkonto zum Lesen und Schreiben in Cloud Storage hinzu:

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

  5. Erlauben Sie dem Kubernetes-ServiceAccount, die Identität des IAM-Dienstkontos zu übernehmen. Fügen Sie dazu eine IAM-Richtlinienbindung zwischen den beiden Dienstkonten hinzu. Durch diese Bindung kann das Kubernetes-ServiceAccount als IAM-Dienstkonto verwendet werden, sodass das Kubernetes-Dienstkonto Daten in Cloud Storage lesen und schreiben kann.

    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. Kennzeichnen Sie das Kubernetes-Dienstkonto mit der E-Mail-Adresse des IAM-Dienstkontos. Dadurch teilt die Beispielanwendung mit, welches Dienstkonto für den Zugriff auf Google Cloud-Dienste verwendet werden soll. Wenn die Anwendung also Standard-Google API-Clientbibliotheken für den Zugriff auf Google Cloud-Dienste verwendet, verwendet sie dieses IAM-Dienstkonto.

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

Saxml bereitstellen

In diesem Abschnitt stellen Sie den Saxml-Administratorserver und den Saxml-Modellserver bereit.

Saxml-Administratorserver bereitstellen

  1. Erstellen Sie das folgende sax-admin-server.yaml-Manifest:

    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

    Ersetzen Sie BUCKET_NAME durch den Namen des Cloud Storage-Buckets.

  2. Wenden Sie das Manifest an:

    kubectl apply -f sax-admin-server.yaml

  3. Prüfen Sie, ob der Administratorserver-Pod ausgeführt wird:

    kubectl get deployment

    Die Ausgabe sieht in etwa so aus:

    NAME               READY   UP-TO-DATE   AVAILABLE   AGE
sax-admin-server   1/1     1            1           52s

Saxml-Modellserver bereitstellen

Arbeitslasten, die in TPU-Slices mit mehreren Hosts ausgeführt werden, benötigen für jeden Pod eine stabile Netzwerk-ID, um Peers im selben TPU-Slice zu erkennen. Verwenden Sie zum Definieren dieser Kennungen IndexedJob, StatefulSet mit einem monitorlosen Dienst oder JobSet. Damit wird automatisch ein monitorloser Service für alle Jobs erstellt, die zum JobSet gehören. Im folgenden Abschnitt wird gezeigt, wie Sie mehrere Gruppen von Modellserver-Pods mit JobSet verwalten.

  1. Installieren Sie JobSet v0.2.3 oder höher.

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml

    Ersetzen Sie JOBSET_VERSION durch die JobSet-Version. Beispiel: v0.2.3.

  2. Prüfen Sie, ob der JobSet-Controller im Namespace jobset-system ausgeführt wird:

    kubectl get pod -n jobset-system

    Die Ausgabe sieht in etwa so aus:

    NAME                                        READY   STATUS    RESTARTS   AGE
jobset-controller-manager-69449d86bc-hp5r6   2/2     Running   0          2m15s

  3. Zwei Modellserver in zwei TPU-Knotenpools bereitstellen Speichern Sie das folgende sax-model-server-set-Manifest:

    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

    Ersetzen Sie BUCKET_NAME durch den Namen des Cloud Storage-Buckets.

    In diesem Manifest:

    • replicas: 2 ist die Anzahl der Jobreplikate. Jeder Job stellt einen Modellserver dar. Daher eine Gruppe von 8 Pods.
    • parallelism: 8 und completions: 8 entsprechen der Anzahl der Knoten in jedem Knotenpool.
    • backoffLimit: 0 muss null sein, um den Job als fehlgeschlagen zu markieren, wenn ein Pod fehlschlägt.
    • ports.containerPort: 8471 ist der Standardport für die TPU-VM-Kommunikation.
    • name: MEGASCALE_NUM_SLICES hebt die Festlegung der Umgebungsvariable auf, da GKE das Multislic-Training nicht ausführt.

  4. Wenden Sie das Manifest an:

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

  5. Prüfen Sie den Status der Saxml Admin Server- und Model Server-Pods:

    kubectl get pods

    Die Ausgabe sieht in etwa so aus:

    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 diesem Beispiel gibt es 16 Modellservercontainer: sax-model-server-set-sax-model-server-0-0-nj4sm und sax-model-server-set-sax-model-server-1-0-8mlf5 sind die beiden primären Modellserver in jeder Gruppe.

Ihr Saxml-Cluster hat zwei Modellserver, die auf zwei v5e-TPU-Knotenpools mit der Topologie 4x8 bereitgestellt werden.

Saxml-HTTP-Server und Load-Balancer bereitstellen

  1. Verwenden Sie das folgende vordefinierte HTTP-Server-Image. Speichern Sie das folgende sax-http.yaml-Manifest:

    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

    Ersetzen Sie BUCKET_NAME durch den Namen Ihres Cloud Storage-Buckets.

  2. Wenden Sie das sax-http.yaml-Manifest an:

    kubectl apply -f sax-http.yaml

  3. Warten Sie, bis der HTTP-Server-Container erstellt wurde:

    kubectl get pods

    Die Ausgabe sieht in etwa so aus:

    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. Warten Sie, bis dem Dienst eine externe IP-Adresse zugewiesen wurde:

    kubectl get svc

    Die Ausgabe sieht in etwa so aus:

    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

Saxml verwenden

Laden Sie das Modell auf dem Saxml im v5e-TPU-Multihost-Segment und stellen Sie es bereit:

Modell laden

  1. Rufen Sie die IP-Adresse des Load-Balancers für Saxml ab.

    LB_IP=$(kubectl get svc sax-http-lb -o jsonpath='{.status.loadBalancer.ingress[*].ip}')
PORT="8888"

  2. Laden Sie das Testmodell LmCloudSpmd175B in zwei v5e-TPU-Knotenpools:

    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
}'

    Das Testmodell hat keinen detaillierten Prüfpunkt. Die Gewichtungen werden nach dem Zufallsprinzip generiert. Das Laden des Modells kann bis zu zehn Minuten dauern.

    Die Ausgabe sieht in etwa so aus:

    {
    "model": "/sax/test/spmd",
    "path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
    "checkpoint": "None",
    "replicas": 2
}

  3. Prüfen Sie die Modellbereitschaft:

    kubectl logs sax-model-server-set-sax-model-server-0-0-nj4sm

    Die Ausgabe sieht in etwa so aus:

    ...
loading completed.
Successfully loaded model for key: /sax/test/spmd

    Das Modell ist vollständig geladen.

  4. Informationen zum Modell abrufen:

    curl --request GET \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/listcell --data \
'{
    "model": "/sax/test/spmd"
}'

    Die Ausgabe sieht in etwa so aus:

    {
"model": "/sax/test/spmd",
"model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
"checkpoint": "None",
"max_replicas": 2,
"active_replicas": 2
}

Modell bereitstellen

Stellen Sie einen Prompt bereit:

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?"
}'

Die Ausgabe zeigt ein Beispiel für die Modellantwort. Diese Antwort ist möglicherweise nicht aussagekräftig, da das Testmodell zufällige Gewichtungen hat.

Veröffentlichung des Modells aufheben

Führen Sie den folgenden Befehl aus, um die Veröffentlichung des Modells aufzuheben:

curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/unpublish --data \
'{
    "model": "/sax/test/spmd"
}'

Die Ausgabe sieht in etwa so aus:

{
  "model": "/sax/test/spmd"
}

Bereinigen

Damit Ihrem Google Cloud-Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, löschen Sie entweder das Projekt, das die Ressourcen enthält, oder Sie behalten das Projekt und löschen die einzelnen Ressourcen.

Bereitgestellte Ressourcen löschen

  1. Löschen Sie den Cluster, den Sie für diese Anleitung erstellt haben:

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

  2. Löschen Sie das Dienstkonto:

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

  3. Löschen Sie den Cloud Storage-Bucket:

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

Nächste Schritte