透過 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 設定」。

準備環境

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

  2. 設定預設環境變數：

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

    替換下列值：

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

建立 GKE Standard 叢集

使用 Cloud Shell 執行下列操作：

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

    gcloud container clusters create saxml \
    --location=${CONTROL_PLANE_LOCATION} \
    --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 \
    --location=${CONTROL_PLANE_LOCATION} \
    --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 \
    --location=${CONTROL_PLANE_LOCATION} \
    --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 --location=${CONTROL_PLANE_LOCATION}

  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}.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}.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/sax-sa]"

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

    kubectl annotate serviceaccount sax-sa \
  iam.gke.io/gcp-service-account=sax-iam-sa@${PROJECT_ID}.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 副本的數量。每個工作都代表一個模型伺服器。因此，這組有 8 個 Pod。
    • 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"
}