透過 Optimum TPU 在 GKE 上使用 TPU 提供開放原始碼模型

本教學課程說明如何使用 Hugging Face 的 Optimum TPU 服務架構，在 Google Kubernetes Engine (GKE) 上，透過張量處理單元 (TPU) 提供大型語言模型 (LLM) 開放原始碼模型。在本教學課程中，您會從 Hugging Face 下載開放原始碼模型，並使用執行 Optimum TPU 的容器，將模型部署至 GKE 標準叢集。

如果您在部署及提供 AI/機器學習工作負載時，需要代管型 Kubernetes 的精細控制、可擴充性、復原能力、可攜性和成本效益，這份指南就是您的入門資源。

本教學課程適用於 Hugging Face 生態系統中的生成式 AI 客戶、GKE 的新使用者或現有使用者、機器學習工程師、MLOps (DevOps) 工程師，或是對使用 Kubernetes 容器編排功能服務 LLM 有興趣的平台管理員。

提醒您，Google Cloud 提供多種 LLM 推論選項，包括 Vertex AI、GKE 和 Google Compute Engine 等服務，您可以在其中整合 JetStream、vLLM 和其他合作夥伴服務等服務程式庫。舉例來說，您可以使用 JetStream 從專案取得最新的最佳化設定。如果偏好使用 Hugging Face 選項，可以改用 Optimum TPU。

Optimum TPU 支援下列功能：

  • 持續批次處理
  • 權杖串流
  • 使用 Transformer 的貪婪搜尋法和多項式取樣法。

取得模型存取權

您可以使用 Gemma 2B 或 Llama3 8B 模型。本教學課程的重點是這兩個模型，但 Optimum TPU 支援更多模型。

Gemma 2B

如要存取 Gemma 模型並部署至 GKE，請先簽署授權同意聲明，然後產生 Hugging Face 存取權杖。

您必須簽署同意聲明協議，才能使用 Gemma。請按照下列步驟操作：

  1. 前往模型同意聲明頁面。
  2. 使用 Hugging Face 帳戶驗證同意聲明。
  3. 接受模型條款。

產生存取權杖

如果沒有，請產生新的 Hugging Face 權杖：

  1. 依序點選「Your Profile」(你的個人資料) >「Settings」(設定) >「Access Tokens」(存取權杖)。
  2. 按一下「New Token」。
  3. 指定您選擇的名稱，以及至少 Read 的角色。
  4. 按一下 [產生憑證]。
  5. 將產生的權杖複製到剪貼簿。

Llama3 8B

如要在 Hugging Face Repo 中使用 Llama3 8b，請務必簽署同意聲明。

產生存取權杖

如果沒有，請產生新的 Hugging Face 權杖：

  1. 依序點選「Your Profile」(你的個人資料) >「Settings」(設定) >「Access Tokens」(存取權杖)。
  2. 選取「New Token」。
  3. 指定您選擇的名稱，以及至少 Read 的角色。
  4. 選取「產生權杖」。
  5. 將產生的權杖複製到剪貼簿。

建立 GKE 叢集

建立具有 1 個 CPU 節點的 GKE Standard 叢集：

gcloud container clusters create CLUSTER_NAME \
    --project=PROJECT_ID \
    --num-nodes=1 \
    --location=ZONE

建立 TPU 節點集區

建立含有 1 個節點和 8 個晶片的 v5e TPU 節點集區：

gcloud container node-pools create tpunodepool \
    --location=ZONE \
    --num-nodes=1 \
    --machine-type=ct5lp-hightpu-8t \
    --cluster=CLUSTER_NAME

如果 TPU 資源可用，GKE 會佈建節點集區。如果 TPU 資源暫時無法使用，輸出內容會顯示 GCE_STOCKOUT 錯誤訊息。如要排解 TPU 缺貨錯誤，請參閱「TPU 資源不足，無法滿足 TPU 要求」。

設定 kubectl 與叢集通訊：

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${ZONE}

建構容器

執行 make 指令來建構映像檔

cd optimum-tpu && make tpu-tgi

將映像檔推送至 Artifact Registry

gcloud artifacts repositories create optimum-tpu --repository-format=docker --location=REGION_NAME && \
gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
docker image tag huggingface/optimum-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest && \
docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest

為 Hugging Face 憑證建立 Kubernetes 密鑰

建立包含 Hugging Face 權杖的 Kubernetes Secret：

kubectl create secret generic hf-secret \
  --from-literal=hf_api_token=${HF_TOKEN} \
  --dry-run=client -o yaml | kubectl apply -f -

部署 Optimum TPU

本教學課程使用 Kubernetes Deployment 部署 Optimum TPU。Deployment 是 Kubernetes API 物件，可讓您執行多個 Pod 副本，並將這些副本分散到叢集中的節點。

Gemma 2B

  1. 請將下列 Deployment 資訊清單儲存為 optimum-tpu-gemma-2b-2x4.yaml：

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: tgi-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: tgi-tpu
  template:
    metadata:
      labels:
        app: tgi-tpu
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      containers:
      - name: tgi-tpu
        image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
        args:
        - --model-id=google/gemma-2b
        - --max-concurrent-requests=4
        - --max-input-length=8191
        - --max-total-tokens=8192
        - --max-batch-prefill-tokens=32768
        - --max-batch-size=16
        securityContext:
            privileged: true
        env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
        ports:
        - containerPort: 80
        resources:
          limits:
            google.com/tpu: 8
        livenessProbe:
          httpGet:
            path: /health
            port: 80
          initialDelaySeconds: 300
          periodSeconds: 120

---
apiVersion: v1
kind: Service
metadata:
  name: service
spec:
  selector:
    app: tgi-tpu
  ports:
    - name: http
      protocol: TCP
      port: 8080
      targetPort: 80

    這份資訊清單說明瞭最佳化 TPU 部署作業，並在 TCP 通訊埠 8080 上使用內部負載平衡器。

  2. 套用資訊清單

    kubectl apply -f optimum-tpu-gemma-2b-2x4.yaml

Llama3 8B

  1. 將下列資訊清單儲存為 optimum-tpu-llama3-8b-2x4.yaml：

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: tgi-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: tgi-tpu
  template:
    metadata:
      labels:
        app: tgi-tpu
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      containers:
      - name: tgi-tpu
        image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
        args:
        - --model-id=meta-llama/Meta-Llama-3-8B
        - --max-concurrent-requests=4
        - --max-input-length=8191
        - --max-total-tokens=8192
        - --max-batch-prefill-tokens=32768
        - --max-batch-size=16
        env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
        ports:
        - containerPort: 80
        resources:
          limits:
            google.com/tpu: 8
        livenessProbe:
          httpGet:
            path: /health
            port: 80
          initialDelaySeconds: 300
          periodSeconds: 120
---
apiVersion: v1
kind: Service
metadata:
  name: service
spec:
  selector:
    app: tgi-tpu
  ports:
    - name: http
      protocol: TCP
      port: 8080
      targetPort: 80

    這份資訊清單說明瞭最佳化 TPU 部署作業，並在 TCP 通訊埠 8080 上使用內部負載平衡器。

  2. 套用資訊清單

    kubectl apply -f optimum-tpu-llama3-8b-2x4.yaml

查看執行中 Deployment 的記錄：

kubectl logs -f -l app=tgi-tpu

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

2024-07-09T22:39:34.365472Z  WARN text_generation_router: router/src/main.rs:295: no pipeline tag found for model google/gemma-2b
2024-07-09T22:40:47.851405Z  INFO text_generation_router: router/src/main.rs:314: Warming up model
2024-07-09T22:40:54.559269Z  INFO text_generation_router: router/src/main.rs:351: Setting max batch total tokens to 64
2024-07-09T22:40:54.559291Z  INFO text_generation_router: router/src/main.rs:352: Connected
2024-07-09T22:40:54.559295Z  WARN text_generation_router: router/src/main.rs:366: Invalid hostname, defaulting to 0.0.0.0

請務必先完整下載模型，再繼續進行下一個部分。

提供模型

設定通訊埠轉送至模型：

kubectl port-forward svc/service 8080:8080

使用 curl 與模型伺服器互動

驗證已部署的模型：

在新的終端機工作階段中，使用 curl 與模型對話：

curl 127.0.0.1:8080/generate     -X POST     -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":40}}'     -H 'Content-Type: application/json'

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

{"generated_text":"\n\nDeep learning is a subset of machine learning that uses artificial neural networks to learn from data.\n\nArtificial neural networks are inspired by the way the human brain works. They are made up of multiple layers"}