Google 會運用 AI 技術將內容翻譯成你偏好的語言，但可能會出錯。

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

標準

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

本指南將示範如何運用可攜式開放原始碼技術 (Kubernetes、JetStream、Pathways on Cloud 和 LeaderWorkerSet (LWS) API)，在 GKE 上部署及提供 AI/機器學習工作負載，同時善用 GKE 的精細控制、擴充性、韌性、可攜性和成本效益。

背景

大型語言模型越來越大，已無法裝進單一主機 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)，用於加速處理使用 TensorFlow、PyTorch 和 JAX 等架構建構的機器學習和 AI 模型。

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

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

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

Pathways on Cloud

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

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 用戶端。

事前準備

登入 Google Cloud 帳戶。如果您是 Google Cloud新手，歡迎建立帳戶，親自體驗產品的實際應用成效。新客戶還能獲得價值 $300 美元的免費抵免額，能用於執行、測試及部署工作負載。

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, if it is not already enabled.

Roles required to enable APIs

To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). 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, if it is not already enabled.

Roles required to enable APIs

Enable the API

請確認您在專案中具備下列角色： roles/container.admin、roles/iam.serviceAccountAdmin、roles/resourcemanager.projectIamAdmin
檢查角色
1. 前往 Google Cloud 控制台的「IAM」頁面。
  前往「IAM」頁面
2. 選取專案。
3. 在「主體」欄中，找出所有可識別您或您所屬群組的資料列。如要瞭解自己所屬的群組，請與管理員聯絡。
4. 針對指定或包含您的所有資料列，檢查「角色」欄，確認角色清單是否包含必要角色。
授予角色
1. 前往 Google Cloud 控制台的「IAM」頁面。
  前往「IAM」頁面
2. 選取專案。
3. 按一下「Grant access」(授予存取權)。
4. 在「New principals」(新增主體) 欄位中，輸入您的使用者 ID。這通常是指 Google 帳戶的電子郵件地址。
5. 按一下「Select a role」(選取角色)，然後搜尋角色。
6. 如要授予其他角色，請按一下「Add another role」(新增其他角色)，然後新增其他角色。
7. 按一下「Save」(儲存)。

確認您有足夠的配額，可使用十六個 TPU v6e PodSlice Lite 晶片。在本教學課程中，您會使用隨選執行個體。

取得模型存取權

如要存取 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 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
```
替換下列值：
- PROJECT_ID：您的 Google Cloud 專案 ID。
- CLUSTER_NAME：GKE 叢集的名稱。
- BUCKET_NAME：Cloud Storage bucket 的名稱。您不需要指定 gs:// 前置碼。
- CONTROL_PLANE_LOCATION：GKE 叢集、Cloud Storage bucket 和 TPU 節點所在的 Compute Engine 區域。該區域包含提供 TPU v6e 機器類型的可用區 (例如 us-east1、us-east5、europe-west4、asia-northeast1 或 us-south1)。
- NODE_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} \
    --location=${CONTROL_PLANE_LOCATION} \
    --scopes=cloud-platform \
    --machine-type=n2-standard-32

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

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

建立一個 TPU v6e 節點集區，每個節點集區都採用 4x4 拓撲，並包含四個節點：

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

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

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

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

gcloud iam service-accounts create jetstream-pathways

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

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

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

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

kubectl annotate serviceaccount default \
iam.gke.io/gcp-service-account=jetstream-pathways@${PROJECT_ID}.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 bucket 來儲存 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_NAME}" --project "${PROJECT_ID}"  \
    --location "${CONTROL_PLANE_LOCATION}"

部署 LeaderWorkerSet (LWS) API

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

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 bucket。
```
perl -pi -e "s|CHECKPOINT_PATH|gs://${BUCKET_NAME}/int8|g" jetstream-pathways-llama-3-1-405b-4x4.yaml
perl -pi -e "s|PATHWAYS_BUCKET|gs://${PATHWAYS_BUCKET}|g" jetstream-pathways-llama-3-1-405b-4x4.yaml
```
變數和設定的說明。

以下說明將解釋 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 檢查點。
提供模型並與模型互動。

分散式提供

分散式提供程序是一種 LLM 提供程序技術，可將預先填充和解碼階段分割到不同主機。這種做法可提高資源使用率，進而改善處理量和延遲時間。

預先填入：對輸入提示進行正向傳遞，以初始化鍵/值快取。
解碼：逐步產生輸出權杖的程序，每個步驟產生一個權杖，每次疊代產生一個 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_ID} \
  --cluster=${CLUSTER_NAME} \
  --location=${CONTROL_PLANE_LOCATION} \
  --node-locations=${NODE_LOCATION} \
  --machine-type=${MACHINE_TYPE} \
  --num-nodes=${WORKERS_PER_SLICE} \
  --tpu-topology=${TOPOLOGY} \
  --scopes=cloud-platform \
  --workload-metadata=GCE_METADATA
done

查核點轉換

如要將 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 bucket。
```
perl -pi -e "s|CHECKPOINT_PATH|gs://${BUCKET_NAME}/maxtext/llama-2-70b/int8|g" jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
perl -pi -e "s|PATHWAYS_BUCKET|gs://${PATHWAYS_BUCKET}|g" jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
```
套用資訊清單：
```
kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
```
視檢查點大小而定，模型伺服器可能需要一些時間才能還原檢查點。還原檢查點 (包括模型暖機) 可能需要約 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} --location=${CONTROL_PLANE_LOCATION}

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

Pathways on Cloud

JetStream

MaxText

Llama 3.1 405B

架構

事前準備

檢查角色

授予角色

取得模型存取權

準備環境

建立及設定 Google Cloud 資源

建立 GKE 叢集

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

設定 Docker 向 Artifact Registry 進行驗證

查核點轉換

建立 Cloud Storage bucket 來儲存 Pathways 暫存檔案

部署 JetStream-MaxText 和 Pathways

連線至 GKE 叢集

部署 LeaderWorkerSet (LWS) API

部署工作負載資訊清單

套用 Deployment 資訊清單

確認模型伺服器已啟動

提供 Llama 3.1-405b

與模型互動

分散式提供

查核點轉換

使用分散式服務部署 JetStream Pathways

排解問題

清除所用資源

刪除已部署的資源

後續步驟

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