Menyajikan LLM menggunakan TPU multi-host di GKE dengan Saxml

Standard

---

Perhatian: Saxml di GKE tidak lagi diprioritaskan mulai 24 April 2025. Artinya, project tidak akan mendapatkan update lebih lanjut. Deployment Saxml yang ada akan terus berfungsi seperti apa adanya tanpa gangguan. Sebaiknya Anda bermigrasi ke JetStream, framework inferensi open source terbaru dari Google untuk penayangan LLM berperforma tinggi di TPU dan GPU. JetStream menawarkan batching dan kuantisasi berkelanjutan untuk throughput dan efisiensi memori yang lebih baik. Untuk contoh migrasi, lihat Menyajikan Gemma menggunakan TPU di GKE dengan JetStream.

Tutorial ini menunjukkan cara men-deploy dan menyajikan model bahasa besar (LLM) menggunakan node pool slice TPU multi-host di Google Kubernetes Engine (GKE) dengan Saxml untuk arsitektur yang dapat diskalakan secara efisien.

Latar belakang

Saxml adalah sistem eksperimental yang melayani framework Paxml, JAX, dan PyTorch. Anda dapat menggunakan TPU untuk mempercepat pemrosesan data dengan framework ini. Untuk mendemonstrasikan deployment TPU di GKE, tutorial ini menyajikan model pengujian LmCloudSpmd175B32Test 175B. GKE men-deploy model pengujian ini di dua node pool slice TPU v5e dengan topologi 4x8.

Untuk men-deploy model pengujian dengan benar, topologi TPU telah ditentukan berdasarkan ukuran model. Mengingat model N miliar 16 bit kira-kira memerlukan memori sekitar 2 kali (2xN) GB, model LmCloudSpmd175B32Test 175B memerlukan memori sekitar 350 GB. Chip TPU tunggal TPU v5e memiliki kapasitas 16 GB. Untuk mendukung 350 GB, GKE memerlukan 21 chip TPU v5e (350/16= 21). Berdasarkan pemetaan konfigurasi TPU, konfigurasi TPU yang tepat untuk tutorial ini adalah:

• Jenis mesin: ct5lp-hightpu-4t
• Topologi: 4x8 (32 jumlah chip TPU)

Memilih topologi TPU yang tepat untuk menyajikan model penting saat men-deploy TPU di GKE. Untuk mempelajari lebih lanjut, lihat Merencanakan konfigurasi TPU.

Tujuan

Tutorial ini ditujukan untuk engineer MLOps atau DevOps atau administrator platform yang ingin menggunakan kemampuan orkestrasi GKE untuk menyajikan model data.

Tutorial ini membahas langkah-langkah berikut:

1. Siapkan lingkungan Anda dengan cluster GKE Standard. Cluster memiliki dua node pool slice TPU v5e dengan topologi 4x8.
2. Deploy Saxml. Saxml memerlukan server administrator, sekelompok Pod yang berfungsi sebagai server model, server HTTP bawaan, dan load balancer.
3. Gunakan Saxml untuk menyajikan LLM.

Diagram berikut menunjukkan arsitektur yang diimplementasikan dalam tutorial berikut:

Gambar: Contoh arsitektur TPU multi-host di GKE.

Sebelum memulai

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.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

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

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify 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, 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.

     Buka IAM
  2. Pilih project.
  3. Klik Grant access.

  4. Di kolom New principals, masukkan ID pengguna Anda. Biasanya berupa alamat email untuk Akun Google.

  5. Di daftar Select a role, pilih peran.
  6. Untuk memberikan peran tambahan, klik Tambahkan peran lain, lalu tambahkan setiap peran tambahan.
  7. Klik Simpan.
  • Pastikan project Anda memiliki kuota yang cukup untuk Cloud TPU di GKE.

  Menyiapkan lingkungan

  1. Di konsol Google 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: Google Cloud Project ID Anda.
     • COMPUTE_ZONE: Zona tempat ct5lp-hightpu-4t tersedia.

     Dalam perintah ini, BUCKET_NAME menentukan nama bucket penyimpanan Google 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 agar 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}.iam.gserviceaccount.com" \
       --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, ServiceAccount Kubernetes dapat bertindak sebagai akun layanan IAM, sehingga ServiceAccount Kubernetes dapat membaca dan menulis di 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. 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 Google Cloud layanan. Jadi, saat menggunakan Library Klien Google API standar untuk mengakses layanan Google Cloud , aplikasi akan menggunakan akun layanan IAM tersebut.

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

  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
     

     Catatan: Anda dapat menggunakan server HTTP Anda sendiri yang dibuat untuk Saxml. Untuk mempelajari lebih lanjut, lihat Inferensi menggunakan Saxml dan Server HTTP.

  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. Memuat 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 checkpoint yang di-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.

  Berhasil: Anda telah berhasil menyajikan LLM menggunakan TPU multi-host di GKE.

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

  Pembersihan

  Agar tidak perlu membayar biaya pada akun Google Cloud Anda untuk resource yang digunakan dalam tutorial ini, hapus project yang berisi resource tersebut, atau simpan project dan hapus setiap resource.

  Menghapus resource yang di-deploy

  1. Hapus cluster yang Anda buat untuk tutorial ini:

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

  2. Hapus akun layanan:

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

  3. Hapus bucket Cloud Storage:

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

  Langkah berikutnya

  • Pelajari versi TPU saat ini dengan arsitektur sistem Cloud TPU.
  • Pelajari lebih lanjut TPU di GKE.