本頁面由 Cloud Translation API 翻譯而成。

透過 Saxml 在 Vertex AI 上使用 TPU 提供 Gemma 開放原始碼模型

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

背景

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

Gemma

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

詳情請參閱 Gemma 說明文件。

Saxml

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

TPU

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

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

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

事前準備

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

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

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the Vertex AI API.

Enable the API

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

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the Vertex AI API.

Enable the API

In the Google Cloud console, activate Cloud Shell.

Activate Cloud Shell

At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

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

Install the Google Cloud CLI.
If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.
To initialize the gcloud CLI, run the following command:
```
gcloud init
```
請確認您有足夠的 Vertex AI TPU v5e 晶片配額。根據預設，這項配額為 0。如果是 1x1 拓撲，則必須為 1。如果是 2x2，則必須為 4。如要同時執行這兩種拓撲，則必須為 5。
如果您尚未建立 Kaggle 帳戶，請先完成此步驟。

取得模型存取權

請注意，Cloud Shell 可能沒有足夠資源來下載模型權重。如果是，您可以建立 Vertex AI Workbench 執行個體來執行該工作。

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

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

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

產生存取權杖

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

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

在瀏覽器中前往 Kaggle 設定。
在「API」部分下方，按一下「Create New Token」。

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

將存取權杖上傳至 Cloud Shell

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

在 Cloud Shell 中，依序點選「更多」>「上傳」。
選取「檔案」，然後按一下「選擇檔案」。
開啟 kaggle.json 檔案。
按一下「上傳」。

建立 Cloud Storage 值區

建立 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：您使用 Vertex AI 的區域。請注意，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：您使用 Vertex AI 的區域。
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：您使用 Vertex AI 的區域。
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.

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

如要透過 Vertex AI 端點叫用模型，請使用標準推論要求 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：您使用 Vertex AI 的區域。
ENDPOINT_NAME：端點的顯示名稱。
instances.json 的格式如下：{"instances": [{"text_batch": "<your prompt>"},{...}]}

正在清除所用資源

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

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

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 替換為您在上一節中建立模型的區域。

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

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 替換為您在上一節中建立模型的區域。

限制

在 Vertex AI 中，只有 us-west1 支援 Cloud TPU。詳情請參閱位置。

後續步驟

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

透過 Saxml 在 Vertex AI 上使用 TPU 提供 Gemma 開放原始碼模型 透過集合功能整理內容 你可以依據偏好儲存及分類內容。

背景

Gemma

Saxml

TPU

事前準備

取得模型存取權

簽署授權同意聲明協議

產生存取權杖

將存取權杖上傳至 Cloud Shell

建立 Cloud Storage 值區

將模型複製到 Cloud Storage bucket

部署模型

上傳模型

Gemma 2B-it

Gemma 7B-it

建立端點

將模型部署至端點

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

正在清除所用資源

限制

後續步驟

透過 Saxml 在 Vertex AI 上使用 TPU 提供 Gemma 開放原始碼模型