跳到内容

vLLM TPU 入门

Google Cloud TPU（张量处理单元）可加速机器学习工作负载。vLLM 支持 TPU v6e 和 v5e。有关架构、支持的拓扑结构等更多信息，请参阅 TPU 系统架构 以及特定的 TPU 版本页面（v5e 和 v6e）。

要求

  • Google Cloud TPU VM： 访问 TPU VM。有关设置说明，请参阅 Cloud TPU 设置指南。
  • TPU 版本： v7x, v6e, v5e
  • Python： 3.11 或更高版本（示例中使用 3.12）。

安装

有关使用 pip 安装 vllm-tpu 或将其作为 Docker 镜像运行的详细步骤，请参阅 安装指南。

我们建议使用 uv (uv pip install) 而不是标准的 pip，因为它可以提高安装速度。

uv pip install vllm-tpu

验证安装

快速验证安装是否成功以及 vllm-tpu 是否配置正确

python -c '
import jax
import vllm
import importlib.metadata
from vllm.platforms import current_platform

tpu_version = importlib.metadata.version("tpu_inference")
print(f"vllm version: {vllm.__version__}")
print(f"tpu_inference version: {tpu_version}")
print(f"vllm platform: {current_platform.get_device_name()}")
print(f"jax backends: {jax.devices()}")
'
# Expected output:
# vllm version: 0.x.x
# tpu_inference version: 0.x.x
# vllm platform: TPU V6E (or your specific TPU architecture)
# jax backends: [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), ...]

运行 vLLM 服务器

安装 vllm-tpu 后，您可以启动 API 服务器。

  1. 登录 Hugging Face： 您需要一个 Hugging Face 令牌来下载模型。
export TOKEN=YOUR_TOKEN
git config --global credential.helper store
huggingface-cli login --token $TOKEN
  1. 启动服务器： 以下命令使用 Llama-3.1-8B 模型启动服务器。
vllm serve "meta-llama/Llama-3.1-8B" \
    --download_dir /tmp \
    --tensor_parallel_size=1 \
    --max-model-len=2048
  1. 发送请求

服务器运行后，您可以使用 curl 向其发送请求

curl https://:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "meta-llama/Llama-3.1-8B",
        "prompt": "Hello, my name is",
        "max_tokens": 20,
        "temperature": 0.7
    }'

后续步骤：

在 tpu-recipes 仓库 中查看完整的端到端示例方案

延伸阅读