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

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

背景

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

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. In the Cloud de Confiance console, on the project selector page, select or create a Cloud de Confiance 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

  2. Verify that billing is enabled for your Cloud de Confiance project.

  3. Enable the Agent Platform API, if it is not already enabled.

    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

  4. 在 Cloud de Confiance 控制台中啟用 Cloud Shell。

    啟用 Cloud Shell

    Cloud de Confiance 控制台底部會開啟 Cloud Shell 工作階段，並顯示指令列提示。Cloud Shell 是已安裝 Google Cloud CLI 的殼層環境，並已設定適用於您目前專案的值。工作階段可能要幾秒鐘的時間才能初始化。

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

  1. 安裝 Google Cloud CLI。

  2. 設定 gcloud CLI，使用您的聯合身分。

    詳情請參閱「 使用聯合身分登入 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 權杖上傳至專案： Cloud de Confiance by S3NS

  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：Cloud de Confiance專案的 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 費用，請刪除在本教學課程中建立的 Cloud de Confiance 資源：

  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，Cloud TPU 僅支援 us-west1。詳情請參閱位置。

後續步驟

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