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 memerlukan sekitar 2 kali (2xN) GB memori, 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.
Menyiapkan lingkungan
Di konsol Google Cloud , mulai instance Cloud Shell:
Open Cloud ShellTetapkan variabel lingkungan default:
gcloud config set project PROJECT_ID export PROJECT_ID=$(gcloud config get project) export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION export BUCKET_NAME=PROJECT_ID-gke-bucket
Ganti nilai berikut:
- PROJECT_ID: Google Cloud Project ID Anda.
- CONTROL_PLANE_LOCATION: zona Compute Engine bidang kontrol cluster Anda. Pilih 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:
Buat cluster Standar yang menggunakan Workload Identity Federation untuk GKE:
gcloud container clusters create saxml \ --location=${CONTROL_PLANE_LOCATION} \ --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.
Buat node pool pertama bernama
tpu1
:gcloud container node-pools create tpu1 \ --location=${CONTROL_PLANE_LOCATION} \ --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.Buat node pool kedua bernama
tpu2
:gcloud container node-pools create tpu2 \ --location=${CONTROL_PLANE_LOCATION} \ --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.
Konfigurasi
kubectl
untuk berkomunikasi dengan cluster Anda:gcloud container clusters get-credentials saxml --location=${CONTROL_PLANE_LOCATION}
Buat Akun Layanan Kubernetes yang akan digunakan untuk 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 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
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}.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. 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, aplikasi akan menggunakan akun layanan IAM tersebut. Google Cloud
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
Buat manifes
sax-admin-server.yaml
berikut:Ganti
BUCKET_NAME
dengan Cloud Storage yang Anda buat sebelumnya:perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-admin-server.yaml
Terapkan manifes:
kubectl apply -f sax-admin-server.yaml
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.
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
.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
Deploy dua server model di dua node pool slice TPU. Simpan manifes
sax-model-server-set
berikut: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
dancompletions: 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 VMname: MEGASCALE_NUM_SLICES
menghapus setelan variabel lingkungan karena GKE tidak menjalankan pelatihan Multislice.
Terapkan manifes:
kubectl apply -f sax-model-server-set.yaml
Verifikasi status Pod Saxml Admin Server dan Model Server:
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 container 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
Gunakan image server HTTP image bawaan berikut. Simpan manifes
sax-http.yaml
berikut:Ganti
BUCKET_NAME
dengan Cloud Storage yang Anda buat sebelumnya:perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-http.yaml
Terapkan manifes
sax-http.yaml
:kubectl apply -f sax-http.yaml
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 ...
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 sajikan model di Saxml dalam slice multi-host TPU v5e:
Memuat model
Ambil alamat IP load balancer untuk Saxml.
LB_IP=$(kubectl get svc sax-http-lb -o jsonpath='{.status.loadBalancer.ingress[*].ip}') PORT="8888"
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 }
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.
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"
}