小白视角：利用 vllm serve 新的 Embedding Model

书接上文，今天起来后感谢合作者的点醒，大概发现了如何用 vllm 去 serve embedding model，并且成功利用 serve 了 gte-7b。

vllm 如何处理 embedding/completion 请求？

这里观察两个位于 /vllm/engine/async_llm_engine.py 下的函数（为了方便将部分注释删去）：

async def generate(
        self,
        inputs: PromptInputs,
        sampling_params: SamplingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None
    ) -> AsyncGenerator[RequestOutput, None]:

        async for output in await self.add_request(
                request_id,
                inputs,
                sampling_params,
                lora_request=lora_request,
                trace_headers=trace_headers,
                prompt_adapter_request=prompt_adapter_request,
        ):
            yield LLMEngine.validate_output(output, RequestOutput)

    async def encode(
        self,
        inputs: PromptInputs,
        pooling_params: PoolingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
    ) -> AsyncGenerator[EmbeddingRequestOutput, None]:

        async for output in await self.add_request(
                request_id,
                inputs,
                pooling_params,
                lora_request=lora_request,
                trace_headers=trace_headers,
        ):
            yield LLMEngine.validate_output(output, EmbeddingRequestOutput)

实际上就是每次用 OpenAI 的结构调用 embedding or completion 接口时，会分别调用上方的 encode 函数和 generate 函数，得到 embedding 或者 completion。看上去在 vllm 中，任何一个模型都可以接受 embedding 与 completion 请求。

如何魔改 Qwen2ForCausalLM 来支持 embedding 请求？

直接用 vllm serve gte-7b：

CUDA_VISIBLE_DEVICES=0 vllm serve 7embed --dtype auto --api-key \
sk-1dwqsdv4r3wef3rvefg34ef1dwRv --tensor-parallel-size 1  \
 --max-model-len 32768 --enforce-eager \
 --disable-custom-all-reduce --port 7777 --served-model-name e5_7b

然后发送 embedding 请求，会出错误（pooler not implemented）。

我们进一步观察 vllm 里面 support 的 qwen2 模型（vllm/model_executor/models/qwen2.py）：

class Qwen2ForCausalLM(nn.Module, SupportsLoRA):
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # LoRA specific attributes
    supported_lora_modules = [
        "qkv_proj",
        "o_proj",
        "gate_up_proj",
        "down_proj",
    ]
    embedding_modules = {}
    embedding_padding_modules = []

    def __init__(
        self,
        config: Qwen2Config,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        lora_config: Optional[LoRAConfig] = None,
    ) -> None:
        # TODO (@robertgshaw2): see if this can be moved out
        if (cache_config.sliding_window is not None
                and hasattr(config, "max_window_layers")):
            raise ValueError("Sliding window for some but all layers is not "
                             "supported. This model uses sliding window "
                             "but `max_window_layers` = %s is less than "
                             "`num_hidden_layers` = %s. Please open an issue "
                             "to discuss this feature." % (
                                 config.max_window_layers,
                                 config.num_hidden_layers,
                             ))

        super().__init__()

        self.config = config
        self.lora_config = lora_config

        self.quant_config = quant_config
        self.model = Qwen2Model(config, cache_config, quant_config)

        if config.tie_word_embeddings:
            self.lm_head = self.model.embed_tokens
        else:
            self.lm_head = ParallelLMHead(config.vocab_size,
                                          config.hidden_size,
                                          quant_config=quant_config)

        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> torch.Tensor:
        hidden_states = self.model(input_ids, positions, kv_caches,
                                   attn_metadata, intermediate_tensors)
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
        logits = self.logits_processor(self.lm_head, hidden_states,
                                       sampling_metadata)
        return logits

然后观察 SGLang 里面的 qwen2 model 文件（python/sglang/srt/models/qwen2.py）：

class Qwen2ForCausalLM(nn.Module):
    def __init__(
        self,
        config: Qwen2Config,
        quant_config: Optional[QuantizationConfig] = None,
        cache_config: Optional[CacheConfig] = None,
    ) -> None:
        super().__init__()
        self.config = config
        self.quant_config = quant_config
        self.model = Qwen2BaseModel(config, quant_config=quant_config)
        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
        self.logits_processor = LogitsProcessor(config)

    @torch.no_grad()
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        input_metadata: InputMetadata,
        input_embeds: torch.Tensor = None,
    ) -> torch.Tensor:
        hidden_states = self.model(input_ids, positions, input_metadata, input_embeds)
        return self.logits_processor(
            input_ids, hidden_states, self.lm_head.weight, input_metadata
        )

这里发现了非常有趣的事情，和 SGLang 不同的是，vllm 里每个 model 的 forward 函数仅仅返回了 hidden_states，而 hidden_states 的 logits_processor 是在 compute_logits 函数里实现的。SGLang 的 forward 函数却将 vllm 里面的 forward 和 logits_processor 合在了一起，直接一步返回了 logits。基于如上的设计，vllm 的 generate 请求实际上调用的是 compute_logits 函数，SGLang 的 generate 请求调用的是 forward 函数。

叙述到此处并不能体现出二者的区别，但是考虑到 embedding 请求时，这个事情就颇有意思了。gte 这个模型的 architecture 是 Qwen2ForCausalLM，vllm 将 gte 模型根据 architecture 映射到 Qwen2ForCausalLM 这个类后，处理 embedding 请求时会试图调用 pooler 函数。因此，我们只需要在 vllm 已经实现的 Qwen2ForCausalLM 类下加上 pooler 即可：

class Qwen2ForCausalLM(nn.Module, SupportsLoRA):
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # LoRA specific attributes
    supported_lora_modules = [
        "qkv_proj",
        "o_proj",
        "gate_up_proj",
        "down_proj",
    ]
    embedding_modules = {}
    embedding_padding_modules = []

    def __init__(
        self,
        config: Qwen2Config,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        lora_config: Optional[LoRAConfig] = None,
    ) -> None:
        # TODO (@robertgshaw2): see if this can be moved out
        if (cache_config.sliding_window is not None
                and hasattr(config, "max_window_layers")):
            raise ValueError("Sliding window for some but all layers is not "
                             "supported. This model uses sliding window "
                             "but `max_window_layers` = %s is less than "
                             "`num_hidden_layers` = %s. Please open an issue "
                             "to discuss this feature." % (
                                 config.max_window_layers,
                                 config.num_hidden_layers,
                             ))

        super().__init__()

        self.config = config
        self.lora_config = lora_config

        self.quant_config = quant_config
        self.model = Qwen2Model(config, cache_config, quant_config)

        if config.tie_word_embeddings:
            self.lm_head = self.model.embed_tokens
        else:
            self.lm_head = ParallelLMHead(config.vocab_size,
                                          config.hidden_size,
                                          quant_config=quant_config)

        self.logits_processor = LogitsProcessor(config.vocab_size)
        self._pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
        self.sampler = Sampler()

    def pooler(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> Optional[PoolerOutput]:
        return self._pooler(hidden_states, pooling_metadata)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> torch.Tensor:
        hidden_states = self.model(input_ids, positions, kv_caches,
                                   attn_metadata, intermediate_tensors)
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
        logits = self.logits_processor(self.lm_head, hidden_states,
                                       sampling_metadata)
        return logits

而 pooler 相关的代码可以直接从 vllm/model_executor/models/llama_embedding.py 上 copy。

如此以来，加入 pooler 后，任意一个 architecture 为 Qwen2ForCausalLM 都可以支持 embedding 请求了。道理是这样的：

对于任意一个 architecture 为 Qwen2ForCausalLM 的模型（譬如 Qwen/Qwen2-72B-Instruct 和 Alibaba-NLP/gte-Qwen2-7B-instruct，这两个模型一个是 completion 模型，另一个是 embedding 模型，但是 architecture 都是 Qwen2ForCausalLM），这个模型会被映射到 Qwen2ForCausalLM 这个类上。当用户调用 completion 请求时，engine 会调用 compute_logits 函数，而用户调用 embedding 请求时，engine 会调用 pooler 函数。从而即便是 embedding model 和 completion model 被映射到同一个类上，分别将 embedding 请求和 completion 请求对应到不同的函数上就可以避免冲突。

反过来，考虑 SGLang 的实现。在 SGLang 中，对于一个 server 的 embedding 请求和 completion 请求都会调用 class 的 forward 函数。如同我上篇博文所讲，forward 函数没法在不加参数的情况下区分用户究竟想要得到 embedding 还是 completion。所以 SGLang 暂时没法通过类似的方法进行更改。

完成上述的更改后，回到 vllm 的 vllm/model_executor/models/__init__.py 文件中，将 gte 映射到 Qwen2ForCausalLM 即可。

_EMBEDDING_MODELS = {
    "MistralModel": ("llama_embedding", "LlamaEmbeddingModel"),
    "Qwen2ForCausalLM": ("Qwen2ForCausalLM"),
}

回顾与问题

回顾下利用 vllm support gte 的过程，我们发现 vllm 将 embedding 请求与 completion 请求分设接口的设计极大的帮助了我们扩展接口。而 SGLang 由于两种请求没有分设接口，因此出现了同一个 architecture 没法映射到两个 class（实现两个 forward 函数）的冲突。

下一步打算利用这种方式修改 SGLang 的接口，然后让 SGLang support gte 模型。

当然，vlllm 的 support 似乎还没有完善。我成功 support 起 gte 模型后，尝试将得到的 embedding 与 sentence_former 的 embedding 进行对拍，发现 vllm 返回的 embedding 存在两个问题：

1. hidden state 的维度是 sentence_former 的两倍；
2. hidden state 的偶数维全是 0；
3. hidden state 的数值远大于 sentence_former 的数值，我怀疑是归一化问题，但是目前没有解决。

总之，麻辣系统真是博大精深呀！

END

作者：Chayenne Zhao
来源：GiantPandaCV

推荐阅读

• 小白视角：vllm 迁移到 SGLang 的体验与收获
• 图解大模型训练系列：序列并行3，Ring Attention
• 复旦提出CTA-Net |卷积与Transformer的协同，通过轻量级多尺度特征融合提升视觉识别！
• C#与OpenCV C++导出DLL调用与数据交互

欢迎大家点赞留言，更多Arm技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。