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 中）。