【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
相关推荐
Ronin-Lotus1 小时前
深度学习篇--- ResNet-18
人工智能·深度学习·resnet
说私域2 小时前
基于开源 AI 智能名片链动 2+1 模式 S2B2C 商城小程序的新开非连锁品牌店开业引流策略研究
人工智能·小程序·开源
moonsims3 小时前
无人机和无人系统的计算机视觉-人工智能无人机
人工智能·计算机视觉·无人机
钓了猫的鱼儿3 小时前
无人机航拍数据集|第27期 无人机交通目标检测YOLO数据集3717张yolov11/yolov8/yolov5可训练
人工智能·yolo·目标检测
tzc_fly3 小时前
rbio1:以生物学世界模型为软验证器训练科学推理大语言模型
人工智能·语言模型·自然语言处理
北方有星辰zz3 小时前
语音识别:概念与接口
网络·人工智能·语音识别
尚学教辅学习资料3 小时前
SSM从入门到实战:2.5 SQL映射文件与动态SQL
数据库·sql·动态sql·sql映射
阿里-于怀4 小时前
携程旅游的 AI 网关落地实践
人工智能·网关·ai·旅游·携程·higress·ai网关
大新屋4 小时前
MongoDB 分片集群把非分片集合转成分片集合
数据库·mongodb
赴3354 小时前
神经网络和深度学习介绍
人工智能·深度学习·反向传播