透過 vLLM 在 TPU 上提供 Llama-3.1-8B 服務

本教學課程說明如何使用 vLLM TPU 服務架構，在 v6e TPU VM 上提供 meta-llama/Llama-3.1-8B 模型。

目標

1. 設定環境。
2. 使用 Llama-3.1-8B 執行 vLLM。
3. 傳送推論要求。
4. 執行基準工作負載。
5. 清除所用資源。

費用

本教學課程使用的 Google Cloud 計費元件包括：

• Cloud TPU
• Compute Engine

您可以使用 Pricing Calculator 根據預測用量估算費用。

事前準備

在進行本教學課程前，請先按照「設定 Cloud TPU 環境」頁面的操作說明操作。按照說明操作，建立 Google Cloud專案並設定使用 Cloud TPU。您也可以使用現有Google Cloud 專案。如果選擇這麼做，可以略過建立Google Cloud 專案的步驟，直接從「設定環境以使用 Cloud TPU」開始。

如要使用本教學課程，您需要 Hugging Face 存取權杖。您可以在 Hugging Face 申請免費帳戶。建立帳戶後，請產生存取權杖：

1. 在「Welcome to Hugging Face」(歡迎使用 Hugging Face) 頁面中，按一下帳戶顯示圖片，然後選取「Access tokens」(存取權杖)。
2. 在「存取權權杖」頁面，按一下「建立新權杖」。
3. 選取「讀取」權杖類型，然後輸入權杖名稱。
4. 畫面上會顯示存取權杖。將權杖儲存在安全的地方。

注意： 繼續進行本教學課程前，請注意本課程會使用 Hugging Face 上的受限模型。您必須在 meta-llama/Llama-3.1-8B 模型頁面，向創作者要求並取得模型核准。

設定環境

排入佇列的資源

使用佇列資源 API 建立 Cloud TPU v6e VM。對於 Llama-3.1-8B，我們建議使用 v6e-1 TPU。

export PROJECT_ID=<PROJECT>
export TPU_NAME=<TPU_NAME>
export ZONE=<ZONE>
export QR_ID=<QR_ID>
export TPU_TYPE=<TPU_TYPE>

設定變數：

• PROJECT：專案名稱。
• TPU_NAME：您要建立的 TPU VM 機器名稱。
• ZONE：建立新 VM 的雲端區域。
• TPU_TYPE - 您建立的 TPU VM 類型。例如：v6e-1 或 v6e-4。
• QR_ID - 您建立的排入佇列資源名稱。

建立排入佇列的資源要求：

gcloud alpha compute tpus queued-resources create $QR_ID \
 --node-id $TPU_NAME \
 --project $PROJECT_ID \
 --zone $ZONE \
 --accelerator-type $TPU_TYPE \
 --runtime-version v2-alpha-tpuv6e

確認 TPU VM 已準備就緒。

gcloud compute tpus queued-resources describe $QR_ID \
  --project $PROJECT_ID \
  --zone $ZONE

提示： 佇列中資源要求的初始狀態為 WAITING_FOR_RESOURCES。 TPU VM 建立完成後，狀態會變更為 ACTIVE。視資源是否可用而定，這項作業可能需要一段時間。如果等待時間過長，建議嘗試其他可用區。請等待狀態顯示為 ACTIVE，再繼續操作。

舉例來說，當狀態為 ACTIVE 時：

name: projects/your-project-id/locations/your-zone/queuedResources/your-queued-resource-id
  state:
  state: ACTIVE
  tpu:
  nodeSpec:
  - node:
      acceleratorType: v6e-1
      bootDisk: {}
      networkConfig:
          enableExternalIps: true
      queuedResource: projects/your-project-number/locations/your-zone/queuedResources/your-queued-resource-id
      runtimeVersion: v2-alpha-tpuv6e
      schedulingConfig: {}
      serviceAccount: {}
      shieldedInstanceConfig: {}
      useTpuVm: true
      nodeId: your-node-id
      parent: projects/your-project-number/locations/your-zone

預留項目

使用預訂項目建立 Cloud TPU v6e VM。如果是 Llama-3.1-8B，建議使用 v6e-1 TPU。首先，請設定環境變數：

export PROJECT_ID="YOUR_PROJECT_ID"
export TPU_NAME="TPU_MACHINE_NAME"
export ZONE="ZONE_NAME"
export TPU_TYPE=v6e-8
export HF_TOKEN="YOUR_HF_TOKEN"
export RESERVATION="RESERVATION_NAME"

設定變數：

• PROJECT：專案名稱。
• TPU_NAME：您要建立的 TPU VM 機器名稱。
• ZONE：建立新 VM 的雲端區域。
• TPU_TYPE - 您建立的 TPU VM 類型。例如：v6e-1 或 v6e-4。
• RESERVATION：含有 TPU 的預訂名稱。

使用預留項目建立 TPU VM：

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --zone=$ZONE \
    --project $PROJECT_ID \
    --accelerator-type=$TPU_TYPE \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

連線至 TPU VM。

gcloud compute tpus tpu-vm ssh $TPU_NAME \
  --project $PROJECT_ID \
  --zone $ZONE

使用 Llama-3.1-8B 執行 vLLM

1. 設定 Hugging Face 權杖和模型名稱變數。

     export HF_TOKEN="YOUR_HF_TOKEN"
     export MODEL_NAME="meta-llama/Llama-3.1-8B"
   

2. 在 TPU VM 內，以卸離模式執行 vLLM Docker 容器，並啟動 vLLM 伺服器。這項指令會使用 10 GB 的共用記憶體大小。

   export DOCKER_URI="vllm/vllm-tpu:v0.18.0"
   export CONTAINER_NAME="${USER}-vllm"
   export MAX_MODEL_LEN=4096
   export TP=1 # number of chips
   
   sudo docker run -d --name "${CONTAINER_NAME}" \
       --privileged --net=host \
       -v /dev/shm:/dev/shm \
       --shm-size 10gb \
       -e "HF_HOME=/dev/shm" \
       -e "HF_TOKEN=${HF_TOKEN}" \
       -p 8000:8000 "${DOCKER_URI}" \
           vllm serve ${MODEL_NAME} \
               --seed 42 \
               --gpu-memory-utilization 0.98 \
               --max-num-batched-tokens 1024 \
               --max-num-seqs 128 \
               --tensor-parallel-size $TP \
               --max-model-len $MAX_MODEL_LEN

3. 檢查伺服器記錄，確認伺服器正在運作。

   sudo docker logs -f "${CONTAINER_NAME}"

   vLLM 伺服器執行時，您會看到類似以下的輸出內容。輸出內容顯示完畢後，請按 CTRL+C 返回終端機。

   (APIServer pid=7) INFO:     Started server process [7]
   (APIServer pid=7) INFO:     Waiting for application startup.
   (APIServer pid=7) INFO:     Application startup complete.
   

傳送推論要求

vLLM 伺服器啟動後，您就可以向 API 發送要求。詳情請參閱 vLLM API 參考說明文件。

1. 使用 curl 將測試要求傳送至伺服器。

   sudo docker exec "${CONTAINER_NAME}" \
     curl http://localhost:8000/v1/completions \
       -H "Content-Type: application/json" \
       -d '{
           "prompt": "The future of AI is",
           "max_tokens": 200,
           "temperature": 0
         }'

回應會以 JSON 格式傳回。

執行基準工作負載

您可以從第二個終端機，針對正在執行的伺服器執行基準測試。

1. 在容器內安裝 datasets 程式庫。

   sudo docker exec "${CONTAINER_NAME}" pip install datasets

2. 在容器內執行 vllm bench serve 指令。

   sudo docker exec "${CONTAINER_NAME}" \
       vllm bench serve \
           --backend vllm \
           --dataset-name random \
           --num-prompts 1000 \
           --seed 100

效能評定結果如下所示：

============ Serving Benchmark Result ============
Successful requests:                     1000
Failed requests:                         0
Benchmark duration (s):                  73.97
Total input tokens:                      1024000
Total generated tokens:                  128000
Request throughput (req/s):              13.52
Output token throughput (tok/s):         1730.38
Peak output token throughput (tok/s):    2522.00
Peak concurrent requests:                1000.00
Total Token throughput (tok/s):          15573.42
---------------Time to First Token----------------
Mean TTFT (ms):                          34834.97
Median TTFT (ms):                        34486.19
P99 TTFT (ms):                           70234.40
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          47.30
Median TPOT (ms):                        48.57
P99 TPOT (ms):                           48.60
---------------Inter-token Latency----------------
Mean ITL (ms):                           47.31
Median ITL (ms):                         53.49
P99 ITL (ms):                            54.58
==================================================

清除所用資源

為避免因為本教學課程所用資源，導致系統向 Google Cloud 帳戶收取費用，請刪除含有相關資源的專案，或者保留專案但刪除個別資源。

1. 在終端機中輸入 exit，中斷與 TPU VM 的連線。

刪除資源

您可以刪除專案 (這會一併刪除所有資源)，也可以保留專案並刪除資源。

刪除專案

如要刪除 Google Cloud 專案和所有相關聯的資源，請執行下列指令：

    gcloud projects delete $PROJECT_ID

刪除 TPU 資源

排入佇列的資源

刪除 Cloud TPU 資源。下列指令會使用 --force 參數，同時刪除佇列資源要求和 TPU VM。

gcloud alpha compute tpus queued-resources delete $QR_ID \
  --project=$PROJECT_ID \
  --zone=$ZONE \
  --force

預留項目

刪除 Cloud TPU VM。使用下列指令終止 VM，將 TPU 釋回預訂。

gcloud compute tpus tpu-vm delete $TPU_NAME --zone $ZONE --project $PROJECT_ID --quiet

注意： 只要預訂 TPU，就必須支付相關費用。

後續步驟

• 進一步瞭解 Cloud TPU 上的 vLLM。
• 進一步瞭解 Cloud TPU。