Halaman ini diterjemahkan oleh Cloud Translation API.

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 efektivitas 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 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 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 dibangun menggunakan framework seperti TensorFlow, PyTorch, dan JAX.

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

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

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

Pathways on 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 canggih 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 decoder MaxText ditulis dalam Python. Framework 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 Standar GKE yang menyediakan TPU dan menghosting komponen JetStream dan Pathways untuk men-deploy dan menayangkan model.

Diagram berikut menunjukkan komponen arsitektur ini:

Arsitektur cluster GKE dengan node pool TPU multi-host yang berisi komponen JetStream dan Pathways.

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.

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.

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.

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

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

In the Google Cloud console, go to the IAM page.
Go to IAM
Select the project.
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.
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

In the Google Cloud console, go to the IAM page.
Buka IAM
Pilih project.
Klik Grant access.
Di kolom New principals, masukkan ID pengguna Anda. ID pengguna biasanya adalah alamat email untuk Akun Google.
Di daftar Select a role, pilih peran.
Untuk memberikan peran tambahan, klik Add another role, lalu tambahkan setiap peran tambahan.
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.

Mendapatkan akses ke model

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

Tandatangani perjanjian izin lisensi.
Buka halaman download Meta Llama.
Tinjau dan setujui Persyaratan dan Ketentuan model serta dapatkan URL yang diperlukan untuk mendownload model.
Untuk mendownload checkpoint model, temukan ID 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:

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

Membuat cluster GKE

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.

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

Mengonfigurasi akun layanan untuk akses Objek Storage

Konfigurasi akun layanan Kubernetes agar bertindak sebagai akun layanan IAM.

Buat akun layanan IAM untuk aplikasi Anda:

gcloud iam service-accounts create jetstream-pathways

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

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

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 memiliki status, 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

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

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://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
  - Container ini dibangun dari project JetStream OSS. 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 menetapkan 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:

Men-deploy server model JetStream dengan MaxText dan Pathways di GKE menggunakan TPU.
Membuat checkpoint int8 Llama 3.1-405B di gs://BUCKET_NAME.
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 input perintah 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.

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

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

Checkpoint akan digunakan dalam deployment Anda dengan tanda load_parameters_path.

Men-deploy Jalur JetStream dengan penayangan yang dipisah

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

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|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
```
Terapkan manifes:
```
kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
```
Server model mungkin memerlukan waktu beberapa saat untuk memulihkan titik pemeriksaan, bergantung pada ukuran titik pemeriksaan. 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:

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

Memecahkan masalah

Jika Anda mendapatkan pesan Empty reply from server, kemungkinan 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.

Menyajikan LLM menggunakan TPU multi-host di GKE dengan JetStream dan Pathways Tetap teratur dengan koleksi Simpan dan kategorikan konten berdasarkan preferensi Anda.

Latar belakang

Tentang TPU

Pathways on Cloud

JetStream

MaxText

Llama 3.1 405B

Arsitektur

Sebelum memulai

Check for the roles

Grant the roles

Mendapatkan akses ke model

Menyiapkan lingkungan

Membuat dan mengonfigurasi resource Google Cloud

Membuat cluster GKE

Mengonfigurasi akun layanan untuk akses Objek Storage

Mengonfigurasi Docker untuk melakukan autentikasi ke Artifact Registry

Konversi checkpoint

Buat bucket Cloud Storage untuk menyimpan file sementara Pathways

Men-deploy JetStream-MaxText dan Pathways

Menghubungkan ke cluster GKE

Men-deploy LeaderWorkerSet (LWS) API

Men-deploy manifes workload

Terapkan manifes Deployment

Memverifikasi mulai server model

Menyajikan Llama 3.1-405b

Berinteraksi dengan model

Penayangan yang diuraikan

Konversi checkpoint

Men-deploy Jalur JetStream dengan penayangan yang dipisah

Memecahkan masalah

Pembersihan

Menghapus resource yang di-deploy

Langkah berikutnya

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