使用 vLLM 运行 Llama 2 Example

vLLM 是一款功能强大的工具,可用于使用来自各种架构的现代文本生成/完成模型进行推理。在此示例中,我们将它与 fal 的无服务器运行时一起使用,通过将其与快速 AI 加速器卡(如 A100)配对,解锁 vLLM 的高吞吐量推理功能。

开始派对!

让我们像往常一样通过导入 fal 开始派对。但在我们继续设置静态 MODEL_NAME 之前,我们首先必须决定使用哪个模型。

py
import fal
MODEL_NAME = "..."

选择模型

vLLM 支持多种模型,包括 Llama、Llama 2、GPT-J、OPT 等(支持模型的完整列表位于)此处

使用 Llama 2

我们将使用 Llama 2 作为我们的主要 LLM 候选模型,但由于它是一个私有模型,因此我们首先必须通过 Hugging Face hub 请求访问它。如果您不想处理此过程而改用开源模型,请随意跳到下一部分。

  1. 前往 Meta 的 Llama 2 登录页面 并填写表格以请求访问权限(使用与您的 HF 帐户相同的电子邮件地址)。
  2. 然后登录您的 HF 帐户,并请求在 meta-llama org 下访问您想要使用的模型,我们将使用 meta-llama/Llama-2-13b-hf,因此您可以使用此 请求链接 访问权限。
  3. 生成 HF 访问令牌 并在 fal 中设置 HUGGING_FACE_HUB_TOKEN 密钥
bash
$ fal secrets set HUGGING_FACE_HUB_TOKEN="my token"
  1. 继续将模型名称更改为您请求访问的模型。
bash
MODEL_NAME="meta-llama/Llama-2-13b-hf"

可选:使用开源模型

如果您已经有了一个公共模型,您可以在这里使用它(只要它在 Hugging Face Hub 上);否则,作为示例,请随意传递 Open Llama 13B 以继续本教程。

bash
MODEL_NAME = "openlm-research/open_llama_13b"

准备 vLLM 引擎

与我们的大多数示例一样,模型将在缓存函数中初始化,以防止每次调用我们的函数时重新加载。我们要做的就是从 HF 集线器传递模型的名称,然后让 vLLM 处理初始化序列的其余部分。

py
@fal.cached
def prepare_engine(model_name: str = MODEL_NAME):
    """Prepare the vLLM text generation engine. You can change the model to any of the
    supported ones as long as you have access to that model on HF Hub (or they are public)."""

    from vllm import LLM

    engine = LLM(model=model_name)
    return engine

定义 I/O

对于此示例,我们只能接受两个参数(提示本身和最大令牌数),但这部分主要取决于您的想象力,您可以根据需要对其进行自定义。

py
import fal
from pydantic import BaseModel

[...]

class ChatOptions(BaseModel):
    prompt: str
    max_tokens: int = 100

class ChatResponse(BaseModel):
    response: str

将他们组合在一起

现在我们有了模型、引擎和 I/O,我们应该能够将它们组合在一起并准备好开始使用!vLLM(和 torch 2.0)将是我们唯一的依赖项,并且一如既往,它已准备好转变为 API 端点,并启用了 serve 注释。

在参数和可定制性方面,vLLM 提供了非常有趣的选项,因此根据您的需求(和您的输入),不要忘记查看它们。至于这个例子,我们只需要传递max_tokens以避免超出某个限制。

py
@fal.function(
    "virtualenv",
    requirements=[
        "vllm",
        "torch==2.0.1",
    ],
    machine_type="GPU",
    keep_alive=60,
    serve=True,
)
def basic_completion_service(options: ChatOptions) -> ChatResponse:
    from vllm import SamplingParams

    engine = prepare_engine()
    sampling_params = SamplingParams(max_tokens=options.max_tokens)

    result = engine.generate(options.prompt, sampling_params)
    completion = result[0].outputs[0]
    return ChatResponse(response=completion.text)
性能提示

如果您的应用能够在发送提示之前对其进行批处理,则可以通过传递提示列表而不是一次传递一个提示来最大化 vLLM 吞吐量。

与他人分享您的 AI 朋友

让我们部署这个 LLM 并通过我们选择的 HTTP 客户端使用它!

bash
$ fal deploy app.py::basic_completion_service --app-name ai-friend
Registered a new revision for function 'stable-diffusion' (revision='[...]').
URL: https://fal.run/$USER/ai-friend

开始和我们想象中的朋友交谈:

bash
$ curl $APP_URL \
  -H 'Authorization: Key $FAL_KEY' \
  -H 'Content-Type: application/json' \
  -H 'Accept: application/json, */*;q=0.5' \
  -d '{"prompt":"once in a time, there was a cat named "}'

# 它应该打印如下内容:
{
    "result": "..."
}

最后更新于 2024 年 6 月 21 日