Menyajikan LLM menggunakan TPU multi-host di GKE dengan JetStream dan Pathways

Standard

Panduan ini menunjukkan cara menayangkan model bahasa besar (LLM) canggih seperti Llama 3.1 405B di Google Kubernetes Engine (GKE) menggunakan unit pemrosesan tensor (TPU) di beberapa node.

Panduan ini menunjukkan cara menggunakan teknologi open source portabel—Kubernetes, JetStream, Pathways on Cloud, dan LeaderWorkerSet (LWS) API—untuk men-deploy dan menyajikan workload AI/ML di GKE, dengan memanfaatkan kontrol terperinci, skalabilitas, ketahanan, portabilitas, dan efisiensi biaya GKE.

Latar belakang

Model bahasa besar telah bertambah ukurannya dan tidak lagi sesuai dengan slice TPU host tunggal. Untuk inferensi ML, Anda dapat menggunakan Pathways di Cloud untuk menjalankan inferensi multi-host skala besar di GKE di beberapa node TPU yang saling terhubung. Dalam panduan ini, Anda akan mempelajari cara menyediakan cluster GKE dengan slice TPU multi-host dan menggunakan biner Pathways on Cloud, meluncurkan Server JetStream dengan framework MaxText, dan membuat permintaan inferensi multi-host.

Dengan menyajikan LLM menggunakan TPU di GKE dengan JetStream, MaxText, dan Pathways, Anda dapat membuat solusi penayangan yang tangguh dan siap produksi dengan semua manfaat Kubernetes terkelola, termasuk efisiensi biaya, skalabilitas, dan ketersediaan yang lebih tinggi. Bagian ini menjelaskan teknologi utama yang digunakan dalam tutorial ini.

Tentang TPU

TPU adalah sirkuit terintegrasi khusus aplikasi (ASIC) yang dikembangkan khusus oleh Google dan digunakan untuk mempercepat model machine learning dan AI yang dibuat menggunakan framework seperti TensorFlow, PyTorch, dan JAX.

Sebelum menggunakan TPU di GKE, sebaiknya selesaikan jalur pembelajaran berikut:

1. Pelajari ketersediaan versi TPU saat ini dengan arsitektur sistem Cloud TPU.
2. Pelajari TPU di GKE.

Tutorial ini membahas cara menyajikan model Llama 3.1-405B. GKE men-deploy model pada node TPU v6e multi-host dengan topologi TPU yang dikonfigurasi berdasarkan persyaratan model untuk menyajikan perintah dengan latensi rendah.

Pathways on Cloud

Pathways adalah lapisan orkestrasi skala besar untuk akselerator. Pathways dirancang secara eksplisit untuk memungkinkan eksplorasi sistem baru dan ide riset ML, sekaligus mempertahankan performa terbaik untuk model saat ini. Pathways memungkinkan satu proses klien JAX mengoordinasikan komputasi di satu atau beberapa slice TPU besar, sehingga menyederhanakan komputasi ML yang mencakup ratusan atau ribuan chip TPU.

JetStream

JetStream adalah framework penayangan inferensi open source yang dikembangkan oleh Google. JetStream memungkinkan inferensi berperforma tinggi, throughput tinggi, dan yang dioptimalkan untuk memori di TPU dan GPU. JetStream menyediakan pengoptimalan performa lanjutan, termasuk pengelompokan berkelanjutan, pengoptimalan cache KV, dan teknik kuantisasi, untuk memfasilitasi deployment LLM. JetStream memungkinkan penayangan TPU JAX dan PyTorch/XLA untuk mengoptimalkan performa.

MaxText

MaxText adalah implementasi LLM JAX berperforma tinggi, skalabel, dan dapat disesuaikan, yang dibangun di atas library JAX open source seperti Flax, Orbax, dan Optax. Implementasi LLM khusus decoder MaxText ditulis dalam Python. Hal ini memanfaatkan compiler XLA secara intensif untuk mencapai performa tinggi tanpa perlu membuat kernel kustom.

Untuk mengetahui informasi selengkapnya tentang model dan ukuran parameter terbaru yang didukung MaxText, lihat repositori project MaxText.

Llama 3.1 405B

Llama 3.1 405B adalah model bahasa besar dari Meta yang dirancang untuk berbagai tugas natural language processing, termasuk pembuatan teks, terjemahan, dan question answering. GKE menawarkan infrastruktur yang diperlukan untuk mendukung kebutuhan pelatihan dan penyajian terdistribusi model dalam skala ini.

Untuk mengetahui informasi selengkapnya, lihat dokumentasi Llama.

Arsitektur

Bagian ini menjelaskan arsitektur GKE yang digunakan dalam tutorial ini. Arsitektur ini mencakup cluster GKE Standard yang menyediakan TPU dan menghosting komponen JetStream dan Pathways untuk men-deploy dan menayangkan model.

Diagram berikut menunjukkan komponen arsitektur ini:

Arsitektur ini mencakup komponen berikut:

• Cluster regional GKE Standard.
• Node pool slice TPU multi-host yang menghosting komponen JetStream deployment dan Pathways.
• Pathways resource manager mengelola resource akselerator dan mengoordinasikan alokasi akselerator untuk tugas pengguna.
• Pathways client berkoordinasi dengan Pathways resource manager untuk menentukan tempat program yang dikompilasi ditempatkan untuk dieksekusi.
• Pathways worker berjalan dan melakukan komputasi di mesin akselerator, serta mengirimkan data kembali ke workload Anda melalui server proxy IFRT.
• IFRT proxy client mengimplementasikan API Interim Framework Runtime (IFRT) OSS dan bertindak sebagai jembatan komunikasi antara workload dan komponen Pathways Anda.
• IFRT proxy server menerima permintaan dari IFRT proxy client dan meneruskannya ke Pathways client, mendistribusikan tugas.
• Container JetStream-Pathways menyediakan server inferensi berbasis JAX yang menerima permintaan inferensi dan mendelegasikan proses eksekusinya ke Pathways workers
• Komponen Service menyebarkan traffic masuk ke semua replika JetStream HTTP.
• JetStream HTTP adalah server HTTP yang menerima permintaan sebagai wrapper ke format yang diperlukan JetStream dan mengirimkannya ke klien GRPC JetStream.

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.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

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.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the API

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

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.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the API

• Pastikan Anda memiliki peran berikut di project: roles/container.admin, roles/iam.serviceAccountAdmin, roles/resourcemanager.projectIamAdmin

  Memeriksa peran

  1. Di konsol Google Cloud , buka halaman IAM.

     Buka IAM
  2. Pilih project.

  3. Di kolom Akun utama, temukan semua baris yang mengidentifikasi Anda atau grup yang Anda ikuti. Untuk mengetahui grup mana saja yang Anda ikuti, hubungi administrator Anda.

  4. Untuk semua baris yang menentukan atau menyertakan Anda, periksa kolom Peran untuk melihat apakah daftar peran menyertakan peran yang diperlukan.

  Memberikan peran

  1. Di konsol Google Cloud , buka halaman IAM.

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

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

  5. Klik Pilih peran, lalu telusuri peran.
  6. Untuk memberikan peran tambahan, klik Add another role, lalu tambahkan tiap peran tambahan.
  7. Klik Simpan.

• Pastikan Anda memiliki kuota yang cukup untuk enam belas chip PodSlice Lite TPU v6e. Dalam tutorial ini, Anda akan menggunakan instance sesuai permintaan.

Mendapatkan akses ke model

Untuk mendapatkan akses ke checkpoint Meta Llama 3.1-405B untuk deployment ke GKE, ikuti langkah-langkah berikut:

1. Tanda tangani perjanjian izin lisensi.
2. Buka halaman download Meta Llama.
3. Tinjau dan setujui Persyaratan dan Ketentuan model serta dapatkan URL yang diperlukan untuk mendownload model.
4. Untuk mendownload checkpoint model, temukan ID model untuk model yang sesuai. Untuk mengetahui daftar model yang didukung dan ID-nya, lihat dokumentasi CLI llama. Misalnya, gunakan Llama 3.1-405B-Instruct:bf16-mp16 untuk model Llama 3.1-405B.

Menyiapkan lingkungan

Dalam tutorial ini, Anda akan menggunakan Cloud Shell untuk mengelola resource yang dihosting di Google Cloud. Cloud Shell telah diinstal dengan software yang Anda perlukan untuk tutorial ini, termasuk kubectl dan gcloud CLI.

Untuk menyiapkan lingkungan Anda dengan Cloud Shell, ikuti langkah-langkah berikut:

1. Di konsol Google Cloud , luncurkan sesi Cloud Shell dengan mengklik Aktifkan Cloud Shell di konsolGoogle Cloud . Tindakan ini akan meluncurkan sesi di panel bawah konsol Google Cloud .

2. Tetapkan variabel lingkungan default:

   gcloud config set project PROJECT_ID
   gcloud config set billing/quota_project PROJECT_ID
   export PROJECT_ID=$(gcloud config get project)
   export CLUSTER_NAME=CLUSTER_NAME
   export BUCKET_NAME=BUCKET_NAME
   export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
   export NODE_LOCATION=NODE_LOCATION
   export CLUSTER_VERSION=CLUSTER_VERSION
   export MACHINE_TYPE=ct6e-standard-4t
   export TPU_TYPE=v6e
   export TOPOLOGY=4x4
   export WORKERS_PER_SLICE=4
   

   Ganti nilai berikut:

   • PROJECT_ID: Google Cloud Project ID Anda.
   • CLUSTER_NAME: nama cluster GKE Anda.
   • BUCKET_NAME: nama bucket Cloud Storage Anda. Anda tidak perlu menentukan awalan gs://.
   • CONTROL_PLANE_LOCATION: region Compute Engine tempat cluster GKE, bucket Cloud Storage, dan node TPU Anda berada. Region berisi zona tempat jenis mesin TPU v6e tersedia (misalnya, us-east1, us-east5, europe-west4, asia-northeast1, atau us-south1).
   • NODE_LOCATION: zona tempat resource TPU tersedia (misalnya, us-east1-d).
   • CLUSTER_VERSION: versi GKE, yang harus mendukung jenis mesin yang ingin Anda gunakan. Perhatikan bahwa versi GKE default mungkin tidak tersedia untuk TPU target Anda. Untuk mengetahui daftar versi GKE minimum yang tersedia menurut jenis mesin TPU, lihat Ketersediaan TPU di GKE.
   • MACHINE_TYPE: jenis mesin v6e.
   • TPU_TYPE: awalan yang digunakan untuk memberi nama node pool (v6e).
   • TOPOLOGY: topologi TPU v6e.
   • WORKERS_PER_SLICE: jumlah node per node pool atau slice TPU.

Membuat dan mengonfigurasi resource Google Cloud

Untuk membuat resource yang diperlukan, ikuti petunjuk berikut:

Catatan: Anda mungkin perlu membuat reservasi kapasitas untuk menggunakan beberapa akselerator. Untuk mengetahui informasi selengkapnya tentang cara memesan dan menggunakan resource yang dipesan, lihat Menggunakan resource zona yang dipesan.

Membuat cluster GKE

1. Buat cluster GKE Standar regional:

   gcloud container clusters create ${CLUSTER_NAME} \
       --project=${PROJECT_ID} \
       --cluster-version=${CLUSTER_VERSION} \
       --location=${CONTROL_PLANE_LOCATION} \
       --scopes=cloud-platform \
       --machine-type=n2-standard-32
   

   Pembuatan cluster mungkin memerlukan waktu beberapa menit.

   Ganti CLUSTER_VERSION dengan versi cluster yang sesuai.

2. Buat satu node pool TPU v6e dengan topologi 4x4 dan masing-masing empat node:

   gcloud container node-pools create multihost-np \
       --project=${PROJECT_ID} \
       --location=${CONTROL_PLANE_LOCATION} \
       --node-locations=${NODE_LOCATION} \
       --cluster=${CLUSTER_NAME} \
       --machine-type=${MACHINE_TYPE} \
       --num-nodes=${WORKERS_PER_SLICE} \
       --tpu-topology=${TOPOLOGY} \
       --scopes cloud-platform \
       --placement-type=COMPACT \
       --workload-metadata=GCE_METADATA
   

   Catatan: --workload-metadata=GCE_METADATA menunjukkan bahwa Workload Identity Federation untuk GKE dinonaktifkan di node pool. Node pool yang menjalankan Pathways dengan JetStream dan MaxText harus menonaktifkan Workload Identity Federation for GKE.

Mengonfigurasi akun layanan untuk akses Objek Storage

Konfigurasi akun layanan Kubernetes agar bertindak sebagai akun layanan IAM.

1. Buat akun layanan IAM untuk aplikasi Anda:

   gcloud iam service-accounts create jetstream-pathways
   

2. Tambahkan binding kebijakan IAM untuk akun layanan IAM Anda guna mengelola Cloud Storage. Hal ini dilakukan agar akun layanan IAM Anda dapat mengakses bucket penyimpanan tempat checkpoint Anda akan disimpan:

   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
     --member "serviceAccount:jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com" \
     --role roles/storage.objectUser
   
   gcloud projects add-iam-policy-binding ${PROJECT_ID} \
     --member "serviceAccount:jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com" \
     --role roles/storage.insightsCollectorService
   

3. Anotasikan akun layanan Kubernetes dengan alamat email akun layanan IAM.

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

Mengonfigurasi Docker untuk melakukan autentikasi ke Artifact Registry

Konfigurasi Docker untuk melakukan autentikasi ke Artifact Registry sehingga dapat menarik image Pathways yang ada dalam daftar yang diizinkan:

gcloud auth login
gcloud auth configure-docker

Konversi checkpoint

Untuk mengonversi checkpoint Meta Llama 3.1-405B menjadi checkpoint inferensi int8 yang kompatibel dengan MaxText, selesaikan langkah-langkah dalam Konversi checkpoint dengan Llama3.1-405B. Deployment Anda menggunakan titik pemeriksaan dengan tanda load_parameters_path.

Buat bucket Cloud Storage untuk menyimpan file sementara Pathways

Buat bucket Cloud Storage untuk menyimpan file sementara Pathways Anda, seperti cache kompilasi:

export PATHWAYS_BUCKET=PATHWAYS_BUCKET
gcloud storage buckets create gs://$PATHWAYS_BUCKET

Men-deploy JetStream-MaxText dan Pathways

Men-deploy server model JetStream-MaxText dan Pathways.

Menghubungkan ke cluster GKE

gcloud container clusters get-credentials "${CLUSTER_NAME}" --project "${PROJECT_ID}"  \
    --location "${CONTROL_PLANE_LOCATION}"

Men-deploy LeaderWorkerSet (LWS) API

LWS adalah resource kustom yang dirancang untuk men-deploy dan mengelola aplikasi terdistribusi yang stateful, terutama yang memiliki arsitektur leader-worker. Hal ini sangat cocok untuk workload AI/ML di mana model besar di-shard dan ditayangkan di beberapa perangkat pada beberapa node.

VERSION=v0.6.1
kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/$VERSION/manifests.yaml

Tunggu hingga pengontrol LeaderWorkerSet tersedia sepenuhnya:

kubectl wait deploy/lws-controller-manager -n lws-system --for=condition=available --timeout=5m

Outputnya akan mirip dengan berikut ini:

deployment.apps/lws-controller-manager condition met

Validasi bahwa pengontrol LeaderWorkerSet berjalan di namespace lws-system:

kubectl get pod -n lws-system

Outputnya akan mirip dengan berikut ini:

NAME                          READY   STATUS    RESTARTS    AGE
lws-controller-manager-abcd   1/1     Running   0           40s
lws-controller-manager-efgh   1/1     Running   0           40s

Men-deploy manifes workload

1. Simpan manifes berikut sebagai jetstream-pathways-llama-3-1-405b-4x4.yaml:

   apiVersion: leaderworkerset.x-k8s.io/v1
   kind: LeaderWorkerSet
   metadata:
     name: jetstream-pathways
     annotations:
       leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     replicas: 1
     leaderWorkerTemplate:
       leaderTemplate:
         metadata:
           labels:
             app: jetstream-pathways
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 4x4
           tolerations:
           - key: "google.com/tpu"
             operator: "Exists"
             effect: "NoSchedule"
           containers:
           - name: pathways-proxy
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
             args:
             imagePullPolicy: Always
             ports:
             - containerPort: 38681
           - name: pathways-rm
             env:
             - name: HOST_ADDRESS
               value: "$(LWS_LEADER_ADDRESS)"
             - name: TPU_SKIP_MDS_QUERY
               value: "true"
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             args:
             - --server_port=38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             - --node_type=resource_manager
             - --instance_count=1
             - --instance_type=tpuv6e:4x4
             imagePullPolicy: Always
             ports:
             - containerPort: 38677
           - name: jax-tpu
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
             env:
             - name: LOG_LEVEL
               value: "INFO"
             args:
             - MaxText/configs/v5e/inference/llama3_405b_v5e-64.yml
             - model_name=llama3.1-405b
             - load_parameters_path=CHECKPOINT_PATH
             - max_prefill_predict_length=1024
             - max_target_length=2048
             - async_checkpointing=false
             - steps=1
             - ici_fsdp_parallelism=1
             - ici_autoregressive_parallelism=2
             - ici_tensor_parallelism=8
             - scan_layers=false
             - weight_dtype=bfloat16
             - per_device_batch_size=6
             - enable_single_controller=true
             - quantization=int8
             - quantize_kvcache=true
             - checkpoint_is_quantized=true
             - enable_model_warmup=true
             imagePullPolicy: Always
             ports:
             - containerPort: 9000
             startupProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 1
               initialDelaySeconds: 600
               failureThreshold: 10000
             livenessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 10
             readinessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 10
           - name: jetstream-http
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
             imagePullPolicy: Always
             ports:
             - containerPort: 8000
       size: 5
       workerTemplate:
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 4x4
           tolerations:
           - key: "google.com/tpu"
             operator: "Exists"
             effect: "NoSchedule"
           containers:
           - name: worker
             args:
             - --server_port=38679
             - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             imagePullPolicy: Always
             ports:
             - containerPort: 38679
             resources:
               limits:
                 google.com/tpu: "4"
   --- 
   apiVersion: v1
   kind: Service
   metadata:
     name: jetstream-svc
   spec:
     selector:
       app: jetstream-pathways
     ports:
     - protocol: TCP
       name: jetstream-http
       port: 8000
       targetPort: 8000

2. Tetapkan nilai kolom load_parameters_path ke jalur titik pemeriksaan yang dihasilkan dalam proses konversi titik pemeriksaan.

   • Untuk checkpoint bf16, jalur harus serupa dengan gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
   • Untuk checkpoint int8, tampilannya akan serupa dengan gs://OUTPUT_BUCKET_DIRECTORY/int8.

   Tetapkan nilai kolom gcs_scratch_location ke bucket Pathways yang Anda buat sebelumnya.

   perl -pi -e "s|CHECKPOINT_PATH|gs://${BUCKET_NAME}/int8|g" jetstream-pathways-llama-3-1-405b-4x4.yaml
   perl -pi -e "s|PATHWAYS_BUCKET|gs://${PATHWAYS_BUCKET}|g" jetstream-pathways-llama-3-1-405b-4x4.yaml
   

   Deskripsi untuk variabel dan setelan.

   Deskripsi berikut menjelaskan variabel dan setelan yang relevan untuk JetStream, Pathways, dan LeaderWorkerSet.

   • Jetstream
     • Container ini dibangun dari project JetStream OSS. Selain itu, IFRT Proxy juga merangkum library klien berbasis gRPC dari IFRT Proxy, yang menyediakan backend JAX yang dapat mendelegasikan program TPU yang sedang berjalan ke server Pathways Proxy.
     • Untuk mengetahui informasi selengkapnya tentang flag JetStream-MaxText, lihat Deskripsi flag Server JetStream-MaxText.
     • Untuk menggunakan checkpoint terkuantisasi, Anda dapat menetapkan flag quantization, quantize_kvcache, atau checkpoint_is_quantized. Jika titik pemeriksaan adalah bf16, Anda dapat menghapus tanda ini.
   • Jalur: untuk deskripsi komponen Jalur yang relevan, lihat Komponen jalur.
   • LeaderWorkerSet: untuk mengetahui informasi selengkapnya tentang setelan LeaderWorkerSet, lihat definisi API.

Terapkan manifes Deployment

Terapkan manifes untuk men-deploy server:

kubectl apply -f jetstream-pathways-llama-3-1-405b-4x4.yaml

Server model akan dimulai.

Memverifikasi mulai server model

Model 405B mungkin memerlukan waktu sekitar 10 hingga 20 menit untuk memulihkan titik pemeriksaan. Anda mungkin juga menunggu waktu tambahan selama pemanasan model jika mengaktifkan tanda enable_model_warmup.

kubectl logs -f jetstream-pathways-0 -c jax-tpu

Outputnya mirip dengan hal berikut ini:

2025-03-02 02:15:07,682 - JetstreamLogger - INFO - Initializing the driver with 1 prefill engines and 1 generate engines in interleaved mode
2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up prefill thread 0.
2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up transfer thread 0.
2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up generate thread 0.
2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up detokenize thread 0.
2025-03-02 02:15:07,685 - JetstreamLogger - INFO - Driver initialized.
...
...
...
INFO:     Started server process [7]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:9999 (Press CTRL+C to quit)

Menyajikan Llama 3.1-405b

Untuk menayangkan model Llama 3.1-405b, siapkan penerusan port:

kubectl port-forward svc/jetstream-svc 8000:8000

Penerusan port memungkinkan Anda mengakses Layanan dari luar cluster. Anda dapat mengakses Deployment JetStream-Pathways melalui Layanan ClusterIP GKE. Layanan ClusterIP hanya dapat dijangkau dari dalam cluster.

Berinteraksi dengan model

Di terminal baru, jalankan perintah berikut:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8000/generate \
--data \
'{
    "prompt": "What are the top 5 programming languages",
    "max_tokens": 200
}'

Permintaan awal dapat memerlukan waktu beberapa detik untuk diselesaikan karena pemanasan model. Outputnya akan mirip dengan berikut ini:

{
    "response": " for web development?\nThe top 5 programming languages for web development are:\n1. **JavaScript**: JavaScript is the most popular language for web development, used by over 90% of websites for client-side scripting. It's also popular for server-side programming with technologies like Node.js.\n2. **HTML/CSS**: HTML (Hypertext Markup Language) and CSS (Cascading Style Sheets) are not programming languages, but are essential for building websites. HTML is used for structuring content, while CSS is used for styling and layout.\n3. **Python**: Python is a popular language for web development, especially with frameworks like Django and Flask. It's known for its simplicity, flexibility, and large community of developers.\n4. **Java**: Java is a popular language for building enterprise-level web applications, especially with frameworks like Spring and Hibernate. It's known for its platform independence, strong security features, and large community of developers.\n5. **PHP**: PHP is a mature language for web"
}

Anda telah berhasil melakukan hal berikut:

1. Men-deploy server model JetStream dengan MaxText dan Pathways di GKE menggunakan TPU.
2. Membuat checkpoint int8 Llama 3.1-405B di gs://BUCKET_NAME.
3. Menyajikan dan berinteraksi dengan model.

Penayangan yang diuraikan

Penayangan yang tidak digabungkan adalah teknik untuk menayangkan LLM sambil memisahkan tahap pengisian awal dan decoding ke host yang berbeda. Pendekatan ini mengoptimalkan pemanfaatan resource, yang dapat meningkatkan throughput dan latensi.

• Pengisian otomatis: meneruskan perintah input untuk melakukan inisialisasi cache nilai kunci.

• Dekode: prosedur yang menghasilkan token output secara inkremental, satu token per langkah, dan satu nilai cache KV per iterasi.

1. Tetapkan variabel lingkungan default:

   export NODE_POOL_NAME=dis-v6e-8
   export NODE_POOL_SIZE=2
   export MACHINE_TYPE=ct6e-standard-4t
   export TOPOLOGY=2x4
   export WORKERS_PER_SLICE=2
   

2. Buat dua node pool yang menggunakan node v6e-8:

   for i in $(seq 1 ${NODE_POOL_SIZE}); do
     gcloud container node-pools create ${NODE_POOL_NAME}-${i}-np \
     --project=${PROJECT_ID} \
     --cluster=${CLUSTER_NAME} \
     --location=${CONTROL_PLANE_LOCATION} \
     --node-locations=${NODE_LOCATION} \
     --machine-type=${MACHINE_TYPE} \
     --num-nodes=${WORKERS_PER_SLICE} \
     --tpu-topology=${TOPOLOGY} \
     --scopes=cloud-platform \
     --workload-metadata=GCE_METADATA
   done
   

Konversi checkpoint

Untuk mengonversi checkpoint Meta Llama 2-70B menjadi checkpoint inferensi int8 yang kompatibel dengan MaxText, selesaikan langkah-langkah dalam Konversi checkpoint dengan Llama2-70B. Pilih Llama2-70B sebagai model Anda saat Anda menyetujui Persyaratan dan Ketentuan Meta. Deployment Anda menggunakan titik pemeriksaan dengan tanda load_parameters_path.

Ganti parameter berikut di file checkpoint-job.yaml Anda:

- --meta_url=META_URL
- --model_name=llama-2
- --model_path=Llama-2-70b-chat
- --output_directory=gs://BUCKET_NAME/maxtext/llama-2-70b

Pos pemeriksaan akan digunakan dalam deployment Anda dengan tanda load_parameters_path.

Men-deploy Jalur JetStream dengan penayangan yang tidak digabungkan

1. Simpan manifes berikut sebagai jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml:

   apiVersion: leaderworkerset.x-k8s.io/v1
   kind: LeaderWorkerSet
   metadata:
     name: jetstream-pathways
     annotations:
       leaderworkerset.sigs.k8s.io/subgroup-exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     replicas: 1
     leaderWorkerTemplate:
       subGroupPolicy:
         subGroupSize: 2
       leaderTemplate:
         metadata:
           labels:
             app: jetstream-pathways
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 2x4
           tolerations:
           - key: "google.com/tpu"
             operator: "Exists"
             effect: "NoSchedule"
           containers:
           - name: pathways-proxy
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
             args:
             - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
             - --server_port=38681
             - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
             - --xla_jf_auto_cross_replica_sharding=false
             - --xla_tpu_enable_windowed_einsum_for_reduce_scatter=false
             - --xla_tpu_enable_windowed_einsum_for_all_gather=false
             - --xla_tpu_prefer_latch_optimized_rhs_layouts=true
             - --xla_tpu_enable_experimental_fusion_cost_model=false
             - --xla_tpu_dot_dot_fusion_duplicated=false
             - --xla_tpu_dot_dot_fusion=true
             - --xla_jf_conv_input_fusion=true
             - --xla_jf_conv_output_fusion=true
             - --xla_tpu_rwb_fusion=false
             - --xla_tpu_copy_fusion_pad_unpad_ratio=0
             - --xla_tpu_licm_size_inflation_ratio=1
             - --xla_tpu_copy_elision_analysis_allowance=150000
             - --xla_tpu_copy_insertion_use_region_analysis_limit=10000
             - --xla_tpu_order_dot_after_layout=true
             - --xla_jf_rematerialization_percent_shared_memory_limit=100
             - --xla_tpu_use_repeated_instance_for_preferred_prefetch_time=true
             - --xla_tpu_enforce_prefetch_fifo_order=false
             - --xla_tpu_prefetch_interval_picker_size_override=6000000
             - --xla_tpu_async_copy_bandwidth_scaling_factor=1
             - --xla_tpu_nd_short_transfer_max_chunks=-1
             - --xla_tpu_enable_aggressive_broadcast_priority_update=true
             - --xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers=SQRT
             - --xla_tpu_memory_bound_loop_optimizer_options=enabled:true
             - --xla_tpu_enable_copy_fusion=true
             - --xla_tpu_enable_cross_program_prefetch_freeing=false
             - --xla_tpu_enable_dot_strength_reduction=true
             - --xla_tpu_layout_use_dot_grouping=false
             - --xla_tpu_msa_inefficient_use_to_copy_ratio=0.5
             - --xla_tpu_reduce_loop_fusion_dup_with_unfusable_user=false
             - --xla_tpu_vector_load_fusion_window=1024
             - --xla_tpu_vector_store_fusion_window=256
             - --xla_jf_conv_reshape_fusion=false
             - --xla_tpu_input_conv_multi_users=false
             - --xla_tpu_enable_multi_level_input_dot_dot_fusion=false
             - --xla_tpu_enable_multi_level_output_dot_dot_fusion=false
             - --xla_tpu_dot_dot_fusion_separable_convs_only=false
             - --xla_tpu_enable_multi_level_nested_loop_fusion=true
             - --xla_tpu_nested_dot_fusion=true
             - --xla_tpu_enable_multi_level_nested_dot_fusion=false
             - --xla_jf_enable_multi_output_fusion=true
             - --xla_tpu_use_lp_llo_scheduler_for_dot_dot_fusions=false
             - --xla_tpu_enable_flash_attention=true
             imagePullPolicy: Always
             ports:
             - containerPort: 38681
           - name: pathways-rm
             env:       
             - name: HOST_ADDRESS
               value: "$(LWS_LEADER_ADDRESS)"
             - name: TPU_SKIP_MDS_QUERY
               value: "true"
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             args:
             - --server_port=38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             - --node_type=resource_manager
             - --instance_count=2
             - --instance_type=tpuv6e:2x4
             imagePullPolicy: Always
             ports:
             - containerPort: 38677
           - name: jax-tpu
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
             args:
             - MaxText/configs/base.yml
             - tokenizer_path=assets/tokenizer.llama2
             - load_parameters_path=CHECKPOINT_PATH
             - max_prefill_predict_length=1024
             - max_target_length=2048
             - model_name=llama2-70b
             - ici_fsdp_parallelism=1
             - ici_autoregressive_parallelism=1
             - ici_tensor_parallelism=-1
             - scan_layers=false
             - weight_dtype=bfloat16
             - per_device_batch_size=27
             - checkpoint_is_quantized=true 
             - quantization=int8
             - quantize_kvcache=true
             - compute_axis_order=0,2,1,3
             - ar_cache_axis_order=0,2,1,3
             - stack_prefill_result_cache=True
             - inference_server=ExperimentalMaxtextDisaggregatedServer_8
             - inference_benchmark_test=True
             - enable_model_warmup=True
             env:
             - name: LOG_LEVEL
               value: "INFO"
             imagePullPolicy: Always
             securityContext:
               capabilities:
                 add: ["SYS_PTRACE", "NET_ADMIN", "SYS_TIME"]
             ports: 
             - containerPort: 9000
             startupProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 1
               initialDelaySeconds: 240
               failureThreshold: 10000
             livenessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 100
             readinessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 100
           - name: jetstream-http
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
             imagePullPolicy: Always
             ports:
             - containerPort: 8000
       size: 5
       workerTemplate:
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 2x4
           containers:
           - name: worker
             args:
             - --server_port=38679
             - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             imagePullPolicy: Always
             ports:
             - containerPort: 38679
             resources:
               limits:
                 google.com/tpu: "4"
   --- 
   apiVersion: v1
   kind: Service
   metadata:
     name: jetstream-svc
   spec:
     selector:
       app: jetstream-pathways
     ports:
     - protocol: TCP
       name: jetstream-http
       port: 8000
       targetPort: 8000

2. Tetapkan nilai kolom load_parameters_path ke jalur titik pemeriksaan yang dihasilkan dalam proses konversi titik pemeriksaan.

   • Untuk checkpoint bf16, jalur harus serupa dengan gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
   • Untuk checkpoint int8, tampilannya akan serupa dengan gs://OUTPUT_BUCKET_DIRECTORY/int8.

   Tetapkan nilai kolom gcs_scratch_location ke bucket Pathways yang Anda buat sebelumnya.

   perl -pi -e "s|CHECKPOINT_PATH|gs://${BUCKET_NAME}/maxtext/llama-2-70b/int8|g" jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
   perl -pi -e "s|PATHWAYS_BUCKET|gs://${PATHWAYS_BUCKET}|g" jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
   

3. Terapkan manifes:

   kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
   

   Server model mungkin memerlukan waktu beberapa saat untuk memulihkan checkpoint, bergantung pada ukuran checkpoint. Model 70B mungkin memerlukan waktu sekitar 8 menit untuk memulihkan checkpoint, termasuk pemanasan model. Anda dapat mengamati log lebih lanjut untuk menentukan titik kesiapan dengan memverifikasi dimulainya server model dan menayangkan model dengan menyiapkan penerusan port sehingga Anda dapat berinteraksi dengan model.

Anda telah berhasil melakukan hal berikut:

1. Men-deploy server model JetStream dengan MaxText dan Pathways di GKE menggunakan TPU dan penyajian yang terpisah.
2. Membuat checkpoint int8 Llama 2-70B di gs://BUCKET_NAME.
3. Menyajikan dan berinteraksi dengan model.

Memecahkan masalah

• Jika Anda mendapatkan pesan Empty reply from server, mungkin penampung belum selesai mendownload data model. Periksa log Pod lagi untuk melihat pesan Connected, yang menunjukkan bahwa model siap ditayangkan.
• Jika Anda melihat pesan Connection refused, pastikan penerusan port Anda aktif.

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

Agar akun Google Cloud Anda tidak dikenai biaya untuk resource yang Anda buat dalam panduan ini, jalankan perintah berikut dan ikuti perintahnya:

gcloud container clusters delete ${CLUSTER_NAME} --location=${CONTROL_PLANE_LOCATION}

gcloud iam service-accounts delete jetstream-pathways@${PROJECT_ID}.iam.gserviceaccount.com

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

Langkah berikutnya

• Pelajari cara menjalankan model Gemma di GKE dan cara menjalankan workload AI/ML yang dioptimalkan dengan kemampuan orkestrasi platform GKE.
• Pelajari lebih lanjut TPU di GKE.
• Jelajahi repositori GitHub JetStream.
• Jelajahi Vertex AI Model Garden.