LLM mithilfe von mehrere Hosts in GKE mit Saxml bereitstellen

In dieser Anleitung erfahren Sie, wie Sie ein Large Language Model (LLM) unter Verwendung eines TPU-Slice-Knotenpools mit mehreren Hosts auf Google Kubernetes Engine (GKE) mit Saxml für eine effiziente skalierbare Architektur bereitstellen und verfügbar machen.

Hintergrund

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

Für eine korrekte Bereitstellung des Testmodells 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 TPU-Chip TPU v5e hat 16 GB. Zur Unterstützung von 350 GB benötigt GKE 21 v5e-TPU-Chips (350/16= 21). Basierend auf der Zuordnung der TPU-Konfiguration, ist die richtige TPU-Konfiguration für diese Anleitung:

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

Bei der Bereitstellung von TPUs in GKE ist es wichtig, die richtige TPU-Topologie für die Modellbereitstellung auszuwählen. Weitere Informationen finden Sie unter TPU-Konfiguration planen

Lernziele

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

Diese Anleitung umfasst die folgenden Schritte:

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

Das folgende Diagramm zeigt die Architektur, die die folgende Anleitung implementiert:

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

Hinweise

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

      Zu IAM
    2. Wählen Sie das Projekt aus.
    3. Klicken Sie auf Zugriff erlauben.

    4. Geben Sie im Feld Neue Hauptkonten Ihre Nutzer-ID ein. Dies ist in der Regel die E-Mail-Adresse eines Google-Kontos.

    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 der 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-Slice-Knotenpools mit 4x8-Topologie. Jeder Knotenpool stellt acht TPU-Slice-Knoten mit jeweils 4 TPU-Chips dar.

Das 175B-Modell muss auf einem v5e-TPU-Slice mit mehreren Hosts mit Topologiesegment 4x8 (mindestens 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 Details zu den veröffentlichten Modellen.

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 weiß Ihre Beispielanwendung, welches Dienstkonto für den Zugriff auf Google Cloud-Dienste verwendet werden soll. Wenn die App also 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 Ihres Cloud Storage-Buckets.

  2. Wenden Sie das Manifest an:

    kubectl apply -f sax-admin-server.yaml

  3. Prüfen Sie, ob der Administrator-Server-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 eine stabile Netzwerkkennung für jeden Pod, 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. Die folgenden Abschnitte zeigen, 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. Ausführung des JobSet-Controllers im Namespace jobset-system prüfen:

    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. Stellen Sie zwei Modellserver in zwei TPU-Slice-Knotenpools bereit. 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 Ihres 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 im jeweiligen 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 VM-Kommunikation
    • name: MEGASCALE_NUM_SLICES hebt die Festlegung der Umgebungsvariable auf, weil In der GKE kein Multislice-Training ausgeführt wird.

  4. Wenden Sie das Manifest an:

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

  5. Prüfen Sie den Status der Saxml Admininistratorserver- und Modellserver-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 Hauptmodellserver in jeder Gruppe.

Ihr Saxml-Cluster hat zwei Modellserver, die auf zwei v5e-TPU-Slice-Knotenpools mit 4x8-Topologie 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, implementieren und stellen Sie das Modell auf Saxml im v5e TPU-Multihost-Slice 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 LmCloudSpmd175B-Testmodell in zwei v5e-TPU-Slice-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 abgestimmten Prüfpunkt. Die Gewichtungen sind zufällig generiert. Das Laden des Modells kann bis zu 10 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. Rufen Sie Informationen zum Modell ab:

    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 sinnvoll, weil 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