分布式推理和服务
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的总数,来在多台机器上运行推理和提供服务。