Google TPU¶

Tensor Processing Units (TPUs)是谷歌自主研发的专用集成电路(ASICs)，用于加速机器学习工作负载。TPU有多个版本，每个版本具有不同的硬件规格。有关TPU的更多信息，请参阅TPU系统架构。如需了解vLLM支持的TPU版本信息，请参阅：

这些TPU版本允许您配置TPU芯片的物理排列方式。这可以提高吞吐量和网络性能。如需了解更多信息，请参阅：

要使用Cloud TPU，您需要获得Google Cloud Platform项目的TPU配额。TPU配额规定了您可以在GPC项目中使用的TPU数量，具体取决于TPU版本、您希望使用的TPU数量以及配额类型。更多信息请参阅TPU配额。

有关TPU的定价信息，请参阅Cloud TPU定价。

您的TPU虚拟机可能需要额外的持久存储。更多信息，请参阅Cloud TPU数据存储选项。

警告

该设备没有预构建的wheel包，因此您必须使用预构建的Docker镜像或从源代码构建vLLM。

要求¶

Google Cloud TPU 虚拟机
TPU版本：v6e、v5e、v5p、v4
Python: 3.11 或更新版本

配置云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	The TPU version you want to use. Specify the TPU version, for example `v5litepod-4` specifies a v5e TPU with 4 cores, `v6e-1` specifies a v6e TPU with 1 core. For more information, see TPU versions.
RUNTIME_VERSION	The TPU VM runtime version to use. For example, use `v2-alpha-tpuv6e` for a VM loaded with one or more v6e TPU(s). For more information see TPU VM images.
SERVICE_ACCOUNT	The email address for your service account. You can find it in the IAM Cloud Console under Service Accounts. For example: `tpu-service-account@<your_project_ID>.iam.gserviceaccount.com`

使用SSH连接到您的TPU虚拟机：

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

使用Python设置¶

预构建的wheel包¶

目前，没有预构建的TPU轮子。

从源码构建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环境：

conda create -n vllm python=3.12 -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
sudo apt-get install --no-install-recommends --yes libopenblas-base libopenmpi-dev libomp-dev

运行安装脚本：

VLLM_TARGET_DEVICE="tpu" python -m pip install -e .

使用Docker进行设置¶

预构建镜像¶

请参阅deployment-docker-pre-built-image了解如何使用官方Docker镜像，注意将镜像名称vllm/vllm-openai替换为vllm/vllm-tpu。

从源码构建镜像¶

你可以使用 docker/Dockerfile.tpu来构建支持TPU的Docker镜像。

docker build -f docker/Dockerfile.tpu -t vllm-tpu .

使用以下命令运行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）。

提示

如果遇到以下错误：

from torch._C import *  # noqa: F403
ImportError: libopenblas.so.0: cannot open shared object file: No such
file or directory

使用以下命令安装OpenBLAS：

sudo apt-get install --no-install-recommends --yes libopenblas-base libopenmpi-dev libomp-dev