Google TPU¶
張量處理單元 (TPU) 是 Google 定製開發的專用積體電路 (ASIC)，用於加速機器學習工作負載。TPU 有不同版本，每個版本具有不同的硬體規格。有關 TPU 的更多資訊，請參閱 TPU 系統架構。有關 vLLM 支援的 TPU 版本的更多資訊，請參閱
這些 TPU 版本允許您配置 TPU 晶片的物理排列。這可以提高吞吐量和網路效能。有關更多資訊，請參閱
為了使用 Cloud TPU，您的 Google Cloud Platform 專案需要獲得 TPU 配額。TPU 配額指定您可以在 GCP 專案中使用多少 TPU，並按 TPU 版本、要使用的 TPU 數量和配額型別進行指定。有關更多資訊，請參閱 TPU 配額。
有關 TPU 定價資訊，請參閱 Cloud TPU 定價。
您的 TPU VM 可能需要額外的永續性儲存。有關更多資訊，請參閱 Cloud TPU 資料的儲存選項。
警告
此裝置沒有預構建的 wheels，因此您必須使用預構建的 Docker 映象或從原始碼構建 vLLM。
要求¶
- Google Cloud TPU VM
- TPU 版本：v6e, v5e, v5p, v4
- Python：3.11 或更高版本
預配 Cloud TPU¶
您可以使用 Cloud TPU API 或 排隊資源 API（首選）預配 Cloud TPU。本節展示瞭如何使用排隊資源 API 建立 TPU。有關使用 Cloud TPU API 的更多資訊，請參閱 使用 Create Node API 建立 Cloud TPU。排隊資源使您能夠以排隊方式請求 Cloud TPU 資源。當您請求排隊資源時，請求會新增到 Cloud TPU 服務維護的佇列中。當請求的資源可用時，它會被分配到您的 Google Cloud 專案，供您立即獨佔使用。
注意
在以下所有命令中，請將所有大寫的引數名稱替換為適當的值。有關更多資訊，請參閱引數描述表。
透過 GKE 預配 Cloud TPU¶
有關在 GKE 中使用 TPU 的更多資訊，請參閱
配置新環境¶
使用排隊資源 API 預配 Cloud TPU¶
建立包含 4 個 TPU 晶片的 TPU v5e
gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
--node-id TPU_NAME \
--project PROJECT_ID \
--zone ZONE \
--accelerator-type ACCELERATOR_TYPE \
--runtime-version RUNTIME_VERSION \
--service-account SERVICE_ACCOUNT
| 引數名稱 | 描述 |
|---|---|
| QUEUED_RESOURCE_ID | 排隊資源請求的使用者分配 ID。 |
| TPU_NAME | 排隊資源請求分配時建立的 TPU 的使用者分配名稱。 |
| PROJECT_ID | 您的 Google Cloud 專案 |
| ZONE | 您要建立 Cloud TPU 的 GCP 可用區。您使用的值取決於您正在使用的 TPU 版本。有關更多資訊，請參閱 TPU 區域和可用區 |
| ACCELERATOR_TYPE | 您要使用的 TPU 版本。請指定 TPU 版本，例如 v5litepod-4 指定了具有 4 個核心的 v5e TPU，v6e-1 指定了具有 1 個核心的 v6e TPU。有關更多資訊，請參閱 TPU 版本。 |
| RUNTIME_VERSION | 要使用的 TPU VM 執行時版本。例如，對於載入了一個或多個 v6e TPU 的 VM，請使用 v2-alpha-tpuv6e。有關更多資訊，請參閱 TPU VM 映象。 |
| SERVICE_ACCOUNT | 您的服務帳號的電子郵件地址。您可以在 IAM Cloud Console 的服務帳號下找到它。例如：tpu-service-account@<your_project_ID>.iam.gserviceaccount.com |
使用 SSH 連線到您的 TPU VM
使用 Python 進行設定¶
預構建的 Wheels¶
目前，沒有預構建的 TPU wheels。
從原始碼構建 Wheel¶
安裝 Miniconda
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh
source ~/.bashrc
為 vLLM 建立並激活 Conda 環境
克隆 vLLM 倉庫並進入 vLLM 目錄
解除安裝現有的 torch 和 torch_xla 軟體包
安裝構建依賴項
pip install -r requirements/tpu.txt
sudo apt-get install --no-install-recommends --yes libopenblas-base libopenmpi-dev libomp-dev
執行設定指令碼
使用 Docker 進行設定¶
預構建的映象¶
有關使用官方 Docker 映象的說明，請參閱 deployment-docker-pre-built-image，請務必將映象名稱 vllm/vllm-openai 替換為 vllm/vllm-tpu。
從原始碼構建映象¶
您可以使用 docker/Dockerfile.tpu 來構建支援 TPU 的 Docker 映象。
使用以下命令執行 Docker 映象
# Make sure to add `--privileged --net host --shm-size=16G`.
docker run --privileged --net host --shm-size=16G -it vllm-tpu
注意
由於 TPU 依賴於需要靜態形狀的 XLA，vLLM 將可能的輸入形狀進行分桶，併為每個形狀編譯一個 XLA 圖。首次執行時的編譯時間可能需要 20~30 分鐘。但是，此後編譯時間會縮短到約 5 分鐘，因為 XLA 圖會快取到磁碟中（預設在 VLLM_XLA_CACHE_PATH 或 ~/.cache/vllm/xla_cache 中）。