vLLM是什么

• vllm是一个很快的且方便使用的库，可以用于LLM的推理和服务。
• 和huggingface无缝集成，可以直接下载加载hugging的模型
• 比huggingface Transformer普通推理加速24倍，文本生成推理TGI高处3.5倍
• 和chatglm以及LLaMA不同，vLLM仅做GPU推理加速，不做CPU的推理加速

Github: https://github.com/vllm-project/vllm
文档：https://vllm.readthedocs.io/en/latest/models/supported_models.html

vLLM的使用

0 安装

pip install vllm

1 检查模型是否被LLM支持

from vllm import LLM

llm = LLM(model=...)  # e.g. llm = LLM(model="facebook/opt-125m")
output = llm.generate("Hello, my name is") # 单条样本inference
print(output)

2 离线批量inference

from vllm import LLM, SamplingParams

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(model="facebook/opt-125m")
outputs = llm.generate(prompts, sampling_params) # 多条样本inference

# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

API Server

默认监听 8000 端口，--host 和--port 参数可以指定主机和端口。

python -m vllm.entrypoints.api_server

客户端请求，也可以用Client

curl http://localhost:8000/generate \
    -d '{
        "prompt": "San Francisco is a",
        "use_beam_search": true,
        "n": 4,
        "temperature": 0
    }'

兼容 OpenAI 的 API Server

默认监听 8000 端口，--host 和--port 参数可以指定主机和端口。

python -m vllm.entrypoints.openai.api_server --model facebook/opt-125m

客户端请求

curl http://localhost:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "facebook/opt-125m",
        "prompt": "San Francisco is a",
        "max_tokens": 7,
        "temperature": 0
    }'

可以使用 Openai 的 sdk 进行请求

import openai
# Modify OpenAI's API key and API base to use vLLM's API server.
openai.api_key = "EMPTY"
openai.api_base = "http://localhost:8000/v1"
completion = openai.Completion.create(model="facebook/opt-125m", prompt="San Francisco is a")
print("Completion result:", completion)

分布式推断

理论支持论文：Megatron-LM’s tensor parallel algorithm

安装分布式框架 ray pip install ray
tensor_parallel_size 可以指定使用 GPU 的数量。

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

Server 指定 GPU 数量

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

分别在一个主节点和多个工作节点安装 ray 并运行服务。然后在主节点运行上述的 Server，GPU 数量可以指定为集群内所有的 GPU 数量总和。

# On head node
ray start --head

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