使用 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