在 Google Cloud Vertex AI 上部署 Meta Llama 3.1 405B

Meta Llama 3.1 405B 可以使用 Hugging Face 的 Deep Learning Containers (DLCs) 與 Text Generation Inference (TGI) 在 Google Cloud Vertex AI 上部署。這讓使用者能在受管理的環境中執行 Meta 最大的開源 LLM，並利用 A3 加速器最佳化機器進行高效能推論。

硬體需求與記憶體管理

部署 Meta Llama 3.1 405B 需要大量 GPU VRAM，且會根據模型權重的精度而有所不同。載入模型檢查點的大致記憶體需求如下：

模型大小 FP16 FP8 INT4
8B 16 GB 8 GB 4 GB
70B 140 GB 70 GB 35 GB
405B 810 GB 405 GB 203 GB

由於單一 H100 節點（8 x 80GB GPU）約提供 640GB VRAM，405B 模型必須以較低精度（例如 FP8）或多節點設定執行。此部署建議使用 A3 High GPU 機型，提供 8 x NVIDIA H100 80GB GPU、208 個 vCPU 與 1872 GB 記憶體。

在 Vertex AI 上的部署流程

部署模型的流程包括在 Vertex AI Model Registry 中註冊模型，然後將其部署至 Vertex AI Endpoint。

模型註冊

註冊使用 google-cloud-aiplatform Python SDK 完成。此流程需要支援 Meta Llama 3.1 架構的 Hugging Face TGI DLC（其使用與 Llama 3 不同的 RoPE 縮放方法）。主要的環境變數設定包括：

  • MODEL_ID：設定為 meta-llama/Meta-Llama-3.1-405B-Instruct-FP8。
  • HUGGING_FACE_HUB_TOKEN：用於受限 Meta Llama 儲存庫的讀取權杖。
  • NUM_SHARD：設定為 8，以使用 A3 實例上的所有 GPU。
  • HF_XET_HIGH_PERFORMANCE：建議設定為 1，以加速約 400 GiB 模型權重的下載。

端點部署

註冊完成後，模型會使用 a3-highgpu-8g 機型與 8 個 NVIDIA H100 80GB 加速卡部署至 Vertex AI Endpoint。完整的部署流程通常需要 25-30 分鐘，包含資源配置、權重下載（約 10 分鐘）以及載入至 TGI（約 2 分鐘）。

執行線上預測

Vertex AI 會公開 /predict 路由，透過 TGI DLC 提供文字生成服務。由於使用 /generate 端點，輸入必須先以適當的聊天模板格式化，才能送至端點。

輸入格式化

使用者應使用 transformers 套件中的 apply_chat_template 方法，以確保對話正確地為 Meta Llama 3.1 格式化。典型的格式化提示會包含特定的標記，例如 <|begin_of_text|>、<|start_header_id|> 與 <|eot_id|>。

預測方法

預測可透過以下三種主要方式執行：

  1. 同一會話 Python：在部署後立即使用 deployed_model.predict 方法。
  2. 不同會話 Python：透過資源名稱 (projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}) 來實例化 aiplatform.Endpoint。
  3. Vertex AI UI：使用線上預測 UI，提供包含格式化 inputs 與生成 parameters（例如 max_new_tokens、temperature、top_p）的 JSON 負載。

資源管理

為避免不必要的費用，使用完畢後應清理資源。這包括使用 deployed_model.undeploy_all() 從所有端點撤除模型、透過 deployed_model.delete() 刪除端點，最後使用 model.delete() 從註冊表中移除模型。

Sources