小白视角：利用 SGL 来 Serve Embedding Model

书接上文，继续入坑 SGLang 项目：

需求描述

这段时间的项目需要利用 embedding 进行一些简单的聚类算法。作为一个长期把模型当做黑盒使用的前 NLP 小将，很显然我们可以通过如下的方式来获得字符串的 embedding：

from sentence_transformers import SentenceTransformermodel = SentenceTransformer("Alibaba-NLP/gte-Qwen2-7B-instruct", trust_remote_code=True)documents = [
    "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
    "Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.",]document_embeddings = model.encode(documents)

transformers 的使用非常简单，但是可想而知效率非常低下，我们有没有现代方法来 serve Embedding Model 呢？当然，为此我介绍两个框架 Infinity 和 SGLang 并对比二者的优缺点。（注意到 vllm 也是支持 serve Embedding Model 的，但是囿于时间，我没有深入研究）

Infinity

我在此处给出 infinity 的参考用法（说实话我个人觉得 Infinty 的文档做的比较堪忧，甚至找到 serve 方法都很绕）。首先，我们尝试 serve 一个 Alibaba-NLP/gte-Qwen2-7B-instruct 模型，在我本地名为 qw2_7b_emb：

CUDA_VISIBLE_DEVICES=1 infinity_emb v2 --model-id qw2_7b_emb \ --port 7777 --api-key HeiShenHuaWuKongVWo300

接着调用方法：

from openai import OpenAIemb_client = OpenAI(base_url="http://127.0.0.1:7777/", api_key="HeiShenHuaWuKongVWo300")def online_emb(input_str):
    ret = emb_client.embeddings.create(input=input_str, model="qw2_7b_emb").data[0].embedding
    return ret

看上去非常美好，实际上小规模用用也很美好，但是，当进一步加大调用量以压榨显卡性能时，问题就出现了：

粗糙的显存管理

细心观察前面的 serve 指令，我们惊奇的发现居然 serving 的时候没有类似 vllm 或者 SGLang 里面的 memory upper bound 参数（也即 vllm 里面的 gpu_memory_utilization 和 SGLang 里面的 mem_fraction_static）。这个事情实际上反而是寻常的，因为 infinity 的作者没有实现 PageAttention 或者 RadixAttention 这类的算法，而改进的 attention 算法是其他二者的一大特色。

粗糙的显存管理导致了 serving 时显著不稳定。对于 vllm 和 SGLang，启动了 server 后显存占用几乎是固定的（实际上会有少许波动，因此对 SGLang memory upper bound 不宜高于 0.87）。采用比较保守的参数后，serving 服务相当的稳定；然而 Infinity 的 serving 就显著存在稳定性问题。

我在 48G A6000 上用我的上述参数启动了 Infinity server 后，显存占去了 32G（实际上与直接用 sentence_former 的占用一致，我合理怀疑 Infinity 仅仅是大量复用了 HuggingFace 的源码而添加了 API 接口）。然而，大量的请求之后，会直接 OOM。OOM 后，serving 并不会断掉，反而是进入了“可以接到请求但是无法返回任何内容”的垃圾状态。（我个人理解 OpenAI 的 client 利用了类似三次握手的机制，如果第一次握手能成功就会等待，而不成功则会直接返回 Connection Error。对于 SGLang 和 vllm 而言，发生 OOM 后 client 会直接得到 error，而 Infinity 居然能握手一次导致直接陷入垃圾等待状态）

看似完全的模型支持

Infinity 据称可以支持任意的 embedding model，因此我在 7b 的 gte model OOM 后立刻改用了 1.5B 的 gte model。我心想此时应该不会出现 OOM 问题，然而更大的问题出现了，仅仅是将 serve 的模型改为 qw2_1.5b_emb 并更改了 API 请求里的 model 后，embedding 的返回值居然完全成了 [None, ...]

总之，这两点问题的存在，让我直接放弃了这个框架。当然，我的描述可能存在很多不准确内容，如果有对 Infinity 较为熟悉的读者或者作者团队的一员，请一定指正。

改用 SGLang serve Embeeding Model

我立刻改用了 SGLang 来进行 serving，请注意我的模型从前文的 gte-7b 更改为了 intfloat/e5-mistral-7b-instruct（本地命名为 e5_7b，更换模型的原因会在后文叙述）：

CUDA_VISIBLE_DEVICES=3 python -m sglang.launch_server --enable-p2p-check \
 --model-path e5_7b --dtype auto --tensor-parallel-size 1 \
 --context-length 4096 --chunked-prefill-size 512  --port 7798 \
 --host 0.0.0.0 --api-key HeiShenHuaWuKongVWo300 \
 --enable-torch-compile  --mem-fraction-static 0.87

然后我们使用如下的代码来得到 embedding 结果：

from openai import OpenAI

emb_client = OpenAI(base_url="http://127.0.0.1:7798/v1", api_key="HeiShenHuaWuKongVWo300")


def online_emb(input_str):
    ret = emb_client.embeddings.create(input=input_str, model="e5_7b").data[0].embedding
    return ret

请注意 base_url 的结尾有了 v1 字样，事实上这才该是 OpenAI 的标准接口模样。（我不理解为什么 Infinity 居然没有这个 v1 标识）。

然后，SGLang 的优点就比较明确了：

基于这两点，当然我毫不犹豫的使用了 SGLang（囿于时间原因没有继续测试 vllm，当然 vllm 也支持 serve Embedding Model）。

当然，SGLang 目前在 serve Embedding Model 上也存在较多问题：

试试看让 SGLang support gte 模型呢？

当然，我已经尝试过了，不过目前暂未完成

我可以简单描述下这个过程的合理思路：

2. 将模型 clone 到本地，然后不断地在 modeling_qwen.py 文件下的 class 去 raise exception，试图找到真正 load model 的是哪些个类（或者说找到模型的 architecture，而 architecture 的定义在下文给出）。比如说，我截取两个 class 并且塞入 exception：

class Qwen2ForCausalLM(Qwen2PreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config):
        super().__init__(config)
        self.model = Qwen2Model(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        raise "Use Qwen2ForCausalLM"
        # Initialize weights and apply final processing
        self.post_init()


class Qwen2Model(Qwen2PreTrainedModel):
    """
    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`]
    Args:
        config: Qwen2Config
    """

    def __init__(self, config: Qwen2Config):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        raise "Use Qwen2Model"

3. 接着，我们调用如下的代码去启动 gte 模型，查看哪个类被真正的调用了：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("Alibaba-NLP/gte-Qwen2-7B-instruct", trust_remote_code=True)

documents = [
    "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
    "Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.",
]

document_embeddings = model.encode(documents)

注意这里 trust_remote_code 的用法，一定要是 True。这个参数的意义是是否调用本地的 modeling 文件，这个 remote 实际上是针对 HuggingFace 的 remote，也即我们本地。

4. 找到了调用 gte 的真正类后，将这个类按照 SGLang 的 support model 文档进行迁移（这个文档藏得很深 ）。并且，大量的 class 是可以复用的，实际上并不一定要从 HuggingFace 源码迁移到 SGLang 上面，而是考虑复用 SGLang 已经写好的类（譬如 gte 可以复用 SGLang 已经写好的 Qwen2 类）或者从 vllm 迁移到 SGLang 上：

5. 大功告成，恭喜！

为什么我没成功迁移？

虽然前文详述了如何迁移的思路，但是我并没有成功利用 SGLang support 起 ste，没成功的原因比较有意思。这里先贴下我 support 失败但是留下的残骸：

结合我的残骸，我实现了python/sglang/srt/models/qwen2_embedding.py这个文件后，在单测test/srt/models/test_embedding_models.py时，ModelRunner 会进入到python/sglang/srt/models/qwen2.py下的Qwen2ForCausalLM而非python/sglang/srt/models/qwen2_embedding.py下的Qwen2EmbeddingModel。这件事情让我非常惊讶，然后我才发现 gte model 和 e5 的一大区别。

在描述原因之前，我先不严谨的提出 architecture 和 class 的区别。下文中的 architecture 指的是模型的 config 文件下的 architecture 字段；而 class 指的是 SGLang 的 models file 下的一个 class，这个 class 的 forward 函数用于实现某个确定的目的，要么是 get_completion 要么是 get_embedding（除非添加额外的参数）。

给出了 architecture 和 class 的定义后，我可以用一句话概括 support gte 失败的原因：

SGLang 决定某个模型的 class 依靠的是模型的 architecture 到 SGLang 框架所写的 class 之间的单射。当多个 completion model （或者多个 embedding model）共用一个 architecture 时，可以唯一地将这个 architecture 映射到一个 class 上，成功 supprot；而 completion model 和 completion model 共用 architecture 时，这个 architecture 只能被映射到一个 class 上，这个 class 的 forward 函数要么是返回 completion，要么是返回 embedding，不可能在不添加额外参数的情况下有时返回 completion，有时返回 embedding，因此这个 architecture 带来了 class 冲突。

gte 模型的 architecture 和已有 completion 模型冲突

按照前文所述，我进一步参数 gte 模型带来的 architecture 冲突。具体来说，我用modeling_qwen.py的源码按照前文方法 raise error，发觉 gte 调用的模型是Qwen2Model而不是Qwen2ForCausalLM，因此我初步认为 gte 的 architecture 应该是 Qwen2Model。但是，查阅 gte model 的 model config：

在这个 config 里面，gte 的 architecture 是:

  "architectures": [
    "Qwen2ForCausalLM"
  ],

接着，进一步观察：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("gte-Qwen2-7B-instruct", trust_remote_code=True)
print(model)

得到的结果是：

SentenceTransformer(
  (0): Transformer({'max_seq_length': 32768, 'do_lower_case': False}) with Transformer model: Qwen2Model 
  (1): Pooling({'word_embedding_dimension': 3584, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': True, 'include_prompt': True})
  (2): Normalize()
)

这里显示的 Model 是 Qwen2Model。我逐渐发觉了 gte model 的 architecture 的矛盾之处，也即 gte model 实际上的 architecture 是混乱的，在 config 里调用的 architecture 是 Qwen2ForCausalLM。而 SGLang 框架将模型映射到 model class 的方法是依靠 config 下的 architecture 进行映射。

gte 模型 config 里的 architecture 就是Qwen2ForCausalLM，这导致了在目前 SGLang 的设计下没法直接 support gte。因为Qwen2ForCausalLM这个 architecture 已经被映射到已经为 completion 模型实现的Qwen2ForCausalLM类上，没法映射到在我的残骸代码中另一个返回 embedding 的类上，也即我自己定义的Qwen2EmbeddingModel类。

或者说，gte 作为一个 embedding model，其 architecture 和已经被 support 的 completion model 存在冲突。Qwen2ForCausalLM 同时是 completion model 和 embedding model 的 architecture，我们已经将 Qwen2ForCausalLM 这个 architecture 映射到了为 completion model 写的 Qwen2ForCausalLM 类上，因而没法把 Qwen2ForCausalLM 这个 architecture 映射到其他类上。

e5 的 architecture 并不存在冲突

至于 e5，按照这两个途径得到的 class 都是 MistralModel：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("intfloat/e5-mistral-7b-instruct")
# In case you want to reduce the maximum sequence length:
model.max_seq_length = 4096
print(model)

SentenceTransformer(
  (0): Transformer({'max_seq_length': 4096, 'do_lower_case': False}) with Transformer model: MistralModel 
  (1): Pooling({'word_embedding_dimension': 4096, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': True, 'include_prompt': True})
  (2): Normalize()
)

这两个途径下 e5 的 architecture 都是 MistralLM，而且并没有一个 completion model 的 architecture 也是 MistralLM（比较有意思的是，mistralai/Mistral-Large-Instruct-2407 的 architecture 是 MistralForCausalLM，因此 mistral 的 completion model 和 embedding 不存在 architecture 冲突）。

所以，把MistralModel映射到 SGLang 目前实现的LlamaEmbeddingModel上没有问题。或者说 MistralModel 可以被确定的映射到 embedding model 上，并没有任何冲突。

如何去解决 architecture 冲突呢？

如前文所述，gte 暂时不能被 support 是因为其 architecture 和已经被 support 的 completion 模型冲突了。那么，有什么可以方法可以解决这个问题呢？

在讨论解决方法前，回过头来想想，这种冲突真的那么可怕么？对于 SGLang，当然目前是比较可怕的，但是对于 transformers，其实并没有很大问题。因为，对于一个 completion model 而言，直接取出 forward 函数最后一层的 hidden states 返回作为 input text 的 embedding 是完全可行的。而 gte model 就是这个路子：gte 本身就是个 completion model / causal model，只是取了最后一次 hidden state 罢了。因此，gte model 直接把 Qwen2ForCausalLM 作为了 architecture 是偷懒但是仍旧正确的做法，只用继承下 Qwen2ForCausalLM 类然后把 forward 方法重写就行了。

虽然 gte 可以把 Qwen2ForCausalLM 作为 architecture，但是这无疑给 SGLang 依靠 architecture 到 class 的单射的这一设计带来了麻烦。而解决方法似乎呼之欲出——加参数嘛。

比如前文说，一个 class 的 forward 要么返回 completion，要么返回 embedding，但是加上一个类似 get_embedding 的参数就可以决定这个 forward 的返回类型了。听上去没法根据模型名字直接决定返回 embedding 还是 completion 似乎并不优雅，但是反过来想想——任何一个 completion model 都可以把最后一层 hidden state 作为 embedding 返回（还需要 pooling），加上 get_embedding 参数后可以说更符合这一意图。让用户手动确定是否想要用这个 completion model 做 embedding model。

实际上，无论是 vllm 还是 huggingface，都采用了这个思路。vllm 直接有一个参数指出要不要 embed（具体参数容我找找），而 huggingface 通过 sentence_former 和 transformers 这两个库来区分是得到 embedding 还是 completion。

这可能是我目前觉得让 SGLang 能够 support gte 的最直接办法。剩下部分且待下回，对于上文的所有疏漏，烦请批评指正，感谢各位读者！