快速上手vLLM库进行大语言模型推理与服务

2301_80727036

540人浏览 · 2025-06-20 16:13:22

2301_80727036 · 2025-06-20 16:13:22 发布

在本文中，我们将探讨如何使用vLLM库结合Langchain进行大语言模型（LLM）推理和服务。vLLM是一款快速易用的库，提供最新的服务吞吐量，通过PagedAttention高效管理注意力和价值内存，并支持请求的连续批处理和优化的CUDA内核。本文将通过代码示例详细演示其用法。

技术背景介绍

vLLM是一个专为LLM推理优化的库，旨在提高推理的效率和扩展性。它支持多GPU的并行推理，并能够通过量化来减少计算资源的占用。其核心特性包括：

PagedAttention：高效的注意力和价值内存管理。
连续请求批处理：有效提高吞吐量。
优化的CUDA内核：提升计算性能。

核心原理解析

vLLM的设计着重于高吞吐量和资源高效利用。PagedAttention技术确保内存的高效管理，而连续请求批处理则减少了计算的空闲时间。通过这些优化，vLLM在大规模模型推理中表现出色。

代码实现演示

在开始之前，请确保安装了vllm Python包：

%pip install --upgrade --quiet vllm -q

基本LLM推理

我们将使用一个预训练的模型mosaicml/mpt-7b进行文本生成任务：

from langchain_community.llms import VLLM

llm = VLLM(
    model="mosaicml/mpt-7b",
    trust_remote_code=True,  # 对Hugging Face模型必需
    max_new_tokens=128,
    top_k=10,
    top_p=0.95,
    temperature=0.8,
)

print(llm.invoke("What is the capital of France ?"))

输出: The capital of France is Paris.

集成到LLMChain中

我们可以将vLLM与LLMChain结合使用，以实现更复杂的自然语言处理任务：

from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate

template = """Question: {question}

Answer: Let's think step by step."""
prompt = PromptTemplate.from_template(template)

llm_chain = LLMChain(prompt=prompt, llm=llm)

question = "Who was the US president in the year the first Pokemon game was released?"

print(llm_chain.invoke(question))

输出:

The first Pokemon game was released in 1996.
The president was Bill Clinton.
Clinton was president from 1993 to 2001.
The answer is Clinton.

分布式推理

vLLM支持分布式张量并行推理，只需设置tensor_parallel_size参数。例如，使用4个GPU进行推理：

llm = VLLM(
    model="mosaicml/mpt-30b",
    tensor_parallel_size=4,
    trust_remote_code=True,  # 对Hugging Face模型必需
)

llm.invoke("What is the future of AI?")

量化

vLLM支持AWQ量化，可以通过设置quantization参数来启用：

llm_q = VLLM(
    model="TheBloke/Llama-2-7b-Chat-AWQ",
    trust_remote_code=True,
    max_new_tokens=512,
    vllm_kwargs={"quantization": "awq"},
)

OpenAI兼容服务器

vLLM可以部署为一个兼容OpenAI API协议的服务器，从而替代使用OpenAI API的应用：

from langchain_community.llms import VLLMOpenAI

llm = VLLMOpenAI(
    openai_api_key="EMPTY",
    openai_api_base="http://localhost:8000/v1",
    model_name="tiiuae/falcon-7b",
    model_kwargs={"stop": ["."]},
)
print(llm.invoke("Rome is"))