本教學課程說明如何使用 Google Kubernetes Engine (GKE) 的張量處理單元 (TPU)，提供 Gemma 大型語言模型 (LLM) 服務。您將預先建構的容器連同 JetStream 和 MaxText 部署至 GKE。您也可以設定 GKE，在執行階段從 Cloud Storage 載入 Gemma 7B 權重
本教學課程的適用對象為機器學習 (ML) 工程師、平台管理員和操作員，以及有興趣使用 Kubernetes 容器自動化調度管理功能提供 LLM 服務的資料和 AI 專家。如要進一步瞭解我們在Google Cloud 內容中提及的常見角色和範例工作，請參閱「常見的 GKE Enterprise 使用者角色和工作」。
閱讀本頁面之前，請先熟悉下列概念：
- Autopilot 模式和 Standard 模式
- 目前可用的 TPU 版本 (Cloud TPU 系統架構)
- GKE 中的 TPU
背景
本節說明本教學課程中使用的重要技術。
Gemma
Gemma 是一組以開放授權發布的輕量級生成式人工智慧 (AI) 模型，您可以在應用程式、硬體、行動裝置或代管服務中執行這些 AI 模型。您可以使用 Gemma 模型生成文字，也可以調整這些模型，用於執行特定工作。
詳情請參閱 Gemma 說明文件。
TPU
TPU 是 Google 開發的客製化特殊應用積體電路 (ASIC)，用於加速機器學習和 AI 模型，這些模型是使用 TensorFlow、PyTorch 和 JAX 等架構建構而成。
本教學課程涵蓋 Gemma 7B 模型服務。GKE 會在單一主機 TPUv5e 節點上部署模型，並根據模型需求設定 TPU 拓撲，以低延遲方式提供提示。
JetStream
JetStream 是 Google 開發的開放原始碼推論服務架構。JetStream 可在 TPU 和 GPU 上執行高效能、高處理量和記憶體最佳化推論作業。這項技術提供進階效能最佳化功能，包括持續批次處理和量化技術，可協助您部署 LLM。JetStream 可讓 PyTorch/XLA 和 JAX TPU 服務達到最佳效能。
如要進一步瞭解這些最佳化作業，請參閱 JetStream PyTorch 和 JetStream MaxText 專案存放區。
MaxText
MaxText 是高效能、可擴充且可調整的 JAX LLM 實作項目，以 Flax、Orbax 和 Optax 等開放原始碼 JAX 程式庫為基礎建構而成。MaxText 的僅解碼器 LLM 實作是以 Python 編寫，它大量運用 XLA 編譯器，無須建構自訂核心即可達到高效能。
如要進一步瞭解 MaxText 支援的最新模型和參數大小，請參閱 MaxtText 專案存放區。
目標
- 準備 GKE Autopilot 或 Standard 叢集，並根據模型特性採用建議的 TPU 拓撲。
- 在 GKE 上部署 JetStream 元件。
- 取得並發布 Gemma 7B 指令調整模型。
- 提供已發布的模型並與之互動。
架構
本節說明本教學課程使用的 GKE 架構。此架構包含 GKE Autopilot 或 Standard 叢集，可佈建 TPU 並代管 JetStream 元件，以部署及提供模型。
下圖顯示這個架構的元件：
這個架構包含下列元件：
- GKE Autopilot 或 Standard 區域叢集。
- 兩個單一主機 TPU 配量節點集區，用於代管 JetStream 部署作業。
- Service 元件會將傳入流量分散到所有
JetStream HTTP副本。 JetStream HTTP是 HTTP 伺服器，可接受要求做為 JetStream 必要格式的包裝函式，並將要求傳送至 JetStream 的 GRPC 用戶端。Maxengine是 JetStream 伺服器，可透過持續批次處理執行推論。
事前準備
- 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.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the required API.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the required 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.
前往「IAM」頁面 - 選取專案。
- 按一下「授予存取權」。
-
在「New principals」(新增主體) 欄位中，輸入您的使用者 ID。 這通常是 Google 帳戶的電子郵件地址。
- 在「Select a role」(選取角色) 清單中，選取角色。
- 如要授予其他角色，請按一下 「新增其他角色」，然後新增每個其他角色。
- 按一下 [Save]。
-
取得模型存取權
如要存取 Gemma 模型並部署至 GKE，請先簽署授權同意聲明協議。
簽署授權同意聲明協議
您必須簽署同意聲明，才能使用 Gemma。 請按照以下步驟操作：
- 前往 Kaggle.com 的 Gemma 模型同意頁面。
- 如果尚未登入 Kaggle，請先登入。
- 按一下「要求存取權」。
- 在「選擇同意聲明帳戶」部分，選取「透過 Kaggle 帳戶驗證」，使用 Kaggle 帳戶表示同意。
- 接受模型條款及細則。
產生存取權杖
如要透過 Kaggle 存取模型，您需要 Kaggle API 權杖。
如要產生新權杖 (如果沒有)，請按照下列步驟操作：
- 在瀏覽器中前往 Kaggle 設定。
- 在「API」部分下方，按一下「Create New Token」。
系統會下載名為「kaggle.json」的檔案。
準備環境
在本教學課程中，您將使用 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替換下列值：
- PROJECT_ID：您的 Google Cloud 專案 ID。
- CLUSTER_NAME：GKE 叢集名稱。
- BUCKET_NAME：Cloud Storage bucket 的名稱。您不需要指定
gs://前置字元。 - REGION：GKE 叢集、Cloud Storage 值區和 TPU 節點所在的區域。這個區域包含提供 TPU v5e 機器類型的可用區域 (例如
us-west1、us-west4、us-central1、us-east1、us-east5或europe-west4)。如果是 Autopilot 叢集，請確保您所選區域有足夠的 TPU v5e 區域資源。 - (僅限 Standard 叢集) LOCATION：提供 TPU 資源的區域 (例如
us-west4-a)。如果是 Autopilot 叢集，您不需要指定區域，只要指定區域即可。 CLUSTER_VERSION：GKE 版本，必須支援您想使用的機器類型。請注意，預設 GKE 版本可能無法用於目標 TPU。如需各 TPU 機型適用的最低 GKE 版本清單，請參閱「GKE 中的 TPU 支援情形」。
建立及設定 Google Cloud 資源
請按照下列操作說明建立必要資源。
建立 GKE 叢集
您可以在 GKE Autopilot 或 Standard 叢集的 TPU 上提供 Gemma 服務。建議您使用 Autopilot 叢集，享受全代管 Kubernetes 體驗。如要為工作負載選擇最合適的 GKE 作業模式，請參閱「選擇 GKE 作業模式」。
Autopilot
在 Cloud Shell 中執行下列指令：
gcloud container clusters create-auto ${CLUSTER_NAME} \
--project=${PROJECT_ID} \
--region=${REGION} \
--cluster-version=${CLUSTER_VERSION}
標準
建立使用 Workload Identity Federation for GKE 的地區 GKE Standard 叢集。
gcloud container clusters create ${CLUSTER_NAME} \ --enable-ip-alias \ --machine-type=e2-standard-4 \ --num-nodes=2 \ --cluster-version=${CLUSTER_VERSION} \ --workload-pool=${PROJECT_ID}.svc.id.goog \ --location=${REGION}建立叢集可能需要幾分鐘的時間。
執行下列指令，為叢集建立節點集區：
gcloud container node-pools create gemma-7b-tpu-nodepool \ --cluster=${CLUSTER_NAME} \ --machine-type=ct5lp-hightpu-8t \ --project=${PROJECT_ID} \ --num-nodes=2 \ --region=${REGION} \ --node-locations=${LOCATION}GKE 會建立具有
2x4拓撲和兩個節點的 TPU v5e 節點集區。
建立 Cloud Storage 值區
在 Cloud Shell 中執行下列指令：
gcloud storage buckets create gs://${BUCKET_NAME} --location=${REGION}
這會建立 Cloud Storage 值區，用於儲存從 Kaggle 下載的模型檔案。
將存取權杖上傳至 Cloud Shell
在 Cloud Shell 中，您可以將 Kaggle API 權杖上傳至專案： Google Cloud
- 在 Cloud Shell 中，依序點選「更多」>「上傳」。
- 選取「檔案」，然後按一下「選擇檔案」。
- 開啟
kaggle.json檔案。 - 按一下「上傳」。
為 Kaggle 憑證建立 Kubernetes Secret
在 Cloud Shell 中執行下列操作：
設定
kubectl與叢集通訊：gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${REGION}建立 Secret 來儲存 Kaggle 憑證：
kubectl create secret generic kaggle-secret \ --from-file=kaggle.json
使用 Workload Identity Federation for GKE 設定工作負載存取權
將 Kubernetes ServiceAccount 指派給應用程式，並將該 Kubernetes ServiceAccount 設定為 IAM 服務帳戶。
為應用程式建立 IAM 服務帳戶：
gcloud iam service-accounts create wi-jetstream為 IAM 服務帳戶新增 IAM 政策繫結，以便管理 Cloud Storage：
gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member "serviceAccount:wi-jetstream@${PROJECT_ID}.iam.gserviceaccount.com" \ --role roles/storage.objectUser gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member "serviceAccount:wi-jetstream@${PROJECT_ID}.iam.gserviceaccount.com" \ --role roles/storage.insightsCollectorService在兩個服務帳戶之間新增 IAM 政策繫結，允許 Kubernetes ServiceAccount模擬 IAM 服務帳戶。這個繫結可讓 Kubernetes ServiceAccount 做為 IAM 服務帳戶：
gcloud iam service-accounts add-iam-policy-binding wi-jetstream@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/iam.workloadIdentityUser \ --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"使用 IAM 服務帳戶的電子郵件地址註解 Kubernetes 服務帳戶：
kubectl annotate serviceaccount default \ iam.gke.io/gcp-service-account=wi-jetstream@${PROJECT_ID}.iam.gserviceaccount.com
轉換模型檢查點
在本節中，您會建立 Job 來執行下列操作：
- 從 Kaggle 下載基礎 Orbax 檢查點。
- 將檢查點上傳至 Cloud Storage bucket。
- 將檢查點轉換為與 MaxText 相容的檢查點。
- 取消掃描要用於服務的檢查點。
部署模型檢查點轉換作業
請按照下列操作說明下載及轉換 Gemma 7B 模型檢查點檔案。本教學課程使用 Kubernetes Job。Kubernetes 中的 Job 控制器會建立一或多個 Pod，並確保這些 Pod 成功執行特定工作。
建立下列資訊清單做為
job-7b.yaml。apiVersion: batch/v1 kind: Job metadata: name: data-loader-7b spec: ttlSecondsAfterFinished: 30 template: spec: restartPolicy: Never containers: - name: inference-checkpoint image: us-docker.pkg.dev/cloud-tpu-images/inference/inference-checkpoint:v0.2.4 args: - -b=BUCKET_NAME - -m=google/gemma/maxtext/7b-it/2 volumeMounts: - mountPath: "/kaggle/" name: kaggle-credentials readOnly: true resources: requests: google.com/tpu: 8 limits: google.com/tpu: 8 nodeSelector: cloud.google.com/gke-tpu-topology: 2x4 cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice volumes: - name: kaggle-credentials secret: defaultMode: 0400 secretName: kaggle-secret套用資訊清單：
kubectl apply -f job-7b.yaml等待排定 Job 的 Pod 開始執行：
kubectl get pod -w輸出內容大致如下，這可能需要幾分鐘：
NAME READY STATUS RESTARTS AGE data-loader-7b-abcd 0/1 ContainerCreating 0 28s data-loader-7b-abcd 1/1 Running 0 51s如果是 Autopilot 叢集，佈建所需的 TPU 資源可能需要幾分鐘。
查看工作的記錄：
kubectl logs -f jobs/data-loader-7b工作完成後，輸出內容會類似以下內容：
Successfully generated decode checkpoint at: gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items + echo -e '\nCompleted unscanning checkpoint to gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items' Completed unscanning checkpoint to gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items
部署 JetStream
在本節中，您將部署 JetStream 容器，以提供 Gemma 模型。
請按照下列操作說明部署 Gemma 7B 指令微調模型。本教學課程使用 Kubernetes Deployment。Deployment 是 Kubernetes API 物件，可讓您執行多個 Pod 副本，並將這些副本分散到叢集的節點中。
請將下列 Deployment 資訊清單儲存為
jetstream-gemma-deployment.yaml：apiVersion: apps/v1 kind: Deployment metadata: name: maxengine-server spec: replicas: 1 selector: matchLabels: app: maxengine-server template: metadata: labels: app: maxengine-server spec: nodeSelector: cloud.google.com/gke-tpu-topology: 2x4 cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice containers: - name: maxengine-server image: us-docker.pkg.dev/cloud-tpu-images/inference/maxengine-server:v0.2.2 args: - model_name=gemma-7b - tokenizer_path=assets/tokenizer.gemma - per_device_batch_size=4 - max_prefill_predict_length=1024 - max_target_length=2048 - async_checkpointing=false - ici_fsdp_parallelism=1 - ici_autoregressive_parallelism=-1 - ici_tensor_parallelism=1 - scan_layers=false - weight_dtype=bfloat16 - load_parameters_path=gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items - prometheus_port=PROMETHEUS_PORT ports: - containerPort: 9000 resources: requests: google.com/tpu: 8 limits: google.com/tpu: 8 - name: jetstream-http image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.2 ports: - containerPort: 8000 --- apiVersion: v1 kind: Service metadata: name: jetstream-svc spec: selector: app: maxengine-server ports: - protocol: TCP name: jetstream-http port: 8000 targetPort: 8000 - protocol: TCP name: jetstream-grpc port: 9000 targetPort: 9000資訊清單會設定下列重要屬性：
tokenizer_path：模型權杖化工具的路徑。load_parameters_path：Cloud Storage bucket 中儲存檢查點的路徑。per_device_batch_size：每個裝置的解碼批次大小，其中一個 TPU 晶片等於一個裝置。max_prefill_predict_length：執行自動迴歸時，預填內容的長度上限。max_target_length：序列長度上限。model_name：模型名稱 (gemma-7b)。ici_fsdp_parallelism：完全分片資料平行處理 (FSDP) 的分片數量。ici_tensor_parallelism：張量平行處理的資料分割數量。ici_autoregressive_parallelism：自迴歸平行處理的分片數量。prometheus_port：公開 Prometheus 指標的通訊埠。如果不需要指標，請移除這個引數。scan_layers：掃描圖層布林值標記 (布林值)。weight_dtype：權重資料型別 (bfloat16)。
套用資訊清單：
kubectl apply -f jetstream-gemma-deployment.yaml驗證部署作業：
kubectl get deployment輸出結果會與下列內容相似：
NAME READY UP-TO-DATE AVAILABLE AGE maxengine-server 2/2 2 2 ##s如果是 Autopilot 叢集，佈建所需的 TPU 資源可能需要幾分鐘。
查看 HTTP 伺服器記錄，確認模型已載入並編譯。 伺服器可能需要幾分鐘才能完成這項作業。
kubectl logs deploy/maxengine-server -f -c jetstream-http輸出結果會與下列內容相似：
kubectl logs deploy/maxengine-server -f -c jetstream-http INFO: Started server process [1] INFO: Waiting for application startup. INFO: Application startup complete. INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)查看 MaxEngine 記錄，確認編譯已完成。
kubectl logs deploy/maxengine-server -f -c maxengine-server輸出結果會與下列內容相似：
2024-03-29 17:09:08,047 - jax._src.dispatch - DEBUG - Finished XLA compilation of jit(initialize) in 0.26236414909362793 sec 2024-03-29 17:09:08,150 - root - INFO - ---------Generate params 0 loaded.---------
提供模型
在本節中，您將與模型互動。
設定通訊埠轉送
您可以透過在上一個步驟中建立的 ClusterIP 服務存取 JetStream 部署作業。ClusterIP 服務只能從叢集內部存取。因此，如要從叢集外部存取 Service，請完成下列步驟：
如要建立通訊埠轉送工作階段，請執行下列指令：
kubectl port-forward svc/jetstream-svc 8000:8000
使用 curl 與模型互動
開啟新的終端機並執行下列指令，確認您可以存取 JetStream HTTP 伺服器：
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": "\nfor data science in 2023?\n\n**1. Python:**\n- Widely used for data science due to its simplicity, readability, and extensive libraries for data wrangling, analysis, visualization, and machine learning.\n- Popular libraries include pandas, scikit-learn, and matplotlib.\n\n**2. R:**\n- Statistical programming language widely used for data analysis, visualization, and modeling.\n- Popular libraries include ggplot2, dplyr, and caret.\n\n**3. Java:**\n- Enterprise-grade language with strong performance and scalability.\n- Popular libraries include Spark, TensorFlow, and Weka.\n\n**4. C++:**\n- High-performance language often used for data analytics and machine learning models.\n- Popular libraries include TensorFlow, PyTorch, and OpenCV.\n\n**5. SQL:**\n- Relational database language essential for data wrangling and querying large datasets.\n- Popular tools" }
(選用) 透過 Gradio 聊天介面與模型互動
在本節中，您將建構網頁聊天應用程式，與經過指令微調的模型互動。
Gradio 是 Python 程式庫，內含 ChatInterface 包裝函式，可為聊天機器人建立使用者介面。
部署即時通訊介面
在 Cloud Shell 中，將下列資訊清單儲存為
gradio.yaml：apiVersion: apps/v1 kind: Deployment metadata: name: gradio labels: app: gradio spec: replicas: 1 selector: matchLabels: app: gradio template: metadata: labels: app: gradio spec: containers: - name: gradio image: us-docker.pkg.dev/google-samples/containers/gke/gradio-app:v1.0.3 resources: requests: cpu: "512m" memory: "512Mi" limits: cpu: "1" memory: "512Mi" env: - name: CONTEXT_PATH value: "/generate" - name: HOST value: "http://jetstream-svc:8000" - name: LLM_ENGINE value: "max" - name: MODEL_ID value: "gemma" - name: USER_PROMPT value: "<start_of_turn>user\nprompt<end_of_turn>\n" - name: SYSTEM_PROMPT value: "<start_of_turn>model\nprompt<end_of_turn>\n" ports: - containerPort: 7860 --- apiVersion: v1 kind: Service metadata: name: gradio spec: selector: app: gradio ports: - protocol: TCP port: 8080 targetPort: 7860 type: ClusterIP套用資訊清單：
kubectl apply -f gradio.yaml等待部署作業完成：
kubectl wait --for=condition=Available --timeout=300s deployment/gradio
使用即時通訊介面
在 Cloud Shell 中執行下列指令：
kubectl port-forward service/gradio 8080:8080這項操作會建立從 Cloud Shell 到 Gradio 服務的通訊埠轉送。
按一下 Cloud Shell 工作列右上方的「Web Preview」
按鈕。按一下「透過以下通訊埠預覽：8080」。瀏覽器會開啟新分頁。使用 Gradio 對話介面與 Gemma 互動。新增提示，然後按一下「提交」。
排解問題
- 如果收到
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 wi-jetstream@PROJECT_ID.iam.gserviceaccount.com
gcloud storage rm --recursive gs://BUCKET_NAME
後續步驟
- 瞭解如何在 GKE 上執行 Gemma 模型，以及如何運用 GKE 平台的自動化調度管理功能，執行最佳化的 AI/機器學習工作負載。
- 進一步瞭解 GKE 中的 TPU。
- 探索 JetStream GitHub 存放區。
- 探索 Vertex AI Model Garden。