透過 Saxml 在 Gemini Enterprise Agent Platform 上使用 TPU 提供 Gemma 開放原始碼模型

本指南說明如何使用 Saxml，透過 Gemini Enterprise Agent Platform 上的 Tensor 處理單元 (TPU)，提供 Gemma 開放模型大型語言模型 (LLM) 服務。在本指南中，您會將 2B 和 7B 參數指令微調 Gemma 模型下載至 Cloud Storage，並在 Gemini Enterprise Agent Platform 上部署這些模型，該平台會在 TPU 上執行 Saxml。

背景

透過 Saxml 在 Gemini Enterprise Agent Platform 上使用 TPU 提供 Gemma。您可以運用受管理 AI 解決方案，處理低階基礎架構，並以符合成本效益的方式提供大型語言模型服務。本節說明本教學課程中使用的重要技術。

Gemma

Gemma 是一組開放授權的輕量型生成式人工智慧 (AI) 模型，這些 AI 模型可在應用程式、硬體、行動裝置或代管服務中執行。您可以使用 Gemma 模型生成文字，也可以調整這些模型來執行特定工作。

詳情請參閱 Gemma 說明文件。

Saxml

Saxml 是實驗性系統，可為推論作業提供 Paxml、JAX 和 PyTorch 模型。在本教學課程中，我們將說明如何在 Saxml 支援的 TPU 上提供 Gemma 模型，以提高成本效益。GPU 的設定方式類似。Saxml 提供用於建構 Gemini Enterprise Agent Platform 容器的指令碼，我們將在本教學課程中使用這些指令碼。

TPU

TPU 是 Google 開發的客製化特殊應用積體電路 (ASIC)，用於加速 TensorFlow、PyTorch 和 JAX 等資料處理架構。

本教學課程會提供 Gemma 2B 和 Gemma 7B 模型。 Gemini Enterprise Agent Platform 會在下列單一主機 TPU v5e 節點集區中託管這些模型：

  • Gemma 2B：託管於 TPU v5e 節點集區，拓撲為 1x1，代表一個 TPU 晶片。節點的機型為 ct5lp-hightpu-1t。
  • Gemma 7B：託管在 TPU v5e 節點集區中，拓撲為 2x2，代表四個 TPU 晶片。節點的機型為 ct5lp-hightpu-4t。

事前準備

  1. 登入 Google Cloud 帳戶。如果您是 Google Cloud新手，歡迎 建立帳戶，親自評估產品在實際工作環境中的成效。新客戶還能獲得價值 $300 美元的免費抵免額，可用於執行、測試及部署工作負載。

  2. 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

  3. Verify that billing is enabled for your Google Cloud project.

  4. Enable the Agent Platform API.

    Roles required to enable APIs

    To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

    Enable the API

  5. 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

  6. Verify that billing is enabled for your Google Cloud project.

  7. Enable the Agent Platform API.

    Roles required to enable APIs

    To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

    Enable the API

    8.

  8. 在 Google Cloud 控制台中啟用 Cloud Shell。

    啟用 Cloud Shell

    Google Cloud 主控台底部會開啟一個 Cloud Shell 工作階段，並顯示指令列提示。Cloud Shell 是已安裝 Google Cloud CLI 的殼層環境，並已針對您目前的專案設定好相關值。工作階段可能需要幾秒鐘的時間才能完成初始化。

本教學課程假設您使用 Cloud Shell 與 Google Cloud互動。如要使用 Cloud Shell 以外的殼層，請執行下列額外設定：

  1. 安裝 Google Cloud CLI。

  2. 若您採用的是外部識別資訊提供者 (IdP)，請先使用聯合身分登入 gcloud CLI。

  3. 執行下列指令，初始化 gcloud CLI：

    gcloud init
  4. 請確認您有足夠的 TPU v5e 晶片配額，可供 Gemini Enterprise Agent Platform 使用。這項配額預設為 0。如要使用 1x1 拓撲，配額必須為 1；如要使用 2x2，配額必須為 4。如要同時使用這兩種拓撲，配額必須為 5。
  5. 如果您尚未建立 Kaggle 帳戶，請先完成此步驟。

取得模型存取權

請注意，Cloud Shell 可能沒有足夠資源可供下載模型權重。如果發生這種情況，您可以建立 Gemini Enterprise Agent Platform Workbench 執行個體來執行這項工作。

如要存取 Gemma 模型，以便部署至 Gemini Enterprise Agent Platform，請登入 Kaggle 平台、簽署授權同意書，並取得 Kaggle API 權杖。在本教學課程中，您會使用 Kubernetes Secret 儲存 Kaggle 憑證。

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

  1. 前往 Kaggle.com 的模型同意聲明頁面。
  2. 登入 Kaggle (如果尚未登入)。
  3. 按一下「要求存取權」。
  4. 在「選擇同意聲明帳戶」部分，選取「透過 Kaggle 帳戶驗證」，使用 Kaggle 帳戶提供同意聲明。
  5. 接受模型條款及細則。

產生存取權杖

如要透過 Kaggle 存取模型，您需要 Kaggle API 權杖。

如要產生新權杖 (如果沒有)，請按照下列步驟操作：

  1. 在瀏覽器中前往 Kaggle 設定。

  2. 在「API」部分下方，按一下「Create New Token」。

    系統會下載名為「kaggle.json」的檔案。

將存取權杖上傳至 Cloud Shell

在 Cloud Shell 中，您可以將 Kaggle API 權杖上傳至專案： Google Cloud

  1. 在 Cloud Shell 中，依序點選「顯示更多項目」>「上傳」。
  2. 選取「檔案」，然後按一下「選擇檔案」。
  3. 開啟 kaggle.json 檔案。
  4. 按一下「上傳」。

建立 Cloud Storage bucket

建立 Cloud Storage bucket，用來儲存模型檢查點。

在 Cloud Shell 中執行下列指令：

gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME

將 CHECKPOINTS_BUCKET_NAME 替換為儲存模型檢查點的 Cloud Storage bucket 名稱。

將模型複製到 Cloud Storage bucket

在 Cloud Shell 中執行下列指令：

pip install kaggle --break-system-packages

# For Gemma 2B
mkdir -p /data/gemma_2b-it
kaggle models instances versions download google/gemma/pax/2b-it/1 --untar -p /data/gemma_2b-it
gcloud storage cp /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/ --recursive

# For Gemma 7B
mkdir -p /data/gemma_7b-it
kaggle models instances versions download google/gemma/pax/7b-it/1 --untar -p /data/gemma_7b-it
gcloud storage cp /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/ --recursive

部署模型

上傳模型

如要上傳使用 Saxml 容器的 Model 資源，請執行下列gcloud ai models upload指令：

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma2BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Gemma 7B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma7BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

更改下列內容：

  • PROJECT_ID：Google Cloud專案的 ID
  • LOCATION_ID：您使用 Agent Platform 的區域。請注意，TPU 僅適用於 us-west1。
  • DEPLOYED_MODEL_NAME：DeployedModel 的名稱。您也可以使用 Model 的顯示名稱做為 DeployedModel。

建立端點

您必須先將模型部署至端點，才能使用模型進行線上推論。如果您要將模型部署至現有端點，可以跳過這個步驟。以下範例使用 gcloud ai endpoints create 指令：

gcloud ai endpoints create \
  --region=LOCATION \
  --display-name=ENDPOINT_NAME

更改下列內容：

  • LOCATION_ID：您使用 Agent Platform 的區域。
  • ENDPOINT_NAME：端點的顯示名稱。

Google Cloud CLI 工具可能需要幾秒鐘才能建立端點。

將模型部署至端點

端點準備就緒後，請將模型部署至端點。

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai endpoints deploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --model=$MODEL_ID \
  --display-name=DEPLOYED_MODEL_NAME \
  --machine-type=ct5lp-hightpu-4t \
  --traffic-split=0=100

更改下列內容：

  • LOCATION_ID：您使用 Agent Platform 的區域。
  • ENDPOINT_NAME：端點的顯示名稱。
  • DEPLOYED_MODEL_NAME：DeployedModel 的名稱。您也可以使用 Model 的顯示名稱做為 DeployedModel。

Gemma 2B 可部署在較小的 ct5lp-hightpu-1t 機器上，在這種情況下，上傳模型時應指定 --platform_topology=1x1。

Google Cloud CLI 工具可能需要幾分鐘的時間，才能將模型部署至端點。模型成功部署後，這項指令會列印下列輸出內容：

  Deployed a model to the endpoint xxxxx. Id of the deployed model: xxxxx.

從已部署的模型取得線上推論結果

如要透過 Gemini Enterprise Agent Platform 端點叫用模型，請使用標準推論要求 JSON 物件 ，設定推論要求格式。

下列範例使用 gcloud ai endpoints predict 指令：

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

gcloud ai endpoints predict $ENDPOINT_ID \
  --region=LOCATION \
  --http-headers=Content-Type=application/json \
  --json-request instances.json

更改下列內容：

  • LOCATION_ID：您使用 Agent Platform 的區域。
  • ENDPOINT_NAME：端點的顯示名稱。
  • 「instances.json」必須採用下列格式：{"instances": [{"text_batch": "<your prompt>"},{...}]}

正在清除所用資源

為避免產生更多 Gemini Enterprise Agent Platform 費用和 Artifact Registry 費用，請刪除您在本教學課程中建立的 Google Cloud 資源：

  1. 如要從端點取消部署模型並刪除端點，請在殼層中執行下列指令：

    ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

DEPLOYED_MODEL_ID=$(gcloud ai endpoints describe $ENDPOINT_ID \
   --region=LOCATION \
   --format="value(deployedModels.id)")

gcloud ai endpoints undeploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --deployed-model-id=$DEPLOYED_MODEL_ID

gcloud ai endpoints delete $ENDPOINT_ID \
   --region=LOCATION \
   --quiet

    將 LOCATION 替換為您在上一節中建立模型的區域。

  2. 如要刪除模型，請在殼層中執行下列指令：

    MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai models delete $MODEL_ID \
   --region=LOCATION \
   --quiet

    將 LOCATION 替換為您在上一節中建立模型的區域。

限制

  • 在 Gemini Enterprise Agent Platform 上，只有 us-west1 支援 Cloud TPU。詳情請參閱位置。

後續步驟

  • 瞭解如何部署其他 Saxml 模型，例如 Llama2 和 GPT-J。