今天,我要给大家介绍一款名为 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 感兴趣,可以在 这里 查阅更多信息。希望这个工具能对你的工作或学习有所帮助!