跳到主要内容

分布式推理和服务

vLLM 支持分布式tensor-parallel推理和服务托管。 目前,我们支持Megatron-LM的张量并行算法. 我们用ray管理分布式运行时。要运行分布式推理需要安装Ray:

pip install ray

使用LLM类进行多GPU推理, 通过tensor_parallel_size设置你想使用的GPU数量。 例如使用4个GPU推理:

from vllm import LLM
llm = LLM("facebook/opt-13b", tensor_parallel_size=4)
output = llm.generate("San Franciso is a")

启动使用多GPU的服务,启动服务时传--tensor-parallel-size参数。例如, 运行4个GPU的API服务:

python -m vllm.entrypoints.api_server \
--model facebook/opt-13b \
--tensor-parallel-size 4

为了扩展一台以上的服务器, 在vLLM启动前,启动Ray运行时:

# On head node
ray start --head

# On worker nodes
ray start --address=<ray-head-address>

然后, 你可以通过在head节点上启动vLLM进程,并设置tensor_parallel_size为所有机器上GPU的总数,来在多台机器上运行推理和提供服务。