在 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:
請按照「設定 Cloud TPU 環境」指南設定 Google Cloud 專案、設定 Google Cloud CLI、啟用 Cloud TPU API,並確保您有權使用 Cloud TPU。
使用 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 配額
- 使用者專案權限
- 如果您要搭配 XPK 使用 GKE,請參閱「使用者或服務帳戶的 Cloud Console 權限」,瞭解執行 XPK 所需的權限。
配置 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」(服務帳戶) 中找到。
例如: |
請使用 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
安裝依附元件
建立 Miniconda 的目錄:
mkdir -p ~/miniconda3
下載 Miniconda 安裝程式指令碼:
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
安裝 Miniconda:
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
移除 Miniconda 安裝程式指令碼:
rm -rf ~/miniconda3/miniconda.sh
將 Miniconda 新增至
PATH
變數:export PATH="$HOME/miniconda3/bin:$PATH"
重新載入
~/.bashrc
,將變更套用至PATH
變數:source ~/.bashrc
建立 Conda 環境:
conda create -n vllm python=3.11 -y conda activate vllm
複製 vLLM 存放區,然後前往
vllm
目錄:git clone https://github.com/vllm-project/vllm.git && cd vllm
清除現有的 torch 和 torch-xla 套件:
pip uninstall torch torch-xla -y
安裝其他建構依附元件:
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 權杖,請按照下列步驟產生新的權杖:
依序前往「設定檔」>「設定」>「存取權權杖」。
選取「建立新的權杖」。
指定您選擇的名稱和至少具有
Read
權限的角色。選取「產生憑證」。
將產生的權杖複製到剪貼簿,將其設為環境變數,然後使用 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