使用 TPU 安装

使用 TPU 安装#

vLLM 使用 PyTorch XLA 支持 Google Cloud TPU。

要求#

  • Google Cloud TPU VM（单主机和多主机）

  • TPU 版本：v5e、v5p、v4

  • Python：3.10

安装选项：

  1. 使用 Dockerfile 构建 Docker 镜像。

  2. 从源代码构建。

使用 Dockerfile.tpu 构建 Docker 镜像#

Dockerfile.tpu 提供了构建支持 TPU 的 Docker 镜像的方法。

$ docker build -f Dockerfile.tpu -t vllm-tpu .

你可以使用以下命令运行 Docker 镜像：

$ # Make sure to add `--privileged --net host --shm-size=16G`.
$ docker run --privileged --net host --shm-size=16G -it vllm-tpu

从源代码构建#

你也可以从源代码构建并安装 TPU 后端。

首先，安装依赖项：

$ # (Recommended) Create a new conda environment.
$ conda create -n myenv python=3.10 -y
$ conda activate myenv

$ # Clean up the existing torch and torch-xla packages.
$ pip uninstall torch torch-xla -y

$ # Install PyTorch and 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.
$ 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。这只需要几秒钟：

$ VLLM_TARGET_DEVICE="tpu" python setup.py develop

备注

由于 TPU 依赖于需要静态形状的 XLA，vLLM 会将可能的输入形状分组，并为每个不同的形状编译一个 XLA 图。第一次运行时，编译时间可能需要 20~30 分钟。但是，之后编译时间会缩短到约 5 分钟，因为 XLA 图会被缓存到磁盘中（默认情况下在 VLLM_XLA_CACHE_PATH 或 ~/.cache/vllm/xla_cache 中）。

小技巧

如果你遇到以下错误：

from torch._C import *  # noqa: F403
ImportError: libopenblas.so.0: cannot open shared object file: No such file or directory

请使用以下命令安装 OpenBLAS：

$ sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev
目录