LLM mithilfe von mehrere Hosts in GKE mit Saxml bereitstellen

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:

   • PROJECT_ID: Ihre Google Cloud-Projekt-ID.
   • COMPUTE_REGION: Die Compute Engine-Region.
   • COMPUTE_ZONE: Die Zone, in der ct5lp-hightpu-4t verfügbar ist.

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