vLLM：让大型模型推理更快的工具

作者

今天，我要给大家介绍一款名为vLLM的神奇工具。它的主要功能是加速大型语言模型（如OpenAI的GPT-3）的推理速度。如果你是NLP（自然语言处理）领域的研究员或开发者，我相信这个工具可能会对你的工作有所帮助。

为什么我们需要vLLM？

在大规模语言模型运行时，我们常常遇到一个问题：显存不足。这是因为在生成新的单词或者词语（token）时，我们需要存储一些称为 keys 和 values的数据（我们可以把它们看作是模型用来生成新token的”记忆”）在GPU的显存中。然而，这些数据通常非常大，而且大小也会不断变化，这使得显存管理变得困难。传统的管理方式往往会造成显存的60%-80%的浪费。

这就是vLLM要解决的问题。它引入了一种名为PagedAttention的新技术，可以有效地管理这些keys和values，使得显存的使用效率接近最优（浪费比例低于4%）。这样一来，我们就可以使用更大的数据批次进行运算，从而提高GPU的并行计算能力。

vLLM的核心技术：PagedAttention

PagedAttention的工作原理受到了操作系统中虚拟内存和分页的启发。它将每个序列的数据划分为块，每个块包含固定数量的keys和values。这些块不需要连续的内存，因此可以更灵活地对它们进行管理。

此外，PagedAttention还支持内存共享，也就是说，当用一个提示生成多个不同的序列时，可以共享计算量和显存。这种内存共享机制可以大幅降低显存需求（最高可降低55%），从而进一步提升推理速度。

如何使用vLLM？

vLLM的使用非常简单。首先，使用pip命令安装vLLM：

pip install vllm

然后，你就可以使用vLLM来生成文本了：

from vllm import LLM

prompts = ["Hello, my name is", "The capital of France is"]  # 提示
llm = LLM(model="lmsys/vicuna-7b-v1.3")  # 创建一个LLM
outputs = llm.generate(prompts)  # 从提示生成文本

vLLM也支持在线服务。你可以使用以下命令启动一个在线服务：

python -m vllm.entrypoints.openai.api_server --model lmsys/vicuna-7b-v1.3

然后，你就可以通过HTTP请求来调用这个服务了：

curl http://localhost:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "lmsys/vicuna-7b-v1.3",
        "prompt": "San Francisco is a",
        "max_tokens": 7,
        "temperature": 0
    }'

如果你对vLLM感兴趣，可以在这里查阅更多信息。希望这个工具能对你的工作或学习有所帮助！

发表评论