透過 Saxml 在 Vertex AI 上使用多主機 Cloud TPU 提供 Llama 3 開放原始碼模型

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

在本指南中，您會下載 Llama 3 70B 模型權重和權杖化工具，並在 Vertex AI 上部署這些項目，在 TPU 上執行 Saxml。

事前準備

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

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

Enable the APIs

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

Enable the APIs

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

2. 按照 Artifact Registry 說明文件安裝 Docker。
3. 請確認您有足夠的配額，可供 Vertex AI 使用 16 個 TPU v5e 晶片。

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

1. Install the Google Cloud CLI.

2. If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.

3. To initialize the gcloud CLI, run the following command:

   gcloud init

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

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

存取模型並下載模型權重

下列步驟適用於具有 M2 記憶體最佳化 VM 的 Vertex AI Workbench 執行個體。如要瞭解如何變更 Vertex AI Workbench 執行個體的機器類型，請參閱「變更 Vertex AI 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 值區

建立 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 資源上傳至 Vertex AI，請執行下列 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

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

• ：您使用 Vertex AI 的區域。 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 v5e 類型，請參閱 TPU v5e。

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。

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

如要從 Vertex AI 端點取得線上推論結果，請執行 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
   

清除所用資源

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