跳至内容

模型池化

vLLM还支持池化模型，例如嵌入、分类和奖励模型。

在vLLM中，池化模型实现了VllmModelForPooling接口。这些模型使用Pooler在返回之前提取输入的最终隐藏状态。

注意

我们目前支持模型池化主要是出于便利性考虑。这并不能保证比直接使用HF Transformers/Sentence Transformers带来任何性能提升。

我们目前正计划优化vLLM中的池化模型。如果您有任何建议，请在 Issue #21796上发表评论！

配置

模型运行器

通过选项--runner pooling以池化模式运行模型。

提示

在绝大多数情况下无需设置此选项，因为vLLM可以通过--runner auto自动检测要使用的模型运行器。

模型转换

vLLM可以通过选项--convert 来适配各种池化任务的模型。

如果--runner pooling已被设置（手动或自动）但模型未实现VllmModelForPooling接口，vLLM将尝试根据下表中显示的架构名称自动转换模型。

Architecture --convert Supported pooling tasks
*ForTextEncoding, *EmbeddingModel, *Model embed encode, embed
*For*Classification, *ClassificationModel classify encode, classify, score
*ForRewardModeling, *RewardModel reward encode

提示

你可以显式设置--convert 来指定如何转换模型。

任务池化

vLLM中的每个池化模型根据Pooler.get_supported_tasks支持一个或多个任务，从而启用相应的API：

任务 APIs
encode LLM.reward(...)
embed LLM.embed(...), LLM.score(...)*
classify LLM.classify(...)
score LLM.score(...)

* 如果模型不支持score任务，LLM.score(...) API会回退到embed任务。

池化器配置

预定义模型

如果模型定义的Pooler接受pooler_config，您可以通过--override-pooler-config选项覆盖其部分属性。

已转换的模型

如果模型已通过--convert转换（如上所述），默认情况下分配给每个任务的池化器具有以下属性：

任务 池化类型 归一化 Softmax
reward ALL ❌ ❌
embed LAST ✅︎ ❌
classify LAST ❌ ✅︎

加载Sentence Transformers模型时，其Sentence Transformers配置文件(modules.json)会优先于模型的默认设置。

你可以通过--override-pooler-config选项进一步自定义此设置，该选项优先级高于模型和Sentence Transformers的默认配置。

离线推理

LLM 类提供了多种离线推理的方法。初始化模型时的选项列表请参阅configuration。

LLM.embed

embed 方法会为每个提示词输出一个嵌入向量，主要设计用于嵌入模型。

from vllm import LLM

llm = LLM(model="intfloat/e5-small", runner="pooling")
(output,) = llm.embed("Hello, my name is")

embeds = output.outputs.embedding
print(f"Embeddings: {embeds!r} (size={len(embeds)})")

代码示例可在此处找到： examples/offline_inference/basic/embed.py

LLM.classify

classify 方法为每个提示输出一个概率向量。它主要设计用于分类模型。

from vllm import LLM

llm = LLM(model="jason9693/Qwen2.5-1.5B-apeach", runner="pooling")
(output,) = llm.classify("Hello, my name is")

probs = output.outputs.probs
print(f"Class Probabilities: {probs!r} (size={len(probs)})")

代码示例可在此处找到： examples/offline_inference/basic/classify.py

LLM.score

score 方法输出句子对之间的相似度分数。它专为嵌入模型和交叉编码器模型设计。嵌入模型使用余弦相似度，而cross-encoder models在RAG系统中充当候选查询-文档对的重新排序器。

注意

vLLM只能执行RAG的模型推理组件（例如嵌入、重排序）。要在更高层次上处理RAG，您应该使用集成框架如LangChain。

from vllm import LLM

llm = LLM(model="BAAI/bge-reranker-v2-m3", runner="pooling")
(output,) = llm.score("What is the capital of France?",
                      "The capital of Brazil is Brasilia.")

score = output.outputs.score
print(f"Score: {score}")

代码示例可在此处找到： examples/offline_inference/basic/score.py

LLM.reward

vLLM中所有奖励模型均可使用reward方法。该方法直接返回提取的隐藏状态。

from vllm import LLM

llm = LLM(model="internlm/internlm2-1_8b-reward", runner="pooling", trust_remote_code=True)
(output,) = llm.reward("Hello, my name is")

data = output.outputs.data
print(f"Data: {data!r}")

代码示例可在此处找到： examples/offline_inference/basic/reward.py

LLM.encode

vLLM中所有池化模型均可使用encode方法，该方法直接返回提取的隐藏状态。

注意

请在使用LLM.encode时选择更具体的方法或直接设置任务：

  • 对于嵌入任务，使用 LLM.embed(...) 或 pooling_task="embed"。
  • 对于分类逻辑值，使用LLM.classify(...)或pooling_task="classify"。
  • 对于奖励，使用 LLM.reward(...) 或 pooling_task="reward"。
  • 对于相似度评分，使用LLM.score(...)。
from vllm import LLM

llm = LLM(model="intfloat/e5-small", runner="pooling")
(output,) = llm.encode("Hello, my name is", pooling_task="embed")

data = output.outputs.data
print(f"Data: {data!r}")

在线服务

我们的OpenAI兼容服务器提供了与离线API对应的端点：

套娃嵌入

Matryoshka Embeddings 或 Matryoshka Representation Learning (MRL) 是一种用于训练嵌入模型的技术。它允许用户在性能和成本之间进行权衡。

警告

并非所有嵌入模型都采用Matryoshka表示学习进行训练。为避免误用dimensions参数，vLLM会对试图更改不支持Matryoshka嵌入模型的输出维度的请求返回错误。

例如，在使用BAAI/bge-m3模型时设置dimensions参数会导致以下错误。

{"object":"error","message":"Model \"BAAI/bge-m3\" does not support matryoshka representation, changing output dimensions will lead to poor results.","type":"BadRequestError","param":null,"code":400}

手动启用Matryoshka嵌入

目前没有官方接口用于指定对Matryoshka Embeddings的支持。在vLLM中，如果config.json里的is_matryoshka参数设为True，则允许将输出改为任意维度。使用matryoshka_dimensions可以控制允许的输出维度。

对于支持Matryoshka嵌入但未被vLLM识别的模型，请手动使用hf_overrides={"is_matryoshka": True}、hf_overrides={"matryoshka_dimensions": []}（离线）或--hf_overrides '{"is_matryoshka": true}'、--hf_overrides '{"matryoshka_dimensions": []}'（在线）覆盖配置。

以下是一个启用Matryoshka Embeddings来服务模型的示例。

vllm serve Snowflake/snowflake-arctic-embed-m-v1.5 --hf_overrides '{"matryoshka_dimensions":[256]}'

离线推理

您可以通过在PoolingParams中使用dimensions参数来更改支持Matryoshka Embeddings的嵌入模型的输出维度。

from vllm import LLM, PoolingParams

llm = LLM(model="jinaai/jina-embeddings-v3",
          runner="pooling",
          trust_remote_code=True)
outputs = llm.embed(["Follow the white rabbit."],
                    pooling_params=PoolingParams(dimensions=32))
print(outputs[0].outputs)

代码示例可在此处找到： examples/offline_inference/embed_matryoshka_fy.py

在线推理

使用以下命令启动vllm服务器。

vllm serve jinaai/jina-embeddings-v3 --trust-remote-code

您可以通过使用dimensions参数来更改支持Matryoshka Embeddings的嵌入模型的输出维度。

curl http://127.0.0.1:8000/v1/embeddings \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
    "input": "Follow the white rabbit.",
    "model": "jinaai/jina-embeddings-v3",
    "encoding_format": "float",
    "dimensions": 32
  }'

预期输出：

{"id":"embd-5c21fc9a5c9d4384a1b021daccaf9f64","object":"list","created":1745476417,"model":"jinaai/jina-embeddings-v3","data":[{"index":0,"object":"embedding","embedding":[-0.3828125,-0.1357421875,0.03759765625,0.125,0.21875,0.09521484375,-0.003662109375,0.1591796875,-0.130859375,-0.0869140625,-0.1982421875,0.1689453125,-0.220703125,0.1728515625,-0.2275390625,-0.0712890625,-0.162109375,-0.283203125,-0.055419921875,-0.0693359375,0.031982421875,-0.04052734375,-0.2734375,0.1826171875,-0.091796875,0.220703125,0.37890625,-0.0888671875,-0.12890625,-0.021484375,-0.0091552734375,0.23046875]}],"usage":{"prompt_tokens":8,"total_tokens":8,"completion_tokens":0,"prompt_tokens_details":null}}

一个OpenAI客户端示例可以在这里找到： examples/online_serving/openai_embedding_matryoshka_fy.py

优云智算