透過 Saxml 在 Gemini Enterprise Agent Platform 上使用多主機 Cloud TPU 提供 Llama 3 開放原始碼模型

Llama 3 是 Meta 的開放原始碼大型語言模型 (LLM)。本指南說明如何使用 Gemini Enterprise Agent Platform 上的多主機 Tensor Processing Unit (TPU)，透過 Saxml 提供 Llama 3 LLM 服務。

在本指南中，您將下載 Llama 3 70B 模型權重和權杖化工具，並在 Gemini Enterprise Agent Platform 上部署這些項目，該平台會在 TPU 上執行 Saxml。

事前準備

建議您使用記憶體最佳化 M2 VM 下載模型，並轉換為 Saxml。這是因為模型轉換程序需要大量記憶體，如果選擇的機器類型記憶體不足，可能會導致轉換失敗。

  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 and Artifact Registry APIs.

    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 APIs

  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 and Artifact Registry APIs.

    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 APIs

    8.

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

    啟用 Cloud Shell

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

  9. 按照 Artifact Registry 說明文件安裝 Docker。
  10. 請確認您有足夠的配額，可供 Gemini Enterprise Agent Platform 使用 16 個 TPU v5e 晶片。

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

  1. 安裝 Google Cloud CLI。

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

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

    gcloud init

如果使用 Cloud Shell 以外的殼層部署模型，請確認 Google Cloud CLI 版本高於 475.0.0。如要更新 Google Cloud CLI，請執行 gcloud components update 指令。

如果您使用 Vertex AI SDK 部署模型，請確認您使用的是 1.50.0 以上版本。

存取模型並下載模型權重

下列步驟適用於具有 M2 記憶體最佳化 VM 的 Agent Platform Workbench 執行個體。如要瞭解如何變更 Agent Platform Workbench 執行個體的機器類型，請參閱「變更 Agent Platform Workbench 執行個體的機器類型」。

  1. 前往 Llama 模型同意聲明頁面。

  2. 選取 Llama 3、填寫同意聲明表單，並接受條款及細則。

  3. 前往收件匣查看是否收到內含已簽署網址的電子郵件。

  4. 執行下列指令，從 GitHub 下載 download.sh 指令碼：

    wget https://raw.githubusercontent.com/meta-llama/llama3/main/download.sh
chmod +x download.sh

  5. 如要下載模型權重，請執行從 GitHub 下載的 download.sh 指令碼。

  6. 系統提示時，請輸入您在上一節收到的電子郵件中的已簽署網址。

  7. 系統提示要下載的模型時，請輸入 70B。

將模型權重轉換為 Saxml 格式

  1. 執行下列指令下載 Saxml：

    git clone https://github.com/google/saxml.git

  2. 執行下列指令，設定 Python 虛擬環境：

    python -m venv .
source bin/activate

  3. 執行下列指令來安裝依附元件：

    pip install --upgrade pip

pip install paxml

pip install praxis

pip install torch

  4. 如要將模型權重轉換為 Saxml 格式，請執行下列指令：

    python3 saxml/saxml/tools/convert_llama_ckpt.py \
    --base PATH_TO_META_LLAMA3 \
    --pax PATH_TO_PAX_LLAMA3 \
    --model-size llama3_70b

    更改下列內容：

    • PATH_TO_META_LLAMA3：包含已下載模型權重的目錄路徑
    • PATH_TO_PAX_LLAMA3：儲存轉換後模型權重的目錄路徑

    轉換後的模型會放入 $PATH_TO_PAX_LLAMA3/checkpoint_00000000 資料夾。

  5. 將權杖化工具檔案從原始目錄複製到名為 vocabs 的子資料夾，如下所示：

    cp $PATH_TO_META_LLAMA3/tokenizer.model $PATH_TO_PAX_LLAMA3/vocabs/tokenizer.model

  6. 在 $PATH_TO_PAX_LLAMA3 資料夾中新增空白的 commit_success.txt 檔案，並在該資料夾中新增 metadata 和 state 子資料夾，如下所示：

    touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/commit_success.txt
touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/commit_success.txt
touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/commit_success.txt

    $PATH_TO_PAX_LLAMA3 資料夾現在包含下列資料夾和檔案：

    $PATH_TO_PAX_LLAMA3/checkpoint_00000000/commit_success.txt
$PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/
$PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/commit_success.txt
$PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/
$PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/commit_success.txt
$PATH_TO_PAX_LLAMA3/vocabs/tokenizer.model

建立 Cloud Storage bucket

建立 Cloud Storage bucket，用來儲存轉換後的模型權重。

  1. 在 Cloud Shell 中執行下列指令，並將 PROJECT_ID 替換為您的專案 ID：

    projectid=PROJECT_ID
gcloud config set project ${projectid}

  2. 如要建立值區，請執行下列指令：

    gcloud storage buckets create gs://WEIGHTS_BUCKET_NAME

    將 WEIGHTS_BUCKET_NAME 替換為您要使用的 bucket 名稱。

將模型權重複製到 Cloud Storage bucket

如要將模型權重複製到 bucket，請執行下列指令：

gcloud storage cp PATH_TO_PAX_LLAMA3/* gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b/ --recursive

上傳模型

預先建構的 Saxml 容器位於 us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest。

如要使用預先建構的 Saxml 容器，將 Model 資源上傳至 Gemini Enterprise Agent Platform，請執行 gcloud ai models upload 指令，如下所示：

gcloud ai models upload \
    --region=LOCATION \
    --display-name=MODEL_DISPLAY_NAME \
    --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
    --artifact-uri='gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b' \
    --container-args='--model_path=saxml.server.pax.lm.params.lm_cloud.LLaMA3_70BFP16x16' \
    --container-args='--platform_chip=tpuv5e' \
    --container-args='--platform_topology=4x4' \
    --container-args='--ckpt_path_suffix=checkpoint_00000000' \
    --container-deployment-timeout-seconds=2700 \
    --container-ports=8502 \
    --project=PROJECT_ID

請將下列項目改為對應的值：

  • ：您使用 Gemini Enterprise Agent Platform 的區域。 LOCATION請注意，TPU 僅適用於 us-west1。
  • MODEL_DISPLAY_NAME：您要為模型設定的顯示名稱
  • PROJECT_ID：您 Google Cloud 專案的 ID

建立線上推論端點

如要建立端點，請執行下列指令：

gcloud ai endpoints create \
    --region=LOCATION \
    --display-name=ENDPOINT_DISPLAY_NAME \
    --project=PROJECT_ID

將 ENDPOINT_DISPLAY_NAME 替換為您要為端點設定的顯示名稱。

將模型部署至端點

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

在本教學課程中，您將部署 Llama 3 70B 模型，並使用 4x4 拓撲將模型分片，供 16 個 Cloud TPU v5e 晶片使用。不過，您可以指定下列任何支援的多主機 Cloud TPU 拓撲：

機器類型 拓撲 TPU 晶片數量 主機數量
ct5lp-hightpu-4t 4x4 16 2
ct5lp-hightpu-4t 4x8 32 4
ct5lp-hightpu-4t 8x8 64 8
ct5lp-hightpu-4t 8x16 128 16
ct5lp-hightpu-4t 16x16 256 32

如果您要部署 Saxml GitHub 存放區中定義的其他 Llama 模型，請確認模型已分割，可配合目標裝置數量，且 Cloud TPU 有足夠的記憶體可載入模型。

如要瞭解如何在單一主機 Cloud TPU 上部署模型，請參閱「部署模型」。

如需支援的 Cloud TPU 類型和區域完整清單，請參閱「Gemini Enterprise 代理程式平台位置」。

  1. 取得線上推論端點的端點 ID：

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

  2. 取得模型的 ID：

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

  3. 將模型部署至端點：

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

    將 DEPLOYED_MODEL_NAME 替換為已部署的名稱。這可以與模型顯示名稱 (MODEL_DISPLAY_NAME) 相同。

    部署作業可能會逾時。

    deploy-model 指令會傳回作業 ID，可用於檢查作業何時完成。您可以輪詢作業狀態，直到回應包含 "done": true 為止。使用下列指令輪詢狀態：

    gcloud ai operations describe \
--region=LOCATION \
OPERATION_ID

    將 OPERATION_ID 替換為上一個指令傳回的作業 ID。

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

如要從 Gemini Enterprise Agent Platform 端點取得線上推論結果，請執行 gcloud ai endpoints predict 指令。

  1. 執行下列指令，建立包含範例推論要求的 request.json 檔案：

    cat << EOF > request.json
{"instances": [{"text_batch": "the distance between Earth and Moon is "}]}
EOF

  2. 如要將線上推論要求傳送至端點，請執行下列指令：

    gcloud ai endpoints predict $ENDPOINT_ID \
    --project=PROJECT_ID \
    --region=LOCATION \
    --json-request=request.json

清除所用資源

如要避免產生更多 Gemini Enterprise Agent Platform 費用，請刪除您在本教學課程中建立的 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

  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