【vLLM 教程】使用 TPU 安装

vLLM 是一款专为大语言模型推理加速而设计的框架,实现了 KV 缓存内存几乎零浪费,解决了内存管理瓶颈问题。

更多 vLLM 中文文档及教程可访问 →vllm.hyper.ai/

vLLM 使用 PyTorch XLA 支持 Google Cloud TPU。

依赖环境

  • Google Cloud TPU VM(单主机和多主机)
  • TPU 版本: v5e、v5p、v4
  • Python: 3.10

安装选项:

  1. href="vllm.hyper.ai/docs/gettin...">使用Dockerfile.tpu构建 Docker 镜像
  2. 从源代码构建

使用Dockerfile.tpu 构建 Docker 镜像

Dockerfile.tpu 用于构建具有 TPU 支持的 docker 镜像。

erlang 复制代码
docker build -f Dockerfile.tpu -t vllm-tpu .

您可以使用以下命令运行 docker 镜像:

css 复制代码
# Make sure to add `--privileged --net host --shm-size=16G`.

# 确保添加 `--privileged --net host --shm-size=16G`。

docker run --privileged --net host --shm-size=16G -it vllm-tpu

从源代码构建

您还可以从源代码构建并安装 TPU 后端。

首先,安装依赖:

bash 复制代码
# (Recommended) Create a new conda environment.
#(推荐)创建一个新的 conda 环境。

conda create -n myenv python=3.10 -y
conda activate myenv

# Clean up the existing torch and torch-xla packages.
# 清理现有的 torch 和 torch-xla 包。

pip uninstall torch torch-xla -y

# Install PyTorch and PyTorch XLA.
# 安装 PyTorch 和 PyTorch XLA。

export DATE="20240828"
export TORCH_VERSION="2.5.0"
pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-${TORCH_VERSION}.dev${DATE}-cp310-cp310-linux_x86_64.whl
pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}.dev${DATE}-cp310-cp310-linux_x86_64.whl

# Install JAX and Pallas.
# 安装 JAX 和 Pallas。

pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html
pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html

# Install other build dependencies.
# 安装其他构建依赖项。

pip install -r requirements-tpu.txt

下一步,从源代码构建 vLLM。这只需要几秒钟:

ini 复制代码
VLLM_TARGET_DEVICE="tpu" python setup.py develop

注意

由于 TPU 依赖于需要静态形状的 XLA,因此 vLLM 会将可能的输入形状进行分桶处理,并为每个不同的形状编译 XLA 图。第一次运行的编译时间可能需要 20~30 分钟。不过由于 XLA 图会缓存在磁盘中(默认在VLLM_XLA_CACHE_PATH~/.cache/vllm/xla_cache 中),之后的编译时间会减少到大约 5 分钟。

提示

如果您遇到以下错误:

python 复制代码
from torch._C import *  # noqa: F403

ImportError: libopenblas.so.0: cannot open shared object file: No such file or directory

请使用以下命令安装 OpenBLAS:

csharp 复制代码
sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev
相关推荐
StarPrayers.几秒前
卷积层(Convolutional Layer)学习笔记
人工智能·笔记·深度学习·学习·机器学习
skywalk81633 分钟前
AutoCoder Nano 是一款轻量级的编码助手, 利用大型语言模型(LLMs)帮助开发者编写, 理解和修改代码。
人工智能
金井PRATHAMA9 分钟前
描述逻辑对人工智能自然语言处理中深层语义分析的影响与启示
人工智能·自然语言处理·知识图谱
却道天凉_好个秋15 分钟前
OpenCV(四):视频采集与保存
人工智能·opencv·音视频
minhuan15 分钟前
构建AI智能体:五十七、LangGraph + Gradio:构建可视化AI工作流的趣味指南
人工智能·语言模型·workflow·langgraph·自定义工作流
Li zlun25 分钟前
MySQL 性能监控与安全管理完全指南
数据库·mysql·安全
WWZZ202535 分钟前
ORB_SLAM2原理及代码解析:SetPose() 函数
人工智能·opencv·算法·计算机视觉·机器人·自动驾驶
lisw0536 分钟前
AIoT(人工智能物联网):融合范式下的技术演进、系统架构与产业变革
大数据·人工智能·物联网·机器学习·软件工程
mtouch3331 小时前
GIS+VR地理信息虚拟现实XR MR AR
大数据·人工智能·ar·无人机·xr·vr·mr
养生技术人1 小时前
Oracle OCP认证考试题目详解082系列第48题
运维·数据库·sql·oracle·database·开闭原则·ocp