【vLLM 教程】使用 TPU 安装

vLLM 是一款专为大语言模型推理加速而计划的框架，实现了 KV 缓存内存几乎零浪费，解决了内存管理瓶颈题目。
更多 vLLM 中文文档及教程可访问 →https://vllm.hyper.ai/
vLLM 使用 PyTorch XLA 支持 Google Cloud TPU。
依靠环境​

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

安装选项：

• href=“https://vllm.hyper.ai/docs/getting-started/installation-with-tpu#%E4%BD%BF%E7%94%A8-dockerfiletpu-%E6%9E%84%E5%BB%BA-docker-%E9%95%9C%E5%83%8F”>使用Dockerfile.tpu构建 Docker 镜像
  
• 从源代码构建
  

使用Dockerfile.tpu 构建 Docker 镜像​

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

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

复制代码

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

1. # Make sure to add `--privileged --net host --shm-size=16G`.
3. # 确保添加 `--privileged --net host --shm-size=16G`。
5. docker run --privileged --net host --shm-size=16G -it vllm-tpu

复制代码

从源代码构建​

您还可以从源代码构建并安装 TPU 后端。
起首，安装依靠：

1. # (Recommended) Create a new conda environment.
2. #（推荐）创建一个新的 conda 环境。
4. conda create -n myenv python=3.10 -y
5. conda activate myenv
7. # Clean up the existing torch and torch-xla packages.
8. # 清理现有的 torch 和 torch-xla 包。
10. pip uninstall torch torch-xla -y
12. # Install PyTorch and PyTorch XLA.
13. # 安装 PyTorch 和 PyTorch XLA。
15. export DATE="20240828"
16. export TORCH_VERSION="2.5.0"
17. pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-${TORCH_VERSION}.dev${DATE}-cp310-cp310-linux_x86_64.whl
18. pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}.dev${DATE}-cp310-cp310-linux_x86_64.whl
20. # Install JAX and Pallas.
21. # 安装 JAX 和 Pallas。
23. pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html
24. 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
26. # Install other build dependencies.
27. # 安装其他构建依赖项。
29. pip install -r requirements-tpu.txt

复制代码

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

1. VLLM_TARGET_DEVICE="tpu" python setup.py develop

复制代码

注意
由于 TPU 依靠于需要静态形状的 XLA，因此 vLLM 会将可能的输入形状举行分桶处理，并为每个差别的形状编译 XLA 图。第一次运行的编译时间可能需要 20~30 分钟。不过由于 XLA 图会缓存在磁盘中（默认在VLLM_XLA_CACHE_PATH 或 ~/.cache/vllm/xla_cache 中），之后的编译时间会淘汰到大约 5 分钟。
提示
如果您遇到以下错误：

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

复制代码

请使用以下命令安装 OpenBLAS：

1. sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。