Menyajikan LLM menggunakan TPU multi-host di GKE dengan Saxml

Menyiapkan lingkungan

1. Di konsol Trusted Cloud , mulai instance Cloud Shell:
   Open Cloud Shell

2. Tetapkan variabel lingkungan default:

     gcloud config set project PROJECT_ID
     export PROJECT_ID=$(gcloud config get project)
     export ZONE=COMPUTE_ZONE
     export BUCKET_NAME=PROJECT_ID-gke-bucket
   

   Ganti nilai berikut:

   • PROJECT_ID: Trusted Cloud Project ID Anda.
   • COMPUTE_ZONE: Zona tempat ct5lp-hightpu-4t tersedia.

   Dalam perintah ini, BUCKET_NAME menentukan nama bucket penyimpanan Trusted Cloud untuk menyimpan konfigurasi server administrator Saxml.

Membuat cluster Standar GKE

Gunakan Cloud Shell untuk melakukan hal berikut:

1. Buat cluster Standar yang menggunakan Workload Identity Federation untuk GKE:

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

   Ganti VERSION dengan nomor versi GKE. GKE mendukung TPU v5e pada versi 1.27.2-gke.2100 dan yang lebih baru. Untuk mengetahui informasi selengkapnya, lihat Ketersediaan TPU di GKE.

   Pembuatan cluster mungkin memerlukan waktu beberapa menit.

2. Buat node pool pertama bernama tpu1:

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

   Nilai untuk tanda --num-nodes dihitung dengan membagi topologi TPU dengan jumlah TPU chip per slice TPU. Dalam hal ini: (4 * 8) / 4.

3. Buat node pool kedua bernama tpu2:

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

   Nilai untuk tanda --num-nodes dihitung dengan membagi topologi TPU dengan jumlah TPU chip per slice TPU. Dalam hal ini: (4 * 8) / 4.

Anda telah membuat resource berikut:

• Cluster Standard dengan empat node CPU.
• Dua node pool slice TPU v5e dengan topologi 4x8. Setiap node pool mewakili delapan node slice TPU dengan masing-masing 4 chip TPU.

Model 175B harus ditayangkan di slice TPU v5e multi-host dengan slice topologi 4x8 (minimal 32 chip TPU v5e).

Membuat bucket Cloud Storage

Buat bucket Cloud Storage untuk menyimpan konfigurasi server administrator Saxml. Server administrator yang sedang berjalan secara berkala menyimpan statusnya dan detail model yang dipublikasikan.

Di Cloud Shell, jalankan perintah berikut:

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

Mengonfigurasi akses workload Anda menggunakan Workload Identity Federation for GKE

Tetapkan ServiceAccount Kubernetes ke aplikasi dan konfigurasikan ServiceAccount Kubernetes tersebut untuk bertindak sebagai akun layanan IAM.

1. Konfigurasi kubectl untuk berkomunikasi dengan cluster Anda:

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

2. Buat Akun Layanan Kubernetes yang akan digunakan untuk aplikasi Anda:

   kubectl create serviceaccount sax-sa --namespace default
   

3. Buat akun layanan IAM untuk aplikasi Anda:

   gcloud iam service-accounts create sax-iam-sa
   

4. Tambahkan binding kebijakan IAM untuk akun layanan IAM Anda agar dapat membaca dan menulis di Cloud Storage:

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

5. Izinkan ServiceAccount Kubernetes untuk meniru identitas akun layanan IAM dengan menambahkan binding kebijakan IAM antara kedua akun layanan tersebut. Dengan binding ini, Akun Layanan Kubernetes dapat bertindak sebagai akun layanan IAM, sehingga Akun Layanan Kubernetes dapat membaca dan menulis di Cloud Storage.

   gcloud iam service-accounts add-iam-policy-binding sax-iam-sa@${PROJECT_ID}. \
     --role roles/iam.workloadIdentityUser \
     --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/sax-sa]"
   

6. Beri anotasi pada akun layanan Kubernetes dengan alamat email akun layanan IAM. Dengan demikian, aplikasi contoh Anda akan mengetahui akun layanan mana yang akan digunakan untuk mengakses Trusted Cloud layanan. Jadi, saat menggunakan Library Klien Google API standar untuk mengakses layanan, aplikasi akan menggunakan akun layanan IAM tersebut. Trusted Cloud

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

Men-deploy Saxml

Di bagian ini, Anda akan men-deploy server administrator Saxml dan server model Saxml.

Deploy server administrator Saxml

1. Buat manifes sax-admin-server.yaml berikut:

   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. Ganti BUCKET_NAME dengan Cloud Storage yang Anda buat sebelumnya:

   perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-admin-server.yaml
   

3. Terapkan manifes:

   kubectl apply -f sax-admin-server.yaml
   

4. Pastikan Pod server administrator sudah aktif dan berjalan:

   kubectl get deployment
   

   Outputnya mirip dengan hal berikut ini:

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

Men-deploy server model Saxml

Beban kerja yang berjalan di slice TPU multi-host memerlukan ID jaringan yang stabil untuk setiap Pod agar dapat menemukan peer di slice TPU yang sama. Untuk menentukan ID ini, gunakan IndexedJob, StatefulSet dengan Service headless atau JobSet yang otomatis membuat Service headless untuk semua Job yang termasuk dalam JobSet. Jobset adalah API beban kerja yang memungkinkan Anda mengelola sekelompok Tugas Kubernetes sebagai satu unit. Kasus penggunaan yang paling umum untuk JobSet adalah pelatihan terdistribusi, tetapi Anda juga dapat menggunakannya untuk menjalankan workload batch.

Bagian berikut menunjukkan cara mengelola beberapa grup Pod server model dengan JobSet.

1. Instal JobSet v0.2.3 atau yang lebih baru.

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

   Ganti JOBSET_VERSION dengan versi JobSet. Contoh, v0.2.3.

2. Validasi bahwa pengontrol JobSet berjalan di namespace jobset-system:

   kubectl get pod -n jobset-system
   

   Outputnya mirip dengan hal berikut ini:

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

3. Deploy dua server model di dua node pool slice TPU. Simpan manifes sax-model-server-set berikut:

   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. Ganti BUCKET_NAME dengan Cloud Storage yang Anda buat sebelumnya:

   perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-model-server-set.yaml
   

   Dalam manifes ini:

   • replicas: 2 adalah jumlah replika Tugas. Setiap tugas mewakili server model. Oleh karena itu, grup 8 Pod.
   • parallelism: 8 dan completions: 8 sama dengan jumlah node di setiap node pool.
   • backoffLimit: 0 harus nol untuk menandai Tugas sebagai gagal jika ada Pod yang gagal.
   • ports.containerPort: 8471 adalah port default untuk komunikasi VM
   • name: MEGASCALE_NUM_SLICES menghapus setelan variabel lingkungan karena GKE tidak menjalankan pelatihan Multislice.

5. Terapkan manifes:

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

6. Verifikasi status Pod Server Admin Saxml dan Server Model:

   kubectl get pods
   

   Outputnya mirip dengan hal berikut ini:

   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
   

Dalam contoh ini, ada 16 penampung server model: sax-model-server-set-sax-model-server-0-0-nj4sm dan sax-model-server-set-sax-model-server-1-0-8mlf5 adalah dua server model utama di setiap grup.

Cluster Saxml Anda memiliki dua server model yang di-deploy di dua node pool slice TPU v5e dengan topologi 4x8 masing-masing.

Men-deploy Saxml HTTP Server dan load balancer

1. Gunakan image server HTTP image bawaan berikut. Simpan manifes sax-http.yaml berikut:

   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. Ganti BUCKET_NAME dengan Cloud Storage yang Anda buat sebelumnya:

   perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-http.yaml
   

3. Terapkan manifes sax-http.yaml:

   kubectl apply -f sax-http.yaml
   

4. Tunggu hingga penampung HTTP Server selesai dibuat:

   kubectl get pods
   

   Outputnya mirip dengan hal berikut ini:

   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. Tunggu hingga Service memiliki alamat IP eksternal yang ditetapkan:

   kubectl get svc
   

   Outputnya mirip dengan hal berikut ini:

   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
   

Menggunakan Saxml

Muat, deploy, dan sajikan model di Saxml dalam slice multi-host TPU v5e:

Memuat model

1. Ambil alamat IP load balancer untuk Saxml.

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

2. Muat model pengujian LmCloudSpmd175B di dua node pool 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
   }'
   

   Model pengujian tidak memiliki titik pemeriksaan yang di-fine-tune, bobot dihasilkan secara acak. Pemuatan model dapat memerlukan waktu hingga 10 menit.

   Outputnya mirip dengan hal berikut ini:

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

3. Periksa kesiapan model:

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

   Outputnya mirip dengan hal berikut ini:

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

   Model telah dimuat sepenuhnya.

4. Mendapatkan informasi tentang model:

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

   Outputnya mirip dengan hal berikut ini:

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

Menyajikan model

Menayangkan permintaan perintah:

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

Output menampilkan contoh respons model. Respons ini mungkin tidak bermakna karena model pengujian memiliki bobot acak.

Membatalkan publikasi model

Jalankan perintah berikut untuk membatalkan publikasi model:

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

Outputnya mirip dengan hal berikut ini:

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