透過 Saxml 在 GKE 上使用多主機 TPU 提供 LLM

本教學課程說明如何使用 Google Kubernetes Engine (GKE) 上的多主機 TPU 配量節點集區，透過 Saxml 部署及提供大型語言模型 (LLM)，打造可擴充的高效率架構。

背景

Saxml 是實驗性系統，可為 Paxml、JAX 和 PyTorch 架構提供服務。您可以使用 TPU，透過這些架構加速資料處理。為展示如何在 GKE 中部署 TPU，本教學課程會提供 175B LmCloudSpmd175B32Test 測試模型。GKE 會在兩個 v5e TPU 配量節點集區上部署這個測試模型，拓撲分別為 4x8。

為正確部署測試模型，系統已根據模型大小定義 TPU 拓撲。由於 N 十億個 16 位元模型大約需要 2 倍 (2xN) GB 的記憶體，因此 175B LmCloudSpmd175B32Test 模型大約需要 350 GB 的記憶體。TPU v5e 單一 TPU 晶片有 16 GB。如要支援 350 GB，GKE 需要 21 個 v5e TPU 晶片 (350/16= 21)。根據 TPU 設定的對應，本教學課程適用的 TPU 設定如下：

  • 機器類型：ct5lp-hightpu-4t
  • 拓撲：4x8 (32 個 TPU 晶片)

在 GKE 中部署 TPU 時，為模型服務選取合適的 TPU 拓撲非常重要。詳情請參閱「規劃 TPU 設定」。

目標

本教學課程適用於 MLOps 或 DevOps 工程師，以及想要使用 GKE 自動化調度管理功能提供資料模型的平台管理員。

本教學課程包含下列步驟：

  1. 使用 GKE Standard 叢集準備環境。叢集有兩個 v5e TPU 配量節點集區，拓撲為 4x8。
  2. 部署 Saxml。Saxml 需要管理員伺服器、一組做為模型伺服器的 Pod、預先建構的 HTTP 伺服器和負載平衡器。
  3. 使用 Saxml 提供 LLM。

下圖顯示下列教學課程實作的架構：

GKE 上的多主機 TPU 架構。
圖：GKE 上多主機 TPU 的架構範例。

事前準備

  • In the Trusted Cloud console, on the project selector page, select or create a Trusted Cloud project.

    Go to project selector

  • Make sure that billing is enabled for your Trusted 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/iam.policyAdmin

    Check for the roles

    1. In the Trusted 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 Trusted Cloud console, go to the IAM page.

      前往「IAM」頁面
    2. 選取專案。
    3. 按一下「授予存取權」。

    4. 在「New principals」(新增主體) 欄位中，輸入您的使用者 ID。 這通常是工作團隊身分集區中的使用者 ID。詳情請參閱「 在 IAM 政策中代表工作人員集區使用者」，或與管理員聯絡。

    5. 在「Select a role」(選取角色) 清單中，選取角色。
    6. 如要授予其他角色，請按一下 「新增其他角色」，然後新增每個其他角色。
    7. 按一下 [Save]。

準備環境

  1. 在 Trusted Cloud 控制台中啟動 Cloud Shell 執行個體：
    開啟 Cloud Shell

  2. 設定預設環境變數：

      gcloud config set project PROJECT_ID
  export PROJECT_ID=$(gcloud config get project)
  export ZONE=COMPUTE_ZONE
  export BUCKET_NAME=PROJECT_ID-gke-bucket

    替換下列值：

    在此指令中，BUCKET_NAME 指定儲存 Saxml 管理員伺服器設定的儲存空間值區名稱。 Trusted Cloud

建立 GKE Standard 叢集

使用 Cloud Shell 執行下列操作：

  1. 建立使用 Workload Identity Federation for GKE 的 Standard 叢集：

    gcloud container clusters create saxml \
    --zone=${ZONE} \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --cluster-version=VERSION \
    --num-nodes=4

    將 VERSION 替換為 GKE 版本號碼。GKE 1.27.2-gke.2100 以上版本支援 TPU v5e。詳情請參閱「GKE 中的 TPU 可用性」。

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

  2. 建立名為 tpu1 的第一個節點集區：

    gcloud container node-pools create tpu1 \
    --zone=${ZONE} \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=4x8 \
    --num-nodes=8 \
    --cluster=saxml

    --num-nodes 旗標的值計算方式為：將 TPU 拓撲除以每個 TPU 配量的 TPU 晶片數量。在本例中，公式為 (4 * 8) / 4。

  3. 建立名為 tpu2 的第二個節點集區：

    gcloud container node-pools create tpu2 \
    --zone=${ZONE} \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=4x8 \
    --num-nodes=8 \
    --cluster=saxml

    --num-nodes 旗標的值計算方式為：將 TPU 拓撲除以每個 TPU 配量的 TPU 晶片數量。在本例中，公式為 (4 * 8) / 4。

您已建立下列資源：

  • 具有四個 CPU 節點的標準叢集。
  • 兩個具有 4x8 拓撲的 v5e TPU 配量節點集區。每個節點集區代表八個 TPU 配量節點，每個節點有 4 個 TPU 晶片。

175B 模型必須在多主機 v5e TPU 切片上提供服務，且至少要有 4x8 拓撲切片 (32 個 v5e TPU 晶片)。

建立 Cloud Storage 值區

建立 Cloud Storage bucket，儲存 Saxml 管理員伺服器設定。執行中的管理員伺服器會定期儲存其狀態和已發布模型的詳細資料。

在 Cloud Shell 中執行下列指令：

gcloud storage buckets create gs://${BUCKET_NAME}

使用 Workload Identity Federation for GKE 設定工作負載存取權

將 Kubernetes ServiceAccount 指派給應用程式，並將該 Kubernetes ServiceAccount 設定為 IAM 服務帳戶。

  1. 設定 kubectl 與叢集通訊：

    gcloud container clusters get-credentials saxml --zone=${ZONE}

  2. 為應用程式建立要使用的 Kubernetes ServiceAccount：

    kubectl create serviceaccount sax-sa --namespace default

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

    gcloud iam service-accounts create sax-iam-sa

  4. 為 IAM 服務帳戶新增 IAM 政策繫結，允許讀取及寫入 Cloud Storage：

    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
  --member "serviceAccount:sax-iam-sa@${PROJECT_ID}.s3ns-system.iam.gserviceaccount.com" \
  --role roles/storage.admin

  5. 在兩個服務帳戶之間新增 IAM 政策繫結，允許 Kubernetes ServiceAccount模擬 IAM 服務帳戶。這個繫結可讓 Kubernetes ServiceAccount 做為 IAM 服務帳戶，以便 Kubernetes ServiceAccount 讀取及寫入 Cloud Storage。

    gcloud iam service-accounts add-iam-policy-binding sax-iam-sa@${PROJECT_ID}.s3ns-system.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/sax-sa]"

  6. 註解 Kubernetes 服務帳戶，並提供 IAM 服務帳戶的電子郵件地址。這樣範例應用程式就能知道要使用哪個服務帳戶存取服務 Trusted Cloud 。因此，當應用程式使用任何標準 Google API 用戶端程式庫存取 Trusted Cloud 服務時，會使用該 IAM 服務帳戶。

    kubectl annotate serviceaccount sax-sa \
  iam.gke.io/gcp-service-account=sax-iam-sa@${PROJECT_ID}.s3ns-system.iam.gserviceaccount.com

部署 Saxml

在本節中，您將部署 Saxml 管理員伺服器和 Saxml 模型伺服器。

部署 Saxml 管理員伺服器

  1. 建立下列 sax-admin-server.yaml 資訊清單：

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-admin-server
spec:
  replicas: 1
  selector:
    matchLabels:
      app: sax-admin-server
  template:
    metadata:
      labels:
        app: sax-admin-server
    spec:
      hostNetwork: false
      serviceAccountName: sax-sa
      containers:
      - name: sax-admin-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.1.0
        securityContext:
          privileged: true
        ports:
        - containerPort: 10000
        env:
        - name: GSBUCKET
          value: BUCKET_NAME

  2. 將 BUCKET_NAME 替換為您先前建立的 Cloud Storage：

    perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-admin-server.yaml

  3. 套用資訊清單：

    kubectl apply -f sax-admin-server.yaml

  4. 確認管理員伺服器 Pod 是否已啟動並執行：

    kubectl get deployment

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

    NAME               READY   UP-TO-DATE   AVAILABLE   AGE
sax-admin-server   1/1     1            1           52s

部署 Saxml 模型伺服器

在多主機 TPU 節點中執行的工作負載，需要每個 Pod 都有穩定的網路 ID，才能在同一個 TPU 節點中探索對等互連。如要定義這些 ID，請使用 IndexedJob、StatefulSet (搭配無頭服務) 或 JobSet，這會自動為屬於 JobSet 的所有 Job 建立無頭服務。Jobset 是工作負載 API，可讓您以單元形式管理一組 Kubernetes Job。JobSet 最常見的用途是分散式訓練，但您也可以用來執行批次工作負載。

以下章節說明如何使用 JobSet 管理多個模型伺服器 Pod 群組。

  1. 安裝 JobSet v0.2.3 以上版本。

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml

    將 JOBSET_VERSION 替換為 JobSet 版本。例如：v0.2.3。

  2. 確認 JobSet 控制器是否在 jobset-system 命名空間中執行：

    kubectl get pod -n jobset-system

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

    NAME                                        READY   STATUS    RESTARTS   AGE
jobset-controller-manager-69449d86bc-hp5r6   2/2     Running   0          2m15s

  3. 在兩個 TPU 節點集區中部署兩個模型伺服器。儲存下列 sax-model-server-set 資訊清單：

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: sax-model-server-set
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: sax-model-server
      replicas: 2
      template:
        spec:
          parallelism: 8
          completions: 8
          backoffLimit: 0
          template:
            spec:
              serviceAccountName: sax-sa
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 4x8
              containers:
              - name: sax-model-server
                image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0
                args: ["--port=10001","--sax_cell=/sax/test", "--platform_chip=tpuv5e"]
                ports:
                - containerPort: 10001
                - containerPort: 8471
                securityContext:
                  privileged: true
                env:
                - name: SAX_ROOT
                  value: "gs://BUCKET_NAME/sax-root"
                - name: MEGASCALE_NUM_SLICES
                  value: ""
                resources:
                  requests:
                    google.com/tpu: 4
                  limits:
                    google.com/tpu: 4

  4. 將 BUCKET_NAME 替換為您先前建立的 Cloud Storage：

    perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-model-server-set.yaml

    在這個資訊清單中：

    • replicas: 2 是 Job 副本的數量。每個工作都代表一個模型伺服器。因此，這組 Pod 共有 8 個。
    • parallelism: 8 和 completions: 8 等於每個節點集區中的節點數量。
    • 如果任何 Pod 失敗，backoffLimit: 0 就必須為零，才能將 Job 標示為失敗。
    • ports.containerPort: 8471 是 VM 通訊的預設通訊埠
    • name: MEGASCALE_NUM_SLICES 會取消設定環境變數，因為 GKE 並未執行 Multislice 訓練。

  5. 套用資訊清單：

    kubectl apply -f sax-model-server-set.yaml

  6. 確認 Saxml 管理伺服器和模型伺服器 Pod 的狀態：

    kubectl get pods

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

    NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server-557c85f488-lnd5d                 1/1     Running   0          35h
sax-model-server-set-sax-model-server-0-0-nj4sm   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-1-sl8w4   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-2-hb4rk   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-3-qv67g   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-4-pzqz6   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-5-nm7mz   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-6-7br2x   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-7-4pw6z   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-0-8mlf5   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-1-h6z6w   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-2-jggtv   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-3-9v8kj   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-4-6vlb2   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-5-h689p   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-6-bgv5k   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-7-cd6gv   1/1     Running   0          24m

在本範例中，共有 16 個模型伺服器容器：sax-model-server-set-sax-model-server-0-0-nj4sm 和 sax-model-server-set-sax-model-server-1-0-8mlf5 是每個群組中的兩個主要模型伺服器。

您的 Saxml 叢集有兩個模型伺服器，分別部署在兩個 v5e TPU 配量節點集區，拓撲分別為 4x8。

部署 Saxml HTTP 伺服器和負載平衡器

  1. 請使用下列預先建構的映像檔 HTTP 伺服器映像檔。儲存下列 sax-http.yaml 資訊清單：

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-http
spec:
  replicas: 1
  selector:
    matchLabels:
      app: sax-http
  template:
    metadata:
      labels:
        app: sax-http
    spec:
      hostNetwork: false
      serviceAccountName: sax-sa
      containers:
      - name: sax-http
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.0.0
        ports:
        - containerPort: 8888
        env:
        - name: SAX_ROOT
          value: "gs://BUCKET_NAME/sax-root"
---
apiVersion: v1
kind: Service
metadata:
  name: sax-http-lb
spec:
  selector:
    app: sax-http
  ports:
  - protocol: TCP
    port: 8888
    targetPort: 8888
  type: LoadBalancer

  2. 將 BUCKET_NAME 替換為您先前建立的 Cloud Storage：

    perl -pi -e 's|BUCKET_NAME|BUCKET_NAME|g' sax-http.yaml

  3. 套用 sax-http.yaml 資訊清單：

    kubectl apply -f sax-http.yaml

  4. 等待 HTTP 伺服器容器建立完成：

    kubectl get pods

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

    NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server-557c85f488-lnd5d                 1/1     Running   0          35h
sax-http-65d478d987-6q7zd                         1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-0-nj4sm   1/1     Running   0          24m
...

  5. 等待系統為 Service 指派外部 IP 位址：

    kubectl get svc

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

    NAME           TYPE           CLUSTER-IP    EXTERNAL-IP   PORT(S)          AGE
sax-http-lb    LoadBalancer   10.48.11.80   10.182.0.87   8888:32674/TCP   7m36s

使用 Saxml

在 v5e TPU 多主機配量的 Saxml 上載入、部署及提供模型：

載入模型

  1. 擷取 Saxml 的負載平衡器 IP 位址。

    LB_IP=$(kubectl get svc sax-http-lb -o jsonpath='{.status.loadBalancer.ingress[*].ip}')
PORT="8888"

  2. 在兩個 v5e TPU 切片節點集區中載入 LmCloudSpmd175B 測試模型：

    curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/publish --data \
'{
    "model": "/sax/test/spmd",
    "model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
    "checkpoint": "None",
    "replicas": 2
}'

    測試模型沒有微調的檢查點，權重是隨機產生。模型載入作業最多可能需要 10 分鐘。

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

    {
    "model": "/sax/test/spmd",
    "path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
    "checkpoint": "None",
    "replicas": 2
}

  3. 檢查模型完備性：

    kubectl logs sax-model-server-set-sax-model-server-0-0-nj4sm

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

    ...
loading completed.
Successfully loaded model for key: /sax/test/spmd

    模型已完全載入。

  4. 取得模型相關資訊：

    curl --request GET \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/listcell --data \
'{
    "model": "/sax/test/spmd"
}'

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

    {
"model": "/sax/test/spmd",
"model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
"checkpoint": "None",
"max_replicas": 2,
"active_replicas": 2
}

提供模型

提供提示要求：

curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/generate --data \
'{
  "model": "/sax/test/spmd",
  "query": "How many days are in a week?"
}'

輸出內容會顯示模型回覆的範例。由於測試模型具有隨機權重，因此這項回覆可能沒有意義。

取消發布模型

執行下列指令來取消發布模型：

curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/unpublish --data \
'{
    "model": "/sax/test/spmd"
}'

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

{
  "model": "/sax/test/spmd"
}

清除所用資源

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

刪除已部署的資源

  1. 刪除您為本教學課程建立的叢集：

    gcloud container clusters delete saxml --zone ${ZONE}

  2. 刪除服務帳戶：

    gcloud iam service-accounts delete sax-iam-sa@${PROJECT_ID}.s3ns-system.iam.gserviceaccount.com

  3. 刪除 Cloud Storage bucket：

    gcloud storage rm -r gs://${BUCKET_NAME}

後續步驟