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

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

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

事前準備

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

  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 and Artifact Registry APIs, if any are 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 APIs

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

    啟用 Cloud Shell

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

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

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

  1. 安裝 Google Cloud CLI。

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

    詳情請參閱「 使用聯合身分登入 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. 如要建立 bucket，請執行下列指令：

    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

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

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

建立線上推論端點

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

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

將 ENDPOINT_DISPLAY_NAME 替換為端點的顯示名稱。

將模型部署至端點

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

在本教學課程中，您會使用 4x4 拓撲，部署為 16 個 Cloud TPU v5e 晶片分片的 Llama 3 70B 模型。不過，您可以指定下列任何支援的多主機 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 Agent Platform 位置」。

  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 費用，請刪除在本教學課程中建立的 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

  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