本頁面由 Cloud Translation API 翻譯而成。

透過 JetStream 和 Pathways 在 GKE 上使用多主機 TPU 提供 LLM

標準

本指南說明如何使用多個節點的張量處理單元 (TPU)，在 Google Kubernetes Engine (GKE) 上提供最先進的大型語言模型 (LLM)，例如 Llama 3.1 405B。

本指南說明如何使用可攜式開放原始碼技術 (Kubernetes、JetStream、Pathways on Cloud 和 LeaderWorkerSet (LWS) API)，充分運用 GKE 的精細控制、擴充性、復原能力、可攜性和成本效益，在 GKE 上部署及提供 AI/ML 工作負載。

背景

大型語言模型越來越大，已無法再裝進單一主機 TPU 切片。如要進行 ML 推論，您可以使用 Pathways on Cloud，在 GKE 上跨多個互連的 TPU 節點，執行大規模的多主機推論。本指南將逐步說明如何使用多主機 TPU 配量佈建 GKE 叢集、使用 Pathways on Cloud 二進位檔、透過 MaxText 架構啟動 JetStream 伺服器，以及發出多主機推論要求。

透過 JetStream、MaxText 和 Pathways，在 GKE 上使用 TPU 提供 LLM，即可建構完善且可用於正式環境的服務解決方案，同時享有代管型 Kubernetes 的所有優點，包括成本效益、擴充性和高可用性。本節說明本教學課程中使用的主要技術。

關於 TPU

TPU 是 Google 開發的客製化特殊應用積體電路 (ASIC)，用於加速機器學習和 AI 模型，這些模型是使用 TensorFlow、PyTorch 和 JAX 等架構建構而成。

在 GKE 中使用 TPU 之前，建議您先完成下列學習路徑：

如要瞭解目前可用的 TPU 版本，請參閱 Cloud TPU 系統架構。
瞭解 GKE 中的 TPU。

本教學課程涵蓋 Llama 3.1-405B 模型服務。GKE 會在多主機 TPU v6e 節點上部署模型，並根據模型需求設定 TPU 拓撲，以低延遲服務提示。

雲端路徑

Pathways 是加速器的大規模自動化調度管理層，Pathways 的設計宗旨是探索新系統和機器學習研究構想，同時維持現有模型的頂尖效能。Pathways 可讓單一 JAX 用戶端程序協調一或多個大型 TPU 切片之間的運算，簡化涵蓋數百或數千個 TPU 晶片的機器學習運算。

JetStream

JetStream 是 Google 開發的開放原始碼推論服務架構。JetStream 可在 TPU 和 GPU 上執行高效能、高處理量和記憶體最佳化推論作業。JetStream 提供進階效能最佳化功能，包括持續批次處理、KV 快取最佳化和量化技術，可簡化 LLM 部署作業。JetStream 可啟用 PyTorch/XLA 和 JAX TPU 服務，進而提升效能。

MaxText

MaxText 是高效能、可擴充且可調整的 JAX LLM 實作項目，以 Flax、Orbax 和 Optax 等開放原始碼 JAX 程式庫為基礎建構而成。MaxText 的僅解碼器 LLM 實作是以 Python 編寫，它大量運用 XLA 編譯器，無須建構自訂核心即可達到高效能。

如要進一步瞭解 MaxText 支援的最新模型和參數大小，請參閱 MaxText 專案存放區。

Llama 3.1 405B

Llama 3.1 405B 是 Meta 的大型語言模型，適用於各種自然語言處理工作，包括生成文字、翻譯和回答問題。GKE 提供基礎架構，支援這類模型的分散式訓練和服務需求。

詳情請參閱 Llama 說明文件。

架構

本節說明本教學課程使用的 GKE 架構。此架構包含 GKE Standard 叢集，可佈建 TPU，並代管 JetStream 和 Pathways 元件，以部署及提供模型。

下圖顯示這個架構的元件：

GKE 叢集架構，其中包含 JetStream 和 Pathways 元件的多主機 TPU 節點集區。

這個架構包含下列元件：

GKE Standard 區域叢集。
多主機 TPU 配量節點集區，用於代管 JetStream 部署作業和 Pathways 元件。
Pathways resource manager 會管理加速器資源，並協調使用者工作的加速器分配作業。
Pathways client 會與 Pathways resource manager 協調，決定編譯後的程式要放在哪裡執行。
Pathways worker 會在加速器機器上執行運算，並透過 IFRT Proxy 伺服器將資料傳回工作負載。
IFRT proxy client 會實作 OSS Interim Framework Runtime (IFRT) API，並做為工作負載和 Pathways 元件之間的通訊橋樑。
IFRT proxy server 會接收 IFRT proxy client 的要求，並將要求轉送至 Pathways client，分配工作。
JetStream-Pathways 容器提供以 JAX 為基礎的推論伺服器，可接收推論要求，並將執行程序委派給 Pathways workers
Service 元件會將傳入流量分散到所有 JetStream HTTP 副本。
JetStream HTTP 是 HTTP 伺服器，可接受要求做為 JetStream 必要格式的包裝函式，並將要求傳送至 JetStream 的 GRPC 用戶端。

事前準備

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.

Go to project selector

Make sure 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.

Go to project selector

Make sure 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.
  前往「IAM」頁面
2. 選取專案。
3. 按一下「授予存取權」。
4. 在「New principals」(新增主體) 欄位中，輸入您的使用者 ID。這通常是 Google 帳戶的電子郵件地址。
5. 在「Select a role」(選取角色) 清單中，選取角色。
6. 如要授予其他角色，請按一下「新增其他角色」，然後新增每個其他角色。
7. 按一下 [Save]。

確認您有足夠的配額，可使用 16 個 TPU v6e PodSlice Lite 晶片。在本教學課程中，您將使用隨選執行個體。
確認您的 Google Cloud 專案已加入 Pathways 的許可清單。

取得模型存取權

如要存取 Meta Llama 3.1-405B 檢查點，以便部署至 GKE，請按照下列步驟操作：

簽署授權同意聲明協議。
前往 Meta Llama 下載頁面。
詳閱並接受範本《條款及細則》，並取得下載範本所需的網址。
如要下載模型檢查點，請找出適當模型的模型 ID。如需支援的型號及其 ID 清單，請參閱 llama CLI 說明文件。舉例來說，Llama 3.1-405B 模型可使用 Llama 3.1-405B-Instruct:bf16-mp16。

準備環境

在本教學課程中，您將使用 Cloud Shell 管理託管於Google Cloud的資源。Cloud Shell 已預先安裝本教學課程所需的軟體，包括 kubectl 和 gcloud CLI。

如要使用 Cloud Shell 設定環境，請按照下列步驟操作：

在 Google Cloud 控制台中，按一下Google Cloud 控制台中的「啟用 Cloud Shell」，啟動 Cloud Shell 工作階段。系統會在 Google Cloud 控制台的底部窗格啟動工作階段。
設定預設環境變數：
```
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 REGION=REGION
export LOCATION=LOCATION
export CLUSTER_VERSION=CLUSTER_VERSION
export MACHINE_TYPE=ct6e-standard-4t
export TPU_TYPE=v6e
export TOPOLOGY=4x4
export WORKERS_PER_SLICE=4
```
替換下列值：
- PROJECT_ID：您的 Google Cloud 專案 ID。
- CLUSTER_NAME：GKE 叢集名稱。
- BUCKET_NAME：Cloud Storage bucket 的名稱。您不需要指定 gs:// 前置字元。
- REGION：GKE 叢集、Cloud Storage 值區和 TPU 節點所在的區域。這個區域包含提供 TPU v6e 機器類型的區域 (例如 us-east1、us-east5、europe-west4、asia-northeast1 或 us-south1)。
- LOCATION：提供 TPU 資源的區域 (例如 us-east1-d)。
- CLUSTER_VERSION：GKE 版本，必須支援您想使用的機器類型。請注意，預設 GKE 版本可能無法用於目標 TPU。如需各 TPU 機型適用的最低 GKE 版本清單，請參閱「GKE 中的 TPU 支援情形」。
- MACHINE_TYPE：v6e 機器類型。
- TPU_TYPE：用於命名節點集區的前置字串 (v6e)。
- TOPOLOGY：TPU v6e 拓撲。
- WORKERS_PER_SLICE：每個節點集區或 TPU 配量的節點數量。

建立及設定 Google Cloud 資源

如要建立必要資源，請按照下列操作說明進行：

建立 GKE 叢集

建立地區 GKE Standard 叢集：

gcloud container clusters create CLUSTER_NAME \
    --project=PROJECT_ID \
    --cluster-version=CLUSTER_VERSION \
    --region=REGION \
    --scopes=cloud-platform \
    --machine-type=n2-standard-32

建立叢集可能需要幾分鐘的時間。

將 CLUSTER_VERSION 替換為適當的叢集版本。

建立一個具有 4x4 拓撲和四個節點的 TPU v6e 節點集區：

gcloud container node-pools create multihost-np \
--project=PROJECT_ID \
--region=REGION \
--node-locations=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

設定服務帳戶，以便存取 Storage 物件

設定 Kubernetes 服務帳戶，使其做為 IAM 服務帳戶。

為應用程式建立 IAM 服務帳戶：

gcloud iam service-accounts create jetstream-pathways

為 IAM 服務帳戶新增 IAM 政策繫結，以便管理 Cloud Storage。這是為了讓 IAM 服務帳戶存取儲存檢查點的儲存空間 bucket：

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

使用 IAM 服務帳戶的電子郵件地址，為 Kubernetes 服務帳戶加上註解。

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

設定 Docker 向 Artifact Registry 進行驗證

設定 Docker 向 Artifact Registry 進行驗證，以便提取已加入許可清單的 Pathways 映像檔：

gcloud auth login
gcloud auth configure-docker

查核點轉換

如要將 Meta Llama 3.1-405B 檢查點轉換為與 MaxText 相容的 int8 推論檢查點，請完成「使用 Llama3.1-405B 轉換檢查點」一文中的步驟。您的部署作業會使用 load_parameters_path 旗標檢查點。

建立 Cloud Storage 值區，儲存 Pathways 暫存檔案

建立 Cloud Storage bucket，用來儲存 Pathways 暫存檔案 (例如編譯快取)：

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

部署 JetStream-MaxText 和 Pathways

部署 JetStream-MaxText 和 Pathways 模型伺服器。

連線至 GKE 叢集

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

部署 LeaderWorkerSet (LWS) API

LWS 是專為部署及管理具狀態的分散式應用程式設計的自訂資源，特別是具有領導者-工作者架構的應用程式。特別適合用於 AI/ML 工作負載，因為大型模型會經過分片，並在多個節點的多部裝置上提供服務。

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

等待 LeaderWorkerSet 控制器完全可用：

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

畫面會顯示如下的輸出內容：

deployment.apps/lws-controller-manager condition met

確認 LeaderWorkerSet 控制器是否在 lws-system 命名空間中執行：

kubectl get pod -n lws-system

畫面會顯示如下的輸出內容：

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

部署工作負載資訊清單

將下列資訊清單儲存為 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

將 load_parameters_path 欄位的值設為檢查點轉換程序中產生的檢查點路徑。
- 如果是 bf16 檢查點，路徑應類似於 gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items。
- 如果是 int8 檢查點，則應類似於 gs://OUTPUT_BUCKET_DIRECTORY/int8。
將 gcs_scratch_location 欄位的值設為您先前建立的 Pathways 值區。
```
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
```
變數和設定的說明。

以下說明 JetStream、Pathways 和 LeaderWorkerSet 的相關變數和設定。
- Jetstream
  - 這個容器是以 OSS JetStream 專案為基礎建構而成。此外，它也封裝了 IFRT Proxy 的 gRPC 用戶端程式庫，提供可將 TPU 程式執行作業委派給 Pathways Proxy 伺服器的 JAX 後端。
  - 如要進一步瞭解 JetStream-MaxText 標記，請參閱 JetStream-MaxText 伺服器標記說明。
  - 如要使用量化檢查點，可以設定 quantization、quantize_kvcache 或 checkpoint_is_quantized 旗標。如果檢查點是 bf16，可以省略這些標記。
- 路徑：如需相關路徑元件的說明，請參閱「路徑元件」。
- LeaderWorkerSet：如要進一步瞭解 LeaderWorkerSet 設定，請參閱API 定義。

套用 Deployment 資訊清單

套用資訊清單來部署伺服器：

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

模型伺服器應會啟動。

確認模型伺服器啟動

405B 型號可能需要約 10 到 20 分鐘才能還原檢查點。如果您啟用 enable_model_warmup 旗標，模型暖機時可能需要等待額外時間。

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

輸出結果會與下列內容相似：

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)

提供 Llama 3.1-405b

如要提供 Llama 3.1-405b 模型，請設定通訊埠轉送：

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

透過通訊埠轉送，您可以從叢集外部存取 Service。您可以透過 GKE 的 ClusterIP 服務存取 JetStream-Pathways 部署作業。ClusterIP 服務只能從叢集內部存取。

與模型互動

在新的終端機中執行下列指令：

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

由於模型暖機，初始要求可能需要幾秒鐘才能完成。畫面會顯示如下的輸出內容：

{
    "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"
}

您已成功完成下列操作：

使用 TPU 在 GKE 上，透過 MaxText 和 Pathways 部署 JetStream 模型伺服器。
在 gs://BUCKET_NAME 建立 Llama 3.1-405B int8 檢查點。
提供模型並與之互動。

不匯總放送

分散式服務是一種技術，可將預填和解碼階段分割到不同主機，藉此提供大型語言模型服務。這種做法可提高資源使用率，進而改善總處理量和延遲。

預先填入：對輸入提示進行正向傳遞，以初始化鍵/值快取。
解碼：逐步產生輸出權杖的程序，每一步產生一個權杖，每次疊代產生一個 KV 快取值。

設定預設環境變數：

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

建立兩個使用 v6e-8 節點的節點集區：

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

查核點轉換

如要將 Meta Llama 2-70B 檢查點轉換為與 MaxText 相容的 int8 推論檢查點，請完成「使用 Llama2-70B 轉換檢查點」一節中的步驟。確認 Meta 條款及細則後，請選取 Llama2-70B 做為模型。您的部署作業會使用 load_parameters_path 旗標檢查點。

在 checkpoint-job.yaml 檔案中替換下列參數：

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

檢查點會與 load_parameters_path 旗標一起用於部署作業。

使用分離式服務部署 JetStream Pathways

將下列資訊清單儲存為 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

將 load_parameters_path 欄位的值設為檢查點轉換程序中產生的檢查點路徑。
- 如果是 bf16 檢查點，路徑應類似於 gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items。
- 如果是 int8 檢查點，則應類似於 gs://OUTPUT_BUCKET_DIRECTORY/int8。
將 gcs_scratch_location 欄位的值設為您先前建立的 Pathways 值區。
```
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
```
套用資訊清單：
```
kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
```
視檢查點大小而定，模型伺服器可能需要一些時間才能還原檢查點。還原 70B 模型的檢查點可能需要約 8 分鐘，包括模型暖機。您可以進一步觀察記錄，確認模型伺服器已啟動，並設定通訊埠轉送，與模型互動，判斷模型何時準備就緒。

您已成功完成下列操作：

使用 TPU 和分離式服務，在 GKE 上透過 MaxText 和 Pathways 部署 JetStream 模型伺服器。
在 gs://BUCKET_NAME 建立 Llama 2-70B int8 檢查點。
提供模型並與模型互動。

排解問題

如果收到 Empty reply from server 訊息，可能是容器尚未完成下載模型資料。再次檢查 Pod 的記錄，確認是否有 Connected 訊息，這表示模型已準備好提供服務。
如果看到 Connection refused 訊息，請確認連接埠轉送功能是否已啟用。

清除所用資源

如要避免系統向您的 Google Cloud 帳戶收取本教學課程中所用資源的相關費用，請刪除含有該項資源的專案，或者保留專案但刪除個別資源。

刪除已部署的資源

如要避免系統向您的 Google Cloud 帳戶收取本指南所建立資源的費用，請執行下列指令並按照提示操作：

gcloud container clusters delete CLUSTER_NAME --region=REGION

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

gcloud storage rm --recursive gs://BUCKET_NAME

後續步驟

瞭解如何在 GKE 上執行 Gemma 模型，以及如何運用 GKE 平台的自動化調度管理功能，執行最佳化的 AI/機器學習工作負載。
進一步瞭解 GKE 中的 TPU。
探索 JetStream GitHub 存放區。
探索 Vertex AI Model Garden。

透過 JetStream 和 Pathways 在 GKE 上使用多主機 TPU 提供 LLM 透過集合功能整理內容 你可以依據偏好儲存及分類內容。

背景

關於 TPU

雲端路徑

JetStream

MaxText

Llama 3.1 405B

架構

事前準備

Check for the roles

Grant the roles

取得模型存取權

準備環境

建立及設定 Google Cloud 資源

建立 GKE 叢集

設定服務帳戶，以便存取 Storage 物件

設定 Docker 向 Artifact Registry 進行驗證

查核點轉換

建立 Cloud Storage 值區，儲存 Pathways 暫存檔案

部署 JetStream-MaxText 和 Pathways

連線至 GKE 叢集

部署 LeaderWorkerSet (LWS) API

部署工作負載資訊清單

套用 Deployment 資訊清單

確認模型伺服器啟動

提供 Llama 3.1-405b

與模型互動

不匯總放送

查核點轉換

使用分離式服務部署 JetStream Pathways

排解問題

清除所用資源

刪除已部署的資源

後續步驟

透過 JetStream 和 Pathways 在 GKE 上使用多主機 TPU 提供 LLM