| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from collections.abc import Iterable |
| 6 | |
| 7 | import torch |
| 8 | import torch.nn as nn |
| 9 | |
| 10 | from vllm.config import VllmConfig |
| 11 | from vllm.model_executor.models.interfaces import SupportsLoRA |
| 12 | from vllm.model_executor.models.qwen3 import Qwen3Model |
| 13 | from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix |
| 14 | from vllm.sequence import IntermediateTensors |
| 15 | |
| 16 | WeightItem = tuple[str, torch.Tensor] |
| 17 | |
| 18 | |
| 19 | class VoyageQwen3BidirectionalEmbedModel(nn.Module, SupportsLoRA): |
| 20 | """Qwen3Model + Voyage embedding head + bidirectional attention. |
| 21 | |
| 22 | Checkpoint conventions (HF): |
| 23 | - MLP: gate_proj + up_proj (unfused) |
| 24 | - Attn: q_proj + k_proj + v_proj (unfused) |
| 25 | - Linear head: linear.weight |
| 26 | - Weights prefixed with "model." (e.g., model.layers.0...) |
| 27 | |
| 28 | vLLM Qwen3Model expects: |
| 29 | - mlp.gate_up_proj (fused) |
| 30 | - self_attn.qkv_proj (fused) |
| 31 | - No "model." prefix |
| 32 | """ |
| 33 | |
| 34 | packed_modules_mapping = { |
| 35 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 36 | "gate_up_proj": ["gate_proj", "up_proj"], |
| 37 | } |
| 38 | |
| 39 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 40 | super().__init__() |
| 41 | self.config = vllm_config.model_config.hf_config |
| 42 | self.model = Qwen3Model( |
| 43 | vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") |
| 44 | ) |
| 45 | |
| 46 | # Embedding head (hidden_size -> num_labels, bias=False) |
| 47 | self.linear = nn.Linear( |
| 48 | self.config.hidden_size, |
| 49 | self.config.num_labels, |
| 50 | bias=False, |
| 51 | ) |
| 52 | |
| 53 | self.make_empty_intermediate_tensors = ( |
| 54 | self.model.make_empty_intermediate_tensors |
| 55 | ) |
| 56 | |
| 57 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 58 | return self.model.embed_input_ids(input_ids) |
| 59 | |
| 60 | def forward( |
| 61 | self, |
| 62 | input_ids: torch.Tensor | None, |
| 63 | positions: torch.Tensor, |
| 64 | intermediate_tensors: IntermediateTensors | None = None, |
| 65 | inputs_embeds: torch.Tensor | None = None, |
| 66 | ) -> torch.Tensor: |
| 67 | out = self.model(input_ids, positions, intermediate_tensors, inputs_embeds) |
| 68 | return self.linear(out) |
| 69 | |
| 70 | def load_weights(self, weights: Iterable[WeightItem]) -> set[str]: |
| 71 | loader = AutoWeightsLoader(self) |
| 72 | return loader.load_weights(weights) |