在 v6e TPU 上執行 vLLM 推論

本教學課程將說明如何在 v6e TPU 上執行 vLLM 推論。並說明如何為 Meta Llama-3.1-8B 模型執行基準測試指令碼。

如要開始在 v6e TPU 上使用 vLLM，請參閱 vLLM 快速入門。

如果您使用的是 GKE，請參閱 GKE 教學課程。

事前準備

您必須簽署同意聲明，才能使用 HuggingFace 存放區中的 Llama3 系列模型。前往 meta-llama/Llama-3.1-8B，填寫同意聲明，然後等待核准。

準備佈建 4 晶片的 TPU v6e：

  1. 請按照「設定 Cloud TPU 環境」指南設定 Google Cloud 專案、設定 Google Cloud CLI、啟用 Cloud TPU API，並確保您有權使用 Cloud TPU。

  2. 使用 Google Cloud 進行驗證，並設定 Google Cloud CLI 的預設專案和區域。

    gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE

安全容量

準備好取得 TPU 容量時，請參閱「Cloud TPU 配額」，進一步瞭解 Cloud TPU 配額。如果您對如何確保容量有其他疑問，請與 Cloud TPU 銷售團隊或客戶服務團隊聯絡。

佈建 Cloud TPU 環境

您可以使用 GKE、GKE 和 XPK，或以佇列資源佈建 TPU VM。

事前準備

  • 請確認專案具備足夠的 TPUS_PER_TPU_FAMILY 配額，該配額會指定您在Google Cloud 專案中可存取的晶片數量上限。
  • 確認專案具備足夠的 TPU 配額，可供下列項目使用：
    • TPU VM 配額
    • IP 位址配額
    • Hyperdisk Balanced 配額
  • 使用者專案權限

配置 TPU v6e

   gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
     --node-id TPU_NAME \
     --project PROJECT_ID \
     --zone ZONE \
     --accelerator-type v6e-4 \
     --runtime-version v2-alpha-tpuv6e \
     --service-account SERVICE_ACCOUNT

指令旗標說明

變數 說明
NODE_ID 在排隊的資源要求分配時，使用者指派的 TPU ID。
PROJECT_ID Google Cloud 專案名稱。使用現有專案，或建立新專案。
可用區 如要進一步瞭解支援的區域，請參閱 TPU 地區和區域文件。
ACCELERATOR_TYPE 如要瞭解支援的加速器類型，請參閱「加速器類型」說明文件。
RUNTIME_VERSION v2-alpha-tpuv6e
SERVICE_ACCOUNT 這是服務帳戶的電子郵件地址，您可以在 Google Cloud 控制台 >「IAM」>「Service Accounts」(服務帳戶) 中找到。

例如：tpu-service-account@<your_project_ID>.iam.gserviceaccount.com

請使用 list 或 describe 指令查詢排隊資源的狀態。

gcloud alpha compute tpus queued-resources describe QUEUED_RESOURCE_ID  \
  --project PROJECT_ID --zone ZONE

如需佇列資源要求狀態的完整清單，請參閱「佇列資源」說明文件。

使用 SSH 連線至 TPU

  gcloud compute tpus tpu-vm ssh TPU_NAME

安裝依附元件

  1. 建立 Miniconda 的目錄：

    mkdir -p ~/miniconda3

  2. 下載 Miniconda 安裝程式指令碼：

    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh

  3. 安裝 Miniconda：

    bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3

  4. 移除 Miniconda 安裝程式指令碼：

    rm -rf ~/miniconda3/miniconda.sh

  5. 將 Miniconda 新增至 PATH 變數：

    export PATH="$HOME/miniconda3/bin:$PATH"

  6. 重新載入 ~/.bashrc，將變更套用至 PATH 變數：

    source ~/.bashrc

  7. 建立 Conda 環境：

    conda create -n vllm python=3.12 -y
conda activate vllm

  8. 複製 vLLM 存放區，然後前往 vllm 目錄：

    git clone https://github.com/vllm-project/vllm.git && cd vllm

  9. 清除現有的 torch 和 torch-xla 套件：

    pip uninstall torch torch-xla -y

  10. 安裝其他建構依附元件：

    pip install -r requirements/tpu.txt
VLLM_TARGET_DEVICE="tpu" python -m pip install --editable .
sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev

取得模型存取權

如果您還沒有 Hugging Face 權杖，請按照下列步驟產生新的權杖：

  1. 依序前往「設定檔」>「設定」>「存取權權杖」。

  2. 選取「建立新的權杖」。

  3. 指定您選擇的名稱和至少具有 Read 權限的角色。

  4. 選取「產生憑證」。

  5. 將產生的權杖複製到剪貼簿，將其設為環境變數，然後使用 huggingface-cli 進行驗證：

    export TOKEN=YOUR_TOKEN
git config --global credential.helper store
huggingface-cli login --token $TOKEN

啟動 vLLM 伺服器

下列指令會將模型權重從 Hugging Face Model Hub 下載到 TPU VM 的 /tmp 目錄，預先編譯一系列輸入形狀，並將模型編譯結果寫入 ~/.cache/vllm/xla_cache。

詳情請參閱 vLLM 說明文件。

cd ~/vllm
vllm serve "meta-llama/Llama-3.1-8B" --download_dir /tmp --swap-space 16 --disable-log-requests --tensor_parallel_size=4 --max-model-len=2048 &> serve.log &

執行 vLLM 基準

執行 vLLM 基準測試指令碼：

export MODEL="meta-llama/Llama-3.1-8B"
pip install pandas
pip install datasets
python benchmarks/benchmark_serving.py \
  --backend vllm \
  --model $MODEL  \
  --dataset-name random \
  --random-input-len 1820 \
  --random-output-len 128 \
  --random-prefix-len 0

清除所用資源

刪除 TPU：

gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \
  --project PROJECT_ID \
  --zone ZONE \
  --force \
  --async