使用vLLM加速大语言模型推理

2024-03-13 12:31:30 浏览数 (1)

介绍

vLLM 是一个快速且易于使用的库,用于 LLM 推理和服务,和 HuggingFace 无缝集成。区别于 chatglm.cpp 和 llama.cpp,仅是在 GPU 上的模型推理加速,没有 CPU 上的加速。

在吞吐量方面,vLLM 的性能比 HuggingFace Transformers (HF) 高出 24 倍,文本生成推理 (TGI) 高出 3.5 倍。

可以使用 ray 框架实现分布式推理:https://vllm.readthedocs.io/en/latest/serving/distributed_serving.html

Github: https://github.com/vllm-project/vllm

文档:https://vllm.readthedocs.io/en/latest/models/supported_models.html

支持 HuggingFace 上的模型

  • Aquila (BAAI/Aquila-7B, BAAI/AquilaChat-7B, etc.)
  • Baichuan (baichuan-inc/Baichuan-7B, baichuan-inc/Baichuan-13B-Chat, etc.)
  • BLOOM (bigscience/bloom, bigscience/bloomz, etc.)
  • Falcon (tiiuae/falcon-7b, tiiuae/falcon-40b, tiiuae/falcon-rw-7b, etc.)
  • GPT-2 (gpt2, gpt2-xl, etc.)
  • GPT BigCode (bigcode/starcoder, bigcode/gpt_bigcode-santacoder, etc.)
  • GPT-J (EleutherAI/gpt-j-6b, nomic-ai/gpt4all-j, etc.)
  • GPT-NeoX (EleutherAI/gpt-neox-20b, databricks/dolly-v2-12b, stabilityai/stablelm-tuned-alpha-7b, etc.)
  • InternLM (internlm/internlm-7b, internlm/internlm-chat-7b, etc.)
  • LLaMA & LLaMA-2 (meta-llama/Llama-2-70b-hf, lmsys/vicuna-13b-v1.3, young-geng/koala, openlm-research/open_llama_13b, etc.)
  • MPT (mosaicml/mpt-7b, mosaicml/mpt-30b, etc.)
  • OPT (facebook/opt-66b, facebook/opt-iml-max-30b, etc.)
  • Qwen (Qwen/Qwen-7B, Qwen/Qwen-7B-Chat, etc.)

安装

代码语言:shell复制
pip install vllm

检查模型是否被 vLLM 支持,返回成功则是支持的。

代码语言:shell复制
from vllm import LLM

llm = LLM(model=...)  # Name or path of your model
output = llm.generate("Hello, my name is")
print(output)

推断

离线批量推断

代码语言:shell复制
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)

# 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 参数可以指定主机和端口。

代码地址:https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/api_server.py

代码语言:shell复制
python -m vllm.entrypoints.api_server

客户端请求,更多示例:https://github.com/vllm-project/vllm/blob/main/examples/api_client.py

代码语言:shell复制
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 参数可以指定主机和端口。

代码地址:https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/api_server.py

代码语言:shell复制
python -m vllm.entrypoints.openai.api_server --model facebook/opt-125m

客户端请求,更多示例:https://github.com/vllm-project/vllm/blob/main/examples/api_client.py

代码语言:shell复制
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 进行请求

代码语言:shell复制
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

代码语言:shell复制
pip install ray

tensor_parallel_size 可以指定使用 GPU 的数量。

代码语言:shell复制
from vllm import LLM
llm = LLM("facebook/opt-13b", tensor_parallel_size=4)
output = llm.generate("San Franciso is a")

Server 指定 GPU 数量

代码语言:shell复制
python -m vllm.entrypoints.api_server 
    --model facebook/opt-13b 
    --tensor-parallel-size 4

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

代码语言:shell复制
# On head node
ray start --head

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

0 人点赞