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 API LeaderWorkerSet (LWS)—untuk men-deploy dan menyajikan workload AI/ML di GKE, dengan memanfaatkan kontrol terperinci, skalabilitas, ketahanan, portabilitas, dan efektivitas biaya GKE.

Penting: Untuk mendapatkan akses ke Pathways, hubungi perwakilan Akun Google Cloud Anda.

Latar belakang

Model bahasa besar telah bertambah ukurannya dan tidak lagi sesuai dengan slice TPU host tunggal. Untuk inferensi ML, Anda dapat menggunakan Pathways on Cloud untuk menjalankan inferensi multi-host berskala 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 di 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 dibangun 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.

Jalur di Cloud

Pathways adalah lapisan orkestrasi berskala besar untuk akselerator. Pathways dirancang secara eksplisit untuk memungkinkan eksplorasi sistem baru dan ide penelitian 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 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 PyTorch/XLA dan JAX untuk mengoptimalkan performa.

MaxText

MaxText adalah implementasi LLM JAX yang berperforma tinggi, skalabel, dan dapat disesuaikan, yang dibangun di atas library JAX open source seperti Flax, Orbax, dan Optax. Implementasi LLM khusus dekoder 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, sehingga mendistribusikan tugas.
• Penampung 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

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/resourcemanager.projectIamAdmin

  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 Berikan akses.

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

  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 Anda memiliki kuota yang cukup untuk enam belas chip TPU v6e PodSlice Lite. Dalam tutorial ini, Anda akan menggunakan instance sesuai permintaan.
  • Pastikan project Google Cloud Anda diizinkan untuk Pathways.

  Mendapatkan akses ke model

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

  1. Tandatangani 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 Activate 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 ini 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 untuk 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} \
       --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
       --role roles/storage.objectUser
     
     gcloud projects add-iam-policy-binding ${PROJECT} \
       --member "serviceAccount:jetstream-pathways@${PROJECT}.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}.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 di 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

  Deploy server model JetStream-MaxText dan Pathways.

  Menghubungkan ke cluster GKE

  gcloud container clusters get-credentials "${CLUSTER}" --project "${PROJECT}" --location "${ZONE}"
  

  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 disajikan 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 mirip dengan gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
     • Untuk titik pemeriksaan int8, nilainya 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://OUTPUT_BUCKET_DIRECTORY/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
       • Penampung ini dibangun dari project OSS JetStream. Selain itu, library ini 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 tanda JetStream-MaxText, lihat Deskripsi tanda Server JetStream-MaxText.
       • Untuk menggunakan checkpoint terkuantisasi, Anda dapat menyetel flag quantization, quantize_kvcache, atau checkpoint_is_quantized. Jika titik pemeriksaan adalah bf16, Anda dapat menghapus tanda ini.
     • Pathways: untuk deskripsi komponen Pathways yang relevan, lihat Komponen Pathways.
     • 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 juga dapat menunggu waktu tambahan selama pemanasan model jika Anda 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 membagi tahap pengisian awal dan decoding ke host yang berbeda. Pendekatan ini mengoptimalkan pemanfaatan resource, yang dapat meningkatkan throughput dan latensi.

  • Pengisian otomatis: meneruskan input perintah untuk melakukan inisialisasi cache key-value.

  • 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 \
       --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 di 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 dipisah

  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 mirip dengan gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
     • Untuk titik pemeriksaan int8, nilainya 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|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 menyajikan 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 kembali log Pod untuk menemukan 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 tidak menimbulkan biaya pada akun Google Cloud Anda 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.