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

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

背景

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

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

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

事前準備

  • 登入 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.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. 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.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

    •

  • 確認您在專案中具備下列角色： roles/container.admin、roles/iam.serviceAccountAdmin、roles/iam.policyAdmin

    檢查角色

    1. 前往 Google Cloud 控制台的「IAM」頁面。

      前往「IAM」頁面
    2. 選取專案。

    3. 在「主體」欄中，找出所有識別您或您所屬群組的資料列。如要瞭解自己所屬的群組，請與管理員聯絡。

    4. 針對指定或包含您的所有列，請檢查「角色」欄，確認角色清單是否包含必要角色。

    授予角色

    1. 前往 Google Cloud 控制台的「IAM」頁面。

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

    4. 在「New principals」(新增主體) 欄位中，輸入您的使用者 ID。 這通常是指 Google 帳戶的電子郵件地址。

    5. 按一下「選取角色」，然後搜尋角色。
    6. 如要授予其他角色，請按一下「Add another role」(新增其他角色)，然後新增其他角色。
    7. 按一下「Save」(儲存)。

準備環境

  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

建立叢集

使用 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 bucket

建立 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 用戶端程式庫存取服務時，會使用該 IAM 服務帳戶。 Google Cloud 

    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 工作。JobSet 最常見的用途是分散式訓練，但您也可以使用 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 等於每個節點集區中的節點數量。
    • backoffLimit: 0 必須為零，才能在任何 Pod 失敗時將 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 叢集有兩個模型伺服器，分別部署在兩個 4x8 拓撲的 v5e TPU 配量節點集區。

部署 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 --location ${CONTROL_PLANE_LOCATION}

  2. 刪除服務帳戶：

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

  3. 刪除 Cloud Storage bucket：

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

後續步驟