开始使用 vLLM TPU¶
Google Cloud TPU（张量处理单元）可加速机器学习工作负载。vLLM 支持 TPU v6e 和 v5e。有关架构、支持的拓扑结构等信息，请参阅 TPU 系统架构和特定的 TPU 版本页面（v5e 和 v6e）。
要求¶
- Google Cloud TPU VM：可以访问 TPU VM。有关设置说明，请参阅 Cloud TPU 设置指南。
- TPU 版本：v6e, v5e
- Python：3.11 或更新版本（示例中使用 3.12）。
安装¶
有关使用 pip 安装 vllm-tpu 或将其作为 Docker 镜像运行的详细步骤，请参阅 安装指南。
运行 vLLM 服务器¶
安装 vllm-tpu 后，您就可以启动 API 服务器了。
- 登录 Hugging Face：您需要一个 Hugging Face 令牌才能下载模型。
export TOKEN=YOUR_TOKEN
git config --global credential.helper store
huggingface-cli login --token $TOKEN
- 启动服务器：以下命令使用 Llama-3.1-8B 模型启动服务器。
vllm serve "meta-llama/Llama-3.1-8B" \
--download_dir /tmp \
--disable-log-requests \
--tensor_parallel_size=1 \
--max-model-len=2048
- 发送请求
服务器运行后，您可以使用 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 仓库中查看完整的、端到端的示例配方。