vLLM 部署 Embedding 模型

[QingyaFan]

为什么选择 vLLM

部署一个 Embedding 服务有很多方式：

1. 轻量级、快速开发：LangChain + FastAPI、Flask
2. 高性能：vLLM、Triton
3. 企业级部署：BentoML、Kubernetes
4. 完全兼容 OpenAI 接口：vLLM、OpenLLM

用vLLM 既能得到高性能，还能得到兼容 OpenAI 的接口，完全可以平替 OpenAI的服务。

模型选择

Embedding模型如何选择？

嵌入向量的维度

嵌入向量（Embedding Vector）在机器学习和深度学习中广泛用于将离散的类别（如单词、项目等）映射到连续的向量空间，嵌入向量维度是一个关键的超参数。（TODO：什么是超参数，和其他参数区别是啥？）。嵌入向量维度越高，那么：

• 嵌入层的参数量越大（几乎每个LLM都有嵌入层）：假设有V个不同的类别（例如词汇表大小为V），每个嵌入向量的维度为 D，那么嵌入层的总参数量为 V * D.
• 模型复杂度越高：更多的参数使有更强的表达能力，能拟合更复杂的数据模式。然而，这也意味着更容易记住训练中的细节和噪声，不仅仅是数据的普遍规律。

嵌入向量维度低则反之，有更强的泛化能力，计算成本低，但可能无法表达复杂的语义。这是一个权衡，根据不同的数据集、计算资源限制，需要做一些实验来确定最优设置。

模型的性能

Embedding 模型的性能评测一般使用 MTEB(Massive Text Embedding Benchmark)，MTEB 提供统一的评测标准和多样的任务集合。

注意 MTEB 中并未包含处理速度的指标，但是我们可以扩充。（TODO：如何扩充？）

嵌入模型的一般选择

对于词嵌入（如Word2Vec、GloVe），常用维度在100到300之间。对于更复杂的上下文嵌入（如BERT的词向量），维度可能在768甚至更高。

嵌入模型的两个重要参数：

1. embedding_length（嵌入维度）：每个输入元素（例如单词、字词或其他标记）在嵌入空间中的表示维度，每个输入元素会被转换成一个长度为 embedding_length 的向量。虽然，更大的嵌入维度使模型能表示更复杂的模式，但同样会消耗更多的资源。
2. context_length（上下文长度）：一次处理过程中能够处理的最大输入长度，通常以标记（tokens）来衡量。也就是说，模型一次能够处理 context_length 个tokens。context_length 主要影响上下文捕捉，更长的上下文长度使模型能够理解和记住更大范围的文本内容，有助于理解长距离依赖关系，提高生成和理解任务的性能。但是，较长的上下文会消耗更多的内存和计算资源，尤其是Transformer 架构下的模型的计算复杂度跟上下文的平方相关。

一般来说，很多预训练模型（如BERT-base）中，嵌入维度经常是768，这是一个经验上的折中。

从 Huggingface 看

1. jinaai/jina-embeddings-v3

jina-embeddings-v3: Multilingual Embeddings With Task LoRA

这里的关键词是 LoRA（Low-Rank Adaptation，低秩适配），这是一种微调预训练LLM的方法。传统的微调方法需要调整模型的大量参数，这在计算资源和存储上很昂贵。LoRA 通过引入低秩矩阵来代替部分原有的高维权重矩阵，从而减少需要训练和存储的参数数量，同时保证模型的性能。

什么是低秩矩阵？

低秩矩阵（Low-Rank Matrix）是秩较低的矩阵。矩阵的秩是指矩阵中线性独立的行或列的最大数量，秩衡量了一个矩阵中信息的独特性和复杂性，一个全秩矩阵的秩等于其行数或列数中的较小值，而一个低秩矩阵的秩显著低于这个最大值。

低秩矩阵的特点是参数较少、信息压缩、近似表示：

• 参数较少：由于低秩矩阵的秩较低，它可以通过较小的子矩阵来表示，这意味着需要存储和计算的参数数量大大减少；
• 信息压缩：低秩矩阵能够压缩原始数据中的主要信息，同时舍弃冗余或噪声部分，有助于提高计算效率和降低存储成本；
• 近似表示：在许多应用中，低秩矩阵被用来近似高秩矩阵，从而在保障性能的前提下实现简化。

jinaai/jina-embeddings-v3 一些关键参数如下：

architectures                           XLMRobertaModel
num_attention_heads            16
max_position_embeddings   8194
torch_dtype                             bfloat16
vocab_size                               250002
matryoshka_dimensions         [32, 64, 128, 256, 512, 768, 1024]

max_position_embeddings 意味着可以处理 8194 个tokens。
matryoshka_dimensions 表明模型在不同的层使用不同的嵌入维度，维度范围从 32 ~ 1024 不等。

2. nvidia/NV-Embed-v2

We present NV-Embed-v2, a generalist embedding model that ranks No. 1 on the Massive Text Embedding Benchmark (MTEB benchmark)(as of Aug 30, 2024) with a score of 72.31 across 56 text embedding tasks. It also holds the No. 1 in the retrieval sub-category (a score of 62.65 across 15 tasks) in the leaderboard, which is essential to the development of RAG technology.
NV-Embed-v2 presents several new designs, including having the LLM attend to latent vectors for better pooled embedding output, and demonstrating a two-staged instruction tuning method to enhance the accuracy of both retrieval and non-retrieval tasks. Additionally, NV-Embed-v2 incorporates a novel hard-negative mining methods that take into account the positive relevance score for better false negatives removal.

nvidia/NV-Embed-v2在检索子类别中排名第一，这项能力对于RAG很关键。检索性能来自于三点设计：1. 潜在向量关注机制；2. 两阶段指令调优；3. 硬负采样方法。模型的一些关键参数如下：

embedding dimension: 4096

3. intfloat/multilingual-e5-large-instruct

This model has 24 layers and the embedding size is 1024.

从 Ollama 看

1. nomic-embed-text

general.architecture                         nomic-bert
nomic-bert.context_length              2048
nomic-bert.embedding_length        768
tokenizer.ggml.model                        bert
token_embd.weight                           F16    [768, 30522]
nomic-bert.attention.head_count    12

2. mxbai-embed-large

general.architecture        bert
bert.embedding_length  1024
bert.context_length         512
tokenizer.ggml.model      bert
token_embd.weight         F16    [1024, 30522]
bert.attention.head_count     16

3. snowflake-arctic-embed

general.architecture        bert
bert.context_length         512
bert.embedding_length  1024
tokenizer.ggml.model      bert
token_embd.weight         F16    [1024, 30522]
bert.attention.head_count    16

以上是Ollama下载量前三的嵌入模型，架构都是 BERT（Bidirectional Encoder Representation from Transformers），head_count 表示用于多头注意力机制（Multi-Head Attention）的注意力头（Attention Heads）数量。

什么是多头注意力

多头注意力机制是Transformer架构中的一个核心组件，多头注意力中的每一头都会计算各个tokens之间的关系，不同头的计算是并行的，每头都是不同的衡量标准。例如一句话的处理中，一个头关注语义关系，另一个头关注语法关系。

一般情况下，隐层的维度会均匀地分为每一头，例如隐层768维，16头，则每一头得到的维度是 768/16 = 48 维用于计算。每一头是相互独立的，都有自己的Q、K、V矩阵及权重。

每个头在并行地完成各自的计算后，通过 concate 操作将各个头的结果向量连接起来，形成一个较大的向量。然后，通过一个线性变换，将拼接后的向量映射回原始的维度。

线性变换是连接不同向量空间的重要工具，通过保持向量加法和标量乘法的结构，线性变换能够简化复杂问题的分析，同时在理论和应用上都有重要的意义。线性变换在数学、物理、工程等领域中具有广泛的应用。例如，在计算机图形学中，图像的旋转、缩放和平移都可以通过线性变换来实现；在微分方程中，线性变换用于线性化系统以便分析其行为；在线性代数中，研究线性变换帮助理解矩阵的性质及其在不同基底下的表现。

RAG中，嵌入模型的向量维度是否需要跟协同的推理模型的向量维度一致

不需要，嵌入模型主要用在检索阶段，生成模型接收的是检索阶段产生的文本结果，因此，两者没什么关系。

使用近似算法 select 出来的内容应该是向量，输入生成模型的是文本，是什么操作使向量到文本的？这是数据库中存储的向量是一个字段，对于同一行的不同字段保存了原文，所以参与运算的是向量，但是 select 出来的结果其实是对应的文本字段。

哪些模型即支持嵌入，又支持推理？

我们常见的模型基本都是既支持嵌入，又支持推理。例如，阿里的Qwen系列、Facebook的Llama系列、OpenAI的GPT系列、Google的PaLM和BERT系列等。但如果使用 vLLM，一个启动的服务只能使用一种能力。

TODO：待部署Embedding模式来验证。

模型部署

能否一个 vLLM 服务同时支持推理和嵌入？

No. 我们用vLLM 更多的用处是部署推理服务，它也支持嵌入，但是一个 vllm 实例只能是二选一。¹ 服务类型通过 --task 参数指定，有三个值：

• auto，默认值，当模型只支持推理或嵌入一种时，会默认选择支持的服务；如果支持两种，就必须指定 generate 或 embedding 了
• generate，推理服务，对外提供 /v1/completions 和 /v1/chat/completions；
• embedding，嵌入服务，对外提供 /v1/embeddings ²。

使用嵌入服务

业界基本遵循 OpenAI 的规范：

POST    https://api.openai.com/v1/embeddings
Request Body:
input, model, encoding_format, dimensions, user
Returns:
A list of [embedding](https://platform.openai.com/docs/api-reference/embeddings/object) objects.

Embedding 对象例子如下：

{
  "object": "embedding",
  "embedding": [
    0.0023064255,
    -0.009327292,
    .... (1536 floats total for ada-002)
    -0.0028842222,
  ],
  "index": 0
}

参考

Footnotes

1. https://docs.vllm.ai/en/v0.6.4.post1/models/engine_args.html ↩

2. https://platform.openai.com/docs/api-reference/embeddings/create ↩