Menyajikan LLM menggunakan TPU multi-host di GKE dengan Saxml

Standard

Tutorial ini menunjukkan cara menyajikan model bahasa besar (LLM) menggunakan Tensor Processing Unit (TPU) di Google Kubernetes Engine (GKE) dengan Saxml.

Latar belakang

Saxml adalah sistem eksperimental yang menayangkan framework Paxml, JAX, dan PyTorch. TPU dapat digunakan untuk mempercepat pemrosesan data dengan {i>framework<i} ini. Untuk mendemonstrasikan deployment TPU di GKE, tutorial ini menampilkan model pengujian LmCloudSpmd175B32Test 175B. GKE men-deploy model pengujian ini pada dua kumpulan node TPU v5e dengan topologi 4x8, masing-masing.

Untuk men-deploy model pengujian dengan benar, topologi TPU telah ditentukan berdasarkan ukuran model. Mengingat bahwa model N miliar 16 bit kira-kira memerlukan memori sekitar 2 kali (2xN) GB, model 175B LmCloudSpmd175B32Test memerlukan memori sekitar 350 GB. Chip tunggal TPU v5e memiliki kapasitas 16 GB. Untuk mendukung 350 GB, GKE membutuhkan 21 chip 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)

Penting untuk memilih topologi TPU yang tepat untuk menyalurkan model saat men-deploy TPU di GKE. Untuk mempelajari lebih lanjut, lihat Merencanakan konfigurasi TPU.

Tujuan

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

Tutorial ini membahas langkah-langkah berikut:

Siapkan lingkungan Anda dengan cluster GKE Standard. Cluster ini memiliki dua kumpulan node TPU v5e dengan topologi 4x8.
Men-deploy Saxml. Saxml memerlukan server administrator, grup Pod yang berfungsi sebagai server model, server HTTP bawaan, dan load balancer.
Menggunakan Saxml untuk menyalurkan LLM.

Diagram berikut menunjukkan arsitektur yang diterapkan tutorial berikut:

Arsitektur TPU multi-host di GKE. — **Gambar:** Contoh arsitektur TPU multi-host di GKE.

Sebelum memulai

Login ke akun Google Cloud Anda. Jika Anda baru menggunakan Google Cloud, buat akun untuk mengevaluasi performa produk kami dalam skenario dunia nyata. Pelanggan baru juga mendapatkan kredit gratis senilai $300 untuk menjalankan, menguji, dan men-deploy workload.

Di konsol Google Cloud, pada halaman pemilih project, pilih atau buat project Google Cloud.

Buka pemilih project

Pastikan penagihan telah diaktifkan untuk project Google Cloud Anda.

Aktifkan API yang diperlukan.

Mengaktifkan API

Di konsol Google Cloud, pada halaman pemilih project, pilih atau buat project Google Cloud.

Buka pemilih project

Pastikan penagihan telah diaktifkan untuk project Google Cloud Anda.

Aktifkan API yang diperlukan.

Mengaktifkan API

Pastikan Anda memiliki peran berikut di project: roles/container.admin, roles/iam.serviceAccountAdmin
Memeriksa peran
1. Di konsol Google Cloud, buka halaman IAM.
  Buka IAM
2. Pilih project.
3. Di kolom Akun utama, cari baris yang berisi alamat email Anda.
  
  Jika alamat email Anda tidak ada di kolom tersebut, berarti Anda tidak memiliki peran apa pun.
4. Di kolom Peran untuk baris yang berisi alamat email Anda, periksa apakah daftar peran menyertakan peran yang diperlukan.
Memberikan peran
1. Di konsol Google Cloud, buka halaman IAM.
  Buka IAM
2. Pilih project.
3. Klik Berikan akses.
4. Di kolom Akun utama baru, masukkan alamat email Anda.
5. Di daftar Pilih peran, 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

Di konsol Google Cloud, mulai instance Cloud Shell:
Buka Cloud Shell

Tetapkan variabel lingkungan default:

  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

Ganti nilai berikut:

PROJECT_ID: ID project Google Cloud Anda.
COMPUTE_REGION: Region Compute Engine.
COMPUTE_ZONE: Zona tempat ct5lp-hightpu-4t tersedia.

Membuat cluster GKE Standar

Gunakan Cloud Shell untuk melakukan hal berikut:

Buat cluster Standar yang menggunakan Workload Identity Federation for 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 di versi 1.27.2-gke.2100 dan yang lebih baru. Untuk mengetahui informasi selengkapnya, lihat Ketersediaan TPU di GKE.

Pembuatan cluster memerlukan waktu beberapa menit.

Buat kumpulan node pertama bernama tpu1:

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

Buat kumpulan node kedua bernama tpu2:

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

Anda telah membuat referensi berikut:

Cluster Standar dengan empat node CPU.
Dua kumpulan node TPU v5e dengan topologi 4x8. Setiap kumpulan node mewakili delapan node TPU dengan masing-masing 4 chip.

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

Membuat bucket Cloud Storage

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

Jalankan perintah berikut di Cloud Shell:

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

Mengonfigurasi akses beban kerja menggunakan Workload Identity Federation for GKE

Menetapkan ServiceAccount Kubernetes ke aplikasi dan mengonfigurasi ServiceAccount Kubernetes tersebut agar berfungsi sebagai akun layanan IAM.

Konfigurasi kubectl untuk berkomunikasi dengan cluster Anda:

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

Buat ServiceAccount Kubernetes yang akan digunakan aplikasi Anda:
```
kubectl create serviceaccount sax-sa --namespace default
```

Buat akun layanan IAM untuk aplikasi Anda:

gcloud iam service-accounts create sax-iam-sa

Tambahkan binding kebijakan IAM agar akun layanan IAM Anda dapat membaca dan menulis ke Cloud Storage:

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

Izinkan Kubernetes ServiceAccount untuk meniru identitas akun layanan IAM dengan menambahkan binding kebijakan IAM antara dua akun layanan. Dengan binding ini, Kubernetes ServiceAccount dapat bertindak sebagai akun layanan IAM, sehingga ServiceAccount Kubernetes dapat membaca dan menulis ke 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]"
```
Beri anotasi pada akun layanan Kubernetes dengan alamat email akun layanan IAM. Hal ini memungkinkan aplikasi contoh Anda mengetahui akun layanan mana yang akan digunakan untuk mengakses layanan Google Cloud. 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.

Men-deploy server administrator Saxml

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

Ganti BUCKET_NAME dengan nama nama bucket Cloud Storage Anda.

Terapkan manifes:
```
kubectl apply -f sax-admin-server.yaml
```

Pastikan bahwa 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

Workload yang berjalan di slice TPU multi-host memerlukan ID jaringan yang stabil agar setiap Pod dapat menemukan peer di slice TPU yang sama. Untuk menentukan ID ini, gunakan IndexedJob, StatefulSet dengan Service headless, atau JobSet yang secara otomatis membuat Layanan headless untuk semua Tugas yang termasuk dalam JobSet. Bagian berikut menunjukkan cara mengelola beberapa grup Pod server model dengan JobSet.

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

Validasi pengontrol JobSet yang 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

Deploy dua server model di dua kumpulan node 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

Ganti BUCKET_NAME dengan nama nama bucket Cloud Storage Anda.

Dalam manifes ini:

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

Terapkan manifes:

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

Verifikasi status Pod Server Admin Saxml dan Pod 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 pada dua kumpulan node TPU v5e dengan topologi 4x8 masing-masing.

Men-deploy Server HTTP Saxml dan load balancer

Gunakan image server HTTP gambar 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

Ganti BUCKET_NAME dengan nama nama bucket Cloud Storage Anda.

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.

Tunggu hingga penampung Server HTTP 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
...

Tunggu hingga Layanan 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 salurkan model di Saxml dalam slice multihost TPU v5e:

Memuat model

Mengambil alamat IP load balancer untuk Saxml.

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

Muat model pengujian LmCloudSpmd175B di dua kumpulan node 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 telah disesuaikan, bobotnya 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
}

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 dimuat sepenuhnya.

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

Menyalurkan 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-nya menunjukkan 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"
}

Pembersihan

Agar tidak dikenakan 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-nya.

Menghapus resource yang di-deploy

Hapus cluster yang Anda buat untuk tutorial ini:

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

Hapus akun layanan:

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

Hapus bucket Cloud Storage:
```
gcloud storage rm -r gs://${GSBUCKET}
```

Langkah selanjutnya

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

Menyajikan LLM menggunakan TPU multi-host di GKE dengan Saxml

Latar belakang

Tujuan

Sebelum memulai

Memeriksa peran

Memberikan peran

Menyiapkan lingkungan

Membuat cluster GKE Standar

Membuat bucket Cloud Storage

Mengonfigurasi akses beban kerja menggunakan Workload Identity Federation for GKE

Men-deploy Saxml

Men-deploy server administrator Saxml

Men-deploy server model Saxml

Men-deploy Server HTTP Saxml dan load balancer

Menggunakan Saxml

Memuat model

Menyalurkan model

Membatalkan publikasi model

Pembersihan

Menghapus resource yang di-deploy

Langkah selanjutnya