@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e4b8713380
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 11:00 UTC

[New Model]: nomic-embed-text-v2-moe (#17785)

wang.yuqi committed May 11, 2025 at 15:59 UTC e4b8713380cebc6a336d20ee32800d8be3e28877
9 files changed +899 -364
docs/source/models/supported_models.md
+17 -3
@@ -622,7 +622,7 @@ Specified using `--task embed`.
622 * [PP](#distributed-serving)
623 - * `BertModel`
624 * BERT-based
625 - * `BAAI/bge-base-en-v1.5`, etc.
625 + * `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc.
626 *
627 *
628 - * `Gemma2Model`
@@ -635,6 +635,16 @@ Specified using `--task embed`.
635 * `parasail-ai/GritLM-7B-vllm`.
636 * ✅︎
637 * ✅︎
638 +- * `GteModel`
639 + * GteModel
640 + * `Snowflake/snowflake-arctic-embed-m-v2.0`.
641 + *
642 + * ︎
643 +- * `NomicBertModel`
644 + * NomicBertModel
645 + * `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc.
646 + * ︎
647 + * ︎
648 - * `LlamaModel`, `LlamaForCausalLM`, `MistralModel`, etc.
649 * Llama-based
650 * `intfloat/e5-mistral-7b-instruct`, etc.
@@ -647,12 +657,12 @@ Specified using `--task embed`.
657 * ✅︎
658 - * `RobertaModel`, `RobertaForMaskedLM`
659 * RoBERTa-based
650 - * `sentence-transformers/all-roberta-large-v1`, `sentence-transformers/all-roberta-large-v1`, etc.
660 + * `sentence-transformers/all-roberta-large-v1`, etc.
661 *
662 *
663 - * `XLMRobertaModel`
664 * XLM-RoBERTa-based
655 - * `intfloat/multilingual-e5-large`, `jinaai/jina-reranker-v2-base-multilingual`, etc.
665 + * `intfloat/multilingual-e5-large`, `jinaai/jina-reranker-v2-base-multilingual`, `Snowflake/snowflake-arctic-embed-l-v2.0`, `jinaai/jina-embeddings-v3`(see note), etc.
666 *
667 *
668 :::
@@ -670,6 +680,10 @@ For both the 1.5B and 7B variants, you also need to enable `--trust-remote-code`
680 See [relevant issue on HF Transformers](https://github.com/huggingface/transformers/issues/34882).
681 :::
682
683 +:::{note}
684 +`jinaai/jina-embeddings-v3` supports multiple tasks through lora, while vllm temporarily only supports text-matching tasks by merging lora weights.
685 +:::
686 +
687 If your model is not in the above list, we will try to automatically convert the model using
688 {func}`~vllm.model_executor.models.adapters.as_embedding_model`. By default, the embeddings
689 of the whole prompt are extracted from the normalized hidden state corresponding to the last token.
tests/models/language/pooling/mteb_utils.py new
+111
@@ -0,0 +1,111 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import math
3 +from collections.abc import Sequence
4 +
5 +import mteb
6 +import numpy as np
7 +import pytest
8 +
9 +from tests.models.utils import EmbedModelInfo
10 +
11 +# Most models on the STS12 task (See #17175):
12 +# - Model implementation and minor changes in tensor dtype
13 +# results in differences less than 1e-4
14 +# - Different model results in differences more than 1e-3
15 +# 1e-4 is a good tolerance threshold
16 +MTEB_EMBED_TASKS = ["STS12"]
17 +MTEB_EMBED_TOL = 1e-4
18 +
19 +
20 +class VllmMtebEncoder(mteb.Encoder):
21 +
22 + def __init__(self, vllm_model):
23 + super().__init__()
24 + self.model = vllm_model
25 + self.rng = np.random.default_rng(seed=42)
26 +
27 + def encode(
28 + self,
29 + sentences: Sequence[str],
30 + *args,
31 + **kwargs,
32 + ) -> np.ndarray:
33 + # Hoping to discover potential scheduling
34 + # issues by randomizing the order.
35 + r = self.rng.permutation(len(sentences))
36 + sentences = [sentences[i] for i in r]
37 + outputs = self.model.encode(sentences, use_tqdm=False)
38 + embeds = np.array(outputs)
39 + embeds = embeds[np.argsort(r)]
40 + return embeds
41 +
42 +
43 +class OpenAIClientMtebEncoder(mteb.Encoder):
44 +
45 + def __init__(self, model_name: str, client):
46 + super().__init__()
47 + self.model_name = model_name
48 + self.client = client
49 + self.rng = np.random.default_rng(seed=42)
50 +
51 + def encode(self, sentences: Sequence[str], *args, **kwargs) -> np.ndarray:
52 + # Hoping to discover potential scheduling
53 + # issues by randomizing the order.
54 + r = self.rng.permutation(len(sentences))
55 + sentences = [sentences[i] for i in r]
56 +
57 + embeddings = self.client.embeddings.create(model=self.model_name,
58 + input=sentences)
59 + outputs = [d.embedding for d in embeddings.data]
60 + embeds = np.array(outputs)
61 + embeds = embeds[np.argsort(r)]
62 + return embeds
63 +
64 +
65 +def run_mteb_embed_task(encoder, tasks):
66 + tasks = mteb.get_tasks(tasks=tasks)
67 + evaluation = mteb.MTEB(tasks=tasks)
68 + results = evaluation.run(encoder, verbosity=0, output_folder=None)
69 +
70 + main_score = results[0].scores["test"][0]["main_score"]
71 + return main_score
72 +
73 +
74 +def run_mteb_embed_task_st(model_name, tasks):
75 + from sentence_transformers import SentenceTransformer
76 + model = SentenceTransformer(model_name)
77 + return run_mteb_embed_task(model, tasks)
78 +
79 +
80 +def mteb_test_embed_models(hf_runner, vllm_runner, model_info: EmbedModelInfo):
81 + if not model_info.enable_test:
82 + # A model family has many models with the same architecture,
83 + # and we don't need to test each one.
84 + pytest.skip("Skipping test.")
85 +
86 + with vllm_runner(model_info.name,
87 + task="embed",
88 + max_model_len=None,
89 + dtype=model_info.dtype) as vllm_model:
90 +
91 + if model_info.architecture:
92 + assert (model_info.architecture
93 + in vllm_model.model.llm_engine.model_config.architectures)
94 +
95 + vllm_main_score = run_mteb_embed_task(VllmMtebEncoder(vllm_model),
96 + MTEB_EMBED_TASKS)
97 + vllm_dtype = vllm_model.model.llm_engine.model_config.dtype
98 + model_dtype = getattr(
99 + vllm_model.model.llm_engine.model_config.hf_config, "torch_dtype",
100 + vllm_dtype)
101 +
102 + with hf_runner(model_info.name,
103 + is_sentence_transformer=True,
104 + dtype=model_dtype) as hf_model:
105 + st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
106 +
107 + print("VLLM:", vllm_dtype, vllm_main_score)
108 + print("SentenceTransformer:", model_dtype, st_main_score)
109 + print("Difference:", st_main_score - vllm_main_score)
110 +
111 + assert math.isclose(st_main_score, vllm_main_score, rel_tol=MTEB_EMBED_TOL)
tests/models/language/pooling/test_nomic.py new
+47
@@ -0,0 +1,47 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import pytest
4 +
5 +from ...utils import EmbedModelInfo, run_embedding_correctness_test
6 +
7 +MODELS = [
8 + EmbedModelInfo("nomic-ai/nomic-embed-text-v1",
9 + architecture="NomicBertModel",
10 + dtype="float32",
11 + enable_test=True),
12 + EmbedModelInfo("nomic-ai/nomic-embed-text-v1.5",
13 + architecture="NomicBertModel",
14 + dtype="float32",
15 + enable_test=False),
16 + EmbedModelInfo("nomic-ai/nomic-embed-text-v2-moe",
17 + architecture="NomicBertModel",
18 + dtype="float32",
19 + enable_test=True)
20 +]
21 +
22 +
23 +@pytest.mark.parametrize("model_info", MODELS)
24 +def test_models_mteb(hf_runner, vllm_runner,
25 + model_info: EmbedModelInfo) -> None:
26 + from .mteb_utils import mteb_test_embed_models
27 + mteb_test_embed_models(hf_runner, vllm_runner, model_info)
28 +
29 +
30 +@pytest.mark.parametrize("model_info", MODELS)
31 +def test_models_correctness(hf_runner, vllm_runner, model_info: EmbedModelInfo,
32 + example_prompts) -> None:
33 + if not model_info.enable_test:
34 + pytest.skip("Skipping test.")
35 +
36 + with vllm_runner(model_info.name,
37 + task="embed",
38 + dtype=model_info.dtype,
39 + max_model_len=None) as vllm_model:
40 + vllm_outputs = vllm_model.encode(example_prompts)
41 +
42 + with hf_runner(
43 + model_info.name,
44 + dtype=model_info.dtype,
45 + is_sentence_transformer=True,
46 + ) as hf_model:
47 + run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs)
tests/models/language/pooling/test_snowflake_arctic_embed.py
+22 -43
@@ -1,12 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -import pytest
2
4 -from ...utils import EmbedModelInfo, check_embeddings_close
3 +import pytest
4
6 -EMBEDDING_PROMPTS = [
7 - 'what is snowflake?', 'Where can I get the best tacos?', 'The Data Cloud!',
8 - 'Mexico City of Course!'
9 -]
5 +from ...utils import EmbedModelInfo, run_embedding_correctness_test
6
7 MODELS = [
8 EmbedModelInfo("Snowflake/snowflake-arctic-embed-xs",
@@ -45,51 +41,34 @@ MODELS = [
41
42
43 @pytest.mark.parametrize("model_info", MODELS)
48 -@pytest.mark.parametrize("dtype", ["half"])
49 -def test_models(
44 +def test_models_mteb(
45 hf_runner,
46 vllm_runner,
52 - example_prompts,
47 model_info: EmbedModelInfo,
54 - dtype: str,
55 - monkeypatch,
48 ) -> None:
57 - if not model_info.enable_test:
58 - # A model family has many models with the same architecture,
59 - # and we don't need to test each one.
60 - pytest.skip("Skipping test.")
61 -
62 - example_prompts = example_prompts + EMBEDDING_PROMPTS
49 + from .mteb_utils import mteb_test_embed_models
50 + mteb_test_embed_models(hf_runner, vllm_runner, model_info)
51
64 - vllm_extra_kwargs = {
65 - "hf_overrides": {
66 - "is_matryoshka": model_info.is_matryoshka
67 - }
68 - }
52
70 - with hf_runner(model_info.name, dtype=dtype,
71 - is_sentence_transformer=True) as hf_model:
72 - hf_outputs = hf_model.encode(example_prompts)
53 +@pytest.mark.parametrize("model_info", MODELS)
54 +def test_models_correctness(
55 + hf_runner,
56 + vllm_runner,
57 + model_info: EmbedModelInfo,
58 + example_prompts,
59 +) -> None:
60 + if not model_info.enable_test:
61 + pytest.skip("Skipping test.")
62
63 with vllm_runner(model_info.name,
64 task="embed",
76 - dtype=dtype,
77 - max_model_len=None,
78 - **vllm_extra_kwargs) as vllm_model:
79 -
80 - assert (vllm_model.model.llm_engine.model_config.is_matryoshka ==
81 - model_info.is_matryoshka)
82 -
83 - if model_info.architecture:
84 - assert (model_info.architecture
85 - in vllm_model.model.llm_engine.model_config.architectures)
86 -
65 + dtype=model_info.dtype,
66 + max_model_len=None) as vllm_model:
67 vllm_outputs = vllm_model.encode(example_prompts)
68
89 - check_embeddings_close(
90 - embeddings_0_lst=hf_outputs,
91 - embeddings_1_lst=vllm_outputs,
92 - name_0="hf",
93 - name_1="vllm",
94 - tol=1e-2,
95 - )
69 + with hf_runner(
70 + model_info.name,
71 + dtype=model_info.dtype,
72 + is_sentence_transformer=True,
73 + ) as hf_model:
74 + run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs)
tests/models/utils.py
+2 -1
@@ -332,9 +332,10 @@ def matryoshka_fy(tensor: torch.Tensor, dimensions: int):
332
333 class EmbedModelInfo(NamedTuple):
334 name: str
335 - is_matryoshka: bool
335 + is_matryoshka: bool = False
336 matryoshka_dimensions: Optional[list[int]] = None
337 architecture: str = ""
338 + dtype: str = "auto"
339 enable_test: bool = True
340
341
vllm/model_executor/models/bert.py
+29 -238
@@ -11,16 +11,13 @@ from vllm.compilation.decorators import support_torch_compile
11 from vllm.config import CacheConfig, PoolerConfig, VllmConfig
12 from vllm.distributed import get_tensor_model_parallel_world_size
13 from vllm.forward_context import get_forward_context
14 -from vllm.model_executor.layers.activation import (get_act_and_mul_fn,
15 - get_act_fn)
14 +from vllm.model_executor.layers.activation import get_act_fn
15 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
17 - MergedColumnParallelLinear,
16 QKVParallelLinear,
17 RowParallelLinear)
18 from vllm.model_executor.layers.pooler import (CrossEncodingPooler, Pooler,
19 PoolingType)
20 from vllm.model_executor.layers.quantization import QuantizationConfig
23 -from vllm.model_executor.layers.rotary_embedding import get_rope
21 from vllm.model_executor.layers.vocab_parallel_embedding import (
22 VocabParallelEmbedding)
23 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -41,24 +38,19 @@ class BertEmbedding(nn.Module):
38 self.size = config.hidden_size
39 self.word_embeddings = VocabParallelEmbedding(config.vocab_size,
40 config.hidden_size)
44 -
41 + self.position_embeddings = VocabParallelEmbedding(
42 + config.max_position_embeddings, config.hidden_size)
43 self.token_type_embeddings = VocabParallelEmbedding(
44 config.type_vocab_size, config.hidden_size)
45 self.LayerNorm = nn.LayerNorm(config.hidden_size,
46 eps=config.layer_norm_eps)
47 + self.position_ids = nn.Parameter(
48 + torch.empty((1, config.max_position_embeddings)), )
49
50 self.position_embedding_type = config.position_embedding_type
51 - if self.position_embedding_type == "absolute":
52 - self.position_embeddings = VocabParallelEmbedding(
53 - config.max_position_embeddings, config.hidden_size)
54 - self.position_ids = nn.Parameter(
55 - torch.empty((1, config.max_position_embeddings)), )
56 - elif self.position_embedding_type == "rotary":
57 - self.position_embeddings = None
58 - self.position_ids = None
59 - else:
60 - raise ValueError("Only 'absolute' and 'rotary' " +
61 - "position_embedding_type is supported")
51 + if self.position_embedding_type != "absolute":
52 + raise ValueError("Only 'absolute' position_embedding_type" +
53 + " is supported")
54
55 def forward(
56 self,
@@ -72,6 +64,9 @@ class BertEmbedding(nn.Module):
64 # Input embeddings.
65 inputs_embeds = self.word_embeddings(input_ids)
66
67 + # Position embeddings.
68 + position_embeddings = self.position_embeddings(position_ids)
69 +
70 if token_type_ids is None:
71 token_type_ids = torch.zeros(input_shape,
72 dtype=torch.long,
@@ -79,12 +74,7 @@ class BertEmbedding(nn.Module):
74
75 token_type_embeddings = self.token_type_embeddings(token_type_ids)
76
82 - embeddings = inputs_embeds + token_type_embeddings
83 -
84 - if self.position_embedding_type == "absolute":
85 - position_embeddings = self.position_embeddings(position_ids)
86 - embeddings += position_embeddings
87 -
77 + embeddings = inputs_embeds + token_type_embeddings + position_embeddings
78 embeddings = self.LayerNorm(embeddings)
79 return embeddings
80
@@ -108,11 +98,7 @@ class BertPooler(nn.Module):
98 @support_torch_compile
99 class BertEncoder(nn.Module):
100
111 - def __init__(self,
112 - vllm_config: VllmConfig,
113 - bias: bool = True,
114 - rotary_kwargs: Optional[dict] = None,
115 - prefix: str = ""):
101 + def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
102 super().__init__()
103 config = vllm_config.model_config.hf_config
104 cache_config = vllm_config.cache_config
@@ -121,19 +107,16 @@ class BertEncoder(nn.Module):
107 BertLayer(config=config,
108 cache_config=cache_config,
109 quant_config=quant_config,
124 - bias=bias,
125 - rotary_kwargs=rotary_kwargs,
110 prefix=f"{prefix}.layer.{layer_idx}")
111 for layer_idx in range(config.num_hidden_layers)
112 ])
113
114 def forward(
115 self,
132 - positions: torch.Tensor,
116 hidden_states: torch.Tensor,
117 ) -> torch.Tensor:
118 for layer in self.layer:
136 - hidden_states = layer(positions, hidden_states)
119 + hidden_states = layer(hidden_states)
120 return hidden_states
121
122
@@ -143,8 +126,6 @@ class BertLayer(nn.Module):
126 config: BertConfig,
127 cache_config: Optional[CacheConfig] = None,
128 quant_config: Optional[QuantizationConfig] = None,
146 - bias: bool = True,
147 - rotary_kwargs: Optional[dict] = None,
129 prefix: str = ""):
130 super().__init__()
131
@@ -154,36 +135,23 @@ class BertLayer(nn.Module):
135 layer_norm_eps=config.layer_norm_eps,
136 cache_config=cache_config,
137 quant_config=quant_config,
157 - bias=bias,
158 - rotary_kwargs=rotary_kwargs,
138 prefix=f"{prefix}.attention")
139
161 - if config.hidden_act in ["silu", "gelu_and_mul"]:
162 - self.intermediate = BertGatedIntermediate(
163 - hidden_size=config.hidden_size,
164 - intermediate_size=config.intermediate_size,
165 - hidden_act=config.hidden_act,
166 - bias=bias,
167 - quant_config=quant_config,
168 - prefix=f"{prefix}.intermediate")
169 - else:
170 - self.intermediate = BertIntermediate(
171 - hidden_size=config.hidden_size,
172 - intermediate_size=config.intermediate_size,
173 - hidden_act=config.hidden_act,
174 - bias=bias,
175 - quant_config=quant_config,
176 - prefix=f"{prefix}.intermediate")
140 + self.intermediate = BertIntermediate(
141 + hidden_size=config.hidden_size,
142 + intermediate_size=config.intermediate_size,
143 + hidden_act=config.hidden_act,
144 + quant_config=quant_config,
145 + prefix=f"{prefix}.intermediate")
146
147 self.output = BertOutput(hidden_size=config.hidden_size,
148 intermediate_size=config.intermediate_size,
149 layer_norm_eps=config.layer_norm_eps,
181 - bias=bias,
150 quant_config=quant_config,
151 prefix=f"{prefix}.output")
152
185 - def forward(self, positions: torch.Tensor, hidden_states: torch.Tensor):
186 - attn_output = self.attention(positions, hidden_states)
153 + def forward(self, hidden_states: torch.Tensor):
154 + attn_output = self.attention(hidden_states)
155 intermediate_output = self.intermediate(attn_output)
156 output = self.output(intermediate_output, attn_output)
157 return output
@@ -198,8 +166,6 @@ class BertAttention(nn.Module):
166 layer_norm_eps: float,
167 cache_config: Optional[CacheConfig] = None,
168 quant_config: Optional[QuantizationConfig] = None,
201 - bias: bool = True,
202 - rotary_kwargs: Optional[dict] = None,
169 prefix: str = "",
170 ):
171 super().__init__()
@@ -208,22 +174,18 @@ class BertAttention(nn.Module):
174 num_attention_heads=num_attention_heads,
175 cache_config=cache_config,
176 quant_config=quant_config,
211 - bias=bias,
212 - rotary_kwargs=rotary_kwargs,
177 prefix=f"{prefix}.output")
178
179 self.output = BertSelfOutput(hidden_size=hidden_size,
180 layer_norm_eps=layer_norm_eps,
217 - bias=bias,
181 quant_config=quant_config,
182 prefix=f"{prefix}.output")
183
184 def forward(
185 self,
223 - positions: torch.Tensor,
186 hidden_states: torch.Tensor,
187 ) -> torch.Tensor:
226 - self_output = self.self(positions, hidden_states)
188 + self_output = self.self(hidden_states)
189 return self.output(self_output, hidden_states)
190
191
@@ -235,8 +197,6 @@ class BertSelfAttention(nn.Module):
197 num_attention_heads: int,
198 cache_config: Optional[CacheConfig] = None,
199 quant_config: Optional[QuantizationConfig] = None,
238 - bias: bool = True,
239 - rotary_kwargs: Optional[dict] = None,
200 prefix: str = "",
201 ):
202 super().__init__()
@@ -261,15 +221,10 @@ class BertSelfAttention(nn.Module):
221 head_size=self.head_dim,
222 total_num_heads=self.total_num_heads,
223 total_num_kv_heads=self.total_num_kv_heads,
264 - bias=bias,
224 + bias=True,
225 quant_config=quant_config,
226 prefix=f"{prefix}.qkv_proj")
227
268 - if rotary_kwargs:
269 - self.rotary_emb = get_rope(**rotary_kwargs)
270 - else:
271 - self.rotary_emb = None
272 -
228 self.attn = Attention(num_heads=self.num_heads,
229 head_size=self.head_dim,
230 scale=self.scaling,
@@ -281,15 +236,10 @@ class BertSelfAttention(nn.Module):
236
237 def forward(
238 self,
284 - positions: torch.Tensor,
239 hidden_states: torch.Tensor,
240 ) -> torch.Tensor:
241 qkv, _ = self.qkv_proj(hidden_states)
242 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
289 -
290 - if self.rotary_emb:
291 - q, k = self.rotary_emb(positions, q, k)
292 -
243 output = self.attn(q, k, v)
244 return output
245
@@ -299,13 +249,12 @@ class BertSelfOutput(nn.Module):
249 def __init__(self,
250 hidden_size: int,
251 layer_norm_eps: float,
302 - bias: bool = True,
252 quant_config: Optional[QuantizationConfig] = None,
253 prefix: str = ""):
254 super().__init__()
255 self.dense = RowParallelLinear(input_size=hidden_size,
256 output_size=hidden_size,
308 - bias=bias,
257 + bias=True,
258 quant_config=quant_config,
259 prefix=f"{prefix}.dense")
260 self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
@@ -323,13 +272,12 @@ class BertIntermediate(nn.Module):
272 hidden_size: int,
273 intermediate_size: int,
274 hidden_act: str,
326 - bias: bool = True,
275 quant_config: Optional[QuantizationConfig] = None,
276 prefix: str = ""):
277 super().__init__()
278 self.dense = ColumnParallelLinear(input_size=hidden_size,
279 output_size=intermediate_size,
332 - bias=bias,
280 + bias=True,
281 quant_config=quant_config,
282 prefix=f"{prefix}.dense")
283 self.intermediate_act_fn = get_act_fn(hidden_act)
@@ -340,46 +288,19 @@ class BertIntermediate(nn.Module):
288 return hidden_states
289
290
343 -class BertGatedIntermediate(nn.Module):
344 - # for NomciBert and GteModel
345 -
346 - def __init__(self,
347 - hidden_size: int,
348 - intermediate_size: int,
349 - hidden_act: str,
350 - bias: bool = True,
351 - quant_config: Optional[QuantizationConfig] = None,
352 - prefix: str = ""):
353 - super().__init__()
354 - self.act_fn = get_act_and_mul_fn(hidden_act)
355 - self.gate_up_proj = MergedColumnParallelLinear(
356 - hidden_size,
357 - [intermediate_size] * 2,
358 - bias=bias,
359 - quant_config=quant_config,
360 - prefix=f"{prefix}.gate_up_proj",
361 - )
362 -
363 - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
364 - gate_up, _ = self.gate_up_proj(hidden_states)
365 - hidden_states = self.act_fn(gate_up)
366 - return hidden_states
367 -
368 -
291 class BertOutput(nn.Module):
292
293 def __init__(self,
294 hidden_size: int,
295 intermediate_size: int,
296 layer_norm_eps: float,
375 - bias: bool = True,
297 quant_config: Optional[QuantizationConfig] = None,
298 prefix: str = ""):
299 super().__init__()
300
301 self.dense = RowParallelLinear(input_size=intermediate_size,
302 output_size=hidden_size,
382 - bias=bias,
303 + bias=True,
304 quant_config=quant_config,
305 prefix=f"{prefix}.dense")
306
@@ -393,33 +314,18 @@ class BertOutput(nn.Module):
314
315
316 class BertModel(nn.Module, SupportsQuant):
396 - packed_modules_mapping = {
397 - "qkv_proj": ["query", "key", "value"],
398 - "gate_up_proj": [
399 - "gate_proj",
400 - "up_proj",
401 - ],
402 - }
317 + packed_modules_mapping = {"qkv_proj": ["query", "key", "value"]}
318
319 def __init__(self,
320 *,
321 vllm_config: VllmConfig,
322 prefix: str = "",
323 embedding_class: type = BertEmbedding,
409 - bias: bool = True,
410 - rotary_kwargs: Optional[dict] = None,
324 add_pooling_layer: bool = False):
325 super().__init__()
413 - """
414 - For BertModel, all linear layers have bias.
415 - For NomicBertModel, all linear layers do not have bias.
416 - """
417 -
326 config = vllm_config.model_config.hf_config
327 self.embeddings = embedding_class(config)
328 self.encoder = BertEncoder(vllm_config=vllm_config,
421 - bias=bias,
422 - rotary_kwargs=rotary_kwargs,
329 prefix=f"{prefix}.encoder")
330 self.pooler = BertPooler(config) if add_pooling_layer else None
331
@@ -441,7 +347,7 @@ class BertModel(nn.Module, SupportsQuant):
347 seq_lens=attn_metadata.seq_lens_tensor,
348 position_ids=position_ids,
349 token_type_ids=token_type_ids)
444 - return self.encoder(position_ids, hidden_states)
350 + return self.encoder(hidden_states)
351
352 def load_weights(self, weights: Iterable[Tuple[str,
353 torch.Tensor]]) -> Set[str]:
@@ -450,8 +356,6 @@ class BertModel(nn.Module, SupportsQuant):
356 ("qkv_proj", "query", "q"),
357 ("qkv_proj", "key", "k"),
358 ("qkv_proj", "value", "v"),
453 - ("gate_up_proj", "gate_proj", 0),
454 - ("gate_up_proj", "up_proj", 1),
359 ]
360
361 params_dict = dict(self.named_parameters())
@@ -497,7 +401,6 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
401 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
402 super().__init__()
403 pooler_config = vllm_config.model_config.pooler_config
500 - self.config = vllm_config.model_config.hf_config
404 self.model = self._build_model(vllm_config=vllm_config,
405 prefix=maybe_prefix(prefix, "model"))
406 self._pooler = self._build_pooler(pooler_config)
@@ -611,115 +514,3 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding,
514 inputs_embeds=inputs_embeds,
515 intermediate_tensors=intermediate_tensors,
516 token_type_ids=token_type_ids)
614 -
615 -
616 -class NomicBertEmbeddingModel(BertEmbeddingModel):
617 -
618 - hf_to_vllm_mapper = WeightsMapper(
619 - orig_to_new_substr={
620 - "emb_ln": "embeddings.LayerNorm",
621 - "layers": "layer",
622 - "attn.Wqkv": "attention.self.qkv_proj",
623 - "attn.out_proj": "attention.output.dense",
624 - 'norm1': "attention.output.LayerNorm",
625 - 'mlp.fc11': "intermediate.up_proj",
626 - 'mlp.fc12': "intermediate.gate_proj",
627 - 'mlp.fc2': "output.dense",
628 - 'norm2': "output.LayerNorm",
629 - })
630 -
631 - def _build_model(self,
632 - vllm_config: VllmConfig,
633 - prefix: str = "") -> BertModel:
634 - config = vllm_config.model_config.hf_config
635 -
636 - assert config.__class__.__name__ == "NomicBertConfig"
637 - assert config.activation_function == "swiglu"
638 -
639 - # Assume NomicBertModel all linear layers do not have bias
640 - assert not config.mlp_fc1_bias
641 - assert not config.mlp_fc2_bias
642 - assert not config.qkv_proj_bias
643 -
644 - config.layer_norm_eps = config.layer_norm_epsilon
645 - config.position_embedding_type = "rotary"
646 - config.intermediate_size = config.n_inner
647 - config.hidden_act = "silu"
648 - config.hidden_size = config.n_embd
649 - config.num_hidden_layers = config.n_layer
650 -
651 - head_dim = config.hidden_size // config.num_attention_heads
652 - rotary_kwargs = {
653 - "head_size": head_dim,
654 - "rotary_dim": getattr(config, "rotary_emb_dim", head_dim),
655 - "max_position": config.max_trained_positions,
656 - "base": config.rotary_emb_base,
657 - "rope_scaling": {
658 - "rope_type": "dynamic",
659 - "factor": config.rotary_scaling_factor
660 - }
661 - }
662 -
663 - return BertModel(vllm_config=vllm_config,
664 - prefix=prefix,
665 - bias=False,
666 - rotary_kwargs=rotary_kwargs,
667 - embedding_class=BertEmbedding)
668 -
669 -
670 -class GteEmbeddingModel(BertEmbeddingModel):
671 - hf_to_vllm_mapper = WeightsMapper(
672 - orig_to_new_substr={
673 - "attention.qkv_proj": "attention.self.qkv_proj",
674 - "attention.o_proj": "attention.output.dense",
675 - 'attn_ln': "attention.output.LayerNorm",
676 - 'mlp.down_proj': "output.dense",
677 - 'mlp_ln': "output.LayerNorm",
678 - })
679 -
680 - def _build_model(self,
681 - vllm_config: VllmConfig,
682 - prefix: str = "") -> BertModel:
683 - config = vllm_config.model_config.hf_config
684 -
685 - assert config.__class__.__name__ == "GteConfig"
686 - assert config.position_embedding_type == "rope"
687 - assert config.hidden_act == "gelu"
688 -
689 - config.position_embedding_type = "rotary"
690 - config.hidden_act = "gelu_and_mul"
691 -
692 - head_dim = config.hidden_size // config.num_attention_heads
693 - rotary_kwargs = {
694 - "head_size": head_dim,
695 - "rotary_dim": getattr(config, "rotary_emb_dim", head_dim),
696 - "max_position": config.max_position_embeddings,
697 - "base": config.rope_theta,
698 - }
699 -
700 - model = BertModel(vllm_config=vllm_config,
701 - prefix=prefix,
702 - rotary_kwargs=rotary_kwargs,
703 - embedding_class=BertEmbedding)
704 -
705 - # GteModel only gate_up_proj does not have bias.
706 - # Hack method learned from vllm/model_executor/models/glm.py
707 - for layer in model.encoder.layer:
708 - layer.intermediate.gate_up_proj.bias = None
709 - layer.intermediate.skip_bias_add = True
710 - return model
711 -
712 - def split_up_gate_proj(self, weights: Iterable[Tuple[str, torch.Tensor]]):
713 - n = "mlp.up_gate_proj"
714 - for name, weight in weights:
715 - if n in name:
716 - up, gate = weight.chunk(2, dim=0)
717 - yield name.replace(n, "intermediate.up_proj"), up
718 - yield name.replace(n, "intermediate.gate_proj"), gate
719 - else:
720 - yield name, weight
721 -
722 - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
723 - weights = self.hf_to_vllm_mapper.apply(weights)
724 - weights = self.split_up_gate_proj(weights)
725 - self.model.load_weights(weights)
vllm/model_executor/models/bert_with_rope.py new
+652
@@ -0,0 +1,652 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Iterable, Optional, Set, Tuple
3 +
4 +import torch
5 +from torch import nn
6 +from transformers import PretrainedConfig
7 +
8 +from vllm.attention import Attention, AttentionType
9 +from vllm.compilation.decorators import support_torch_compile
10 +from vllm.config import CacheConfig, VllmConfig
11 +from vllm.distributed import get_tensor_model_parallel_world_size
12 +from vllm.model_executor.layers.activation import (get_act_and_mul_fn,
13 + get_act_fn)
14 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
15 + MergedColumnParallelLinear,
16 + QKVParallelLinear,
17 + ReplicatedLinear,
18 + RowParallelLinear)
19 +from vllm.model_executor.layers.quantization import QuantizationConfig
20 +from vllm.model_executor.layers.rotary_embedding import get_rope
21 +from vllm.model_executor.layers.vocab_parallel_embedding import (
22 + VocabParallelEmbedding)
23 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
24 +from vllm.model_executor.models import SupportsV0Only
25 +from vllm.model_executor.models.interfaces import SupportsQuant
26 +from vllm.model_executor.models.utils import WeightsMapper
27 +from vllm.sequence import IntermediateTensors
28 +
29 +
30 +class BertWithRopeEmbedding(nn.Module):
31 +
32 + def __init__(self, config: PretrainedConfig):
33 +
34 + super().__init__()
35 + assert config.type_vocab_size > 0
36 + self.word_embeddings = VocabParallelEmbedding(config.vocab_size,
37 + config.hidden_size)
38 + self.token_type_embeddings = VocabParallelEmbedding(
39 + config.type_vocab_size, config.hidden_size)
40 + self.LayerNorm = nn.LayerNorm(config.hidden_size,
41 + eps=config.layer_norm_eps)
42 +
43 + def forward(
44 + self,
45 + input_ids: torch.Tensor,
46 + token_type_ids: Optional[torch.Tensor] = None,
47 + ) -> torch.Tensor:
48 + input_shape = input_ids.size()
49 + inputs_embeds = self.word_embeddings(input_ids)
50 + if token_type_ids is None:
51 + token_type_ids = torch.zeros(input_shape,
52 + dtype=torch.long,
53 + device=inputs_embeds.device)
54 +
55 + token_type_embeddings = self.token_type_embeddings(token_type_ids)
56 + embeddings = inputs_embeds + token_type_embeddings
57 + embeddings = self.LayerNorm(embeddings)
58 + return embeddings
59 +
60 +
61 +class BertWithRopeAttention(nn.Module):
62 +
63 + def __init__(
64 + self,
65 + hidden_size: int,
66 + num_attention_heads: int,
67 + cache_config: Optional[CacheConfig] = None,
68 + quant_config: Optional[QuantizationConfig] = None,
69 + bias: bool = True,
70 + rotary_kwargs: Optional[dict] = None,
71 + prefix: str = "",
72 + ):
73 + super().__init__()
74 +
75 + self.hidden_size = hidden_size
76 + tp_size = get_tensor_model_parallel_world_size()
77 +
78 + self.total_num_heads = num_attention_heads
79 + assert self.total_num_heads % tp_size == 0
80 +
81 + self.num_heads = self.total_num_heads // tp_size
82 + self.total_num_kv_heads = self.total_num_heads
83 + self.head_dim = self.hidden_size // self.total_num_heads
84 + assert self.head_dim * self.total_num_heads == self.hidden_size
85 +
86 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
87 +
88 + self.q_size = self.num_heads * self.head_dim
89 + self.kv_size = self.num_kv_heads * self.head_dim
90 + self.scaling = self.head_dim**-0.5
91 +
92 + self.qkv_proj = QKVParallelLinear(
93 + hidden_size=self.hidden_size,
94 + head_size=self.head_dim,
95 + total_num_heads=self.total_num_heads,
96 + total_num_kv_heads=self.total_num_kv_heads,
97 + bias=bias,
98 + quant_config=quant_config,
99 + prefix=f"{prefix}.qkv_proj")
100 +
101 + self.rotary_emb = get_rope(**rotary_kwargs)
102 +
103 + self.attn = Attention(num_heads=self.num_heads,
104 + head_size=self.head_dim,
105 + scale=self.scaling,
106 + num_kv_heads=self.num_kv_heads,
107 + cache_config=cache_config,
108 + quant_config=quant_config,
109 + prefix=f"{prefix}.attn",
110 + attn_type=AttentionType.ENCODER_ONLY)
111 +
112 + self.out_proj = RowParallelLinear(input_size=hidden_size,
113 + output_size=hidden_size,
114 + bias=bias,
115 + quant_config=quant_config,
116 + prefix=f"{prefix}.dense")
117 +
118 + def forward(
119 + self,
120 + positions: torch.Tensor,
121 + hidden_states: torch.Tensor,
122 + ) -> torch.Tensor:
123 + qkv, _ = self.qkv_proj(hidden_states)
124 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
125 + q, k = self.rotary_emb(positions, q, k)
126 + attn_output = self.attn(q, k, v)
127 + output, _ = self.out_proj(attn_output)
128 + return output
129 +
130 +
131 +class BertWithRopeGatedMLP(nn.Module):
132 +
133 + def __init__(self,
134 + hidden_size: int,
135 + intermediate_size: int,
136 + hidden_act: str,
137 + bias: bool = True,
138 + quant_config: Optional[QuantizationConfig] = None,
139 + prefix: str = ""):
140 + super().__init__()
141 + self.act_fn = get_act_and_mul_fn(hidden_act)
142 + self.gate_up_proj = MergedColumnParallelLinear(
143 + hidden_size,
144 + [intermediate_size] * 2,
145 + bias=bias,
146 + quant_config=quant_config,
147 + prefix=f"{prefix}.gate_up_proj",
148 + )
149 + self.down_proj = RowParallelLinear(input_size=intermediate_size,
150 + output_size=hidden_size,
151 + bias=bias,
152 + quant_config=quant_config,
153 + prefix=f"{prefix}.down_proj")
154 +
155 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
156 + gate_up, _ = self.gate_up_proj(hidden_states)
157 + hidden_states = self.act_fn(gate_up)
158 + hidden_states, _ = self.down_proj(hidden_states)
159 + return hidden_states
160 +
161 +
162 +class BertWithRopeMLP(nn.Module):
163 +
164 + def __init__(self,
165 + hidden_size: int,
166 + intermediate_size: int,
167 + hidden_act: str,
168 + bias: bool = True,
169 + quant_config: Optional[QuantizationConfig] = None,
170 + prefix: str = ""):
171 + super().__init__()
172 + self.act_fn = get_act_fn(hidden_act)
173 + self.up_proj = ColumnParallelLinear(input_size=hidden_size,
174 + output_size=intermediate_size,
175 + bias=bias,
176 + quant_config=quant_config,
177 + prefix=f"{prefix}.up_proj")
178 + self.down_proj = RowParallelLinear(input_size=intermediate_size,
179 + output_size=hidden_size,
180 + bias=bias,
181 + quant_config=quant_config,
182 + prefix=f"{prefix}.down_proj")
183 +
184 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
185 + hidden_states, _ = self.up_proj(hidden_states)
186 + hidden_states = self.act_fn(hidden_states)
187 + hidden_states, _ = self.down_proj(hidden_states)
188 + return hidden_states
189 +
190 +
191 +class NomicRouter(nn.Module):
192 +
193 + def __init__(self, hidden_size: int, moe_num_experts: int, moe_top_k: int):
194 + super().__init__()
195 + self.moe_top_k = moe_top_k
196 + self.layer = ReplicatedLinear(hidden_size, moe_num_experts, bias=False)
197 +
198 + def forward(
199 + self, x: torch.Tensor
200 + ) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
201 + weights = self.layer(x.view(-1, x.shape[-1]))[0].softmax(
202 + dim=-1, dtype=torch.float32)
203 + top_weights, top_experts = torch.topk(weights, self.moe_top_k, dim=-1)
204 + weights = weights.to(x.dtype)
205 + top_weights = top_weights.to(x.dtype)
206 + return weights, top_weights, top_experts # type: ignore
207 +
208 +
209 +class NomicExpertMLP(nn.Module):
210 +
211 + def __init__(self, hidden_size: int, ffn_hidden_size: int,
212 + moe_num_experts: int, ffn_act_fn: str):
213 + super().__init__()
214 + self.hidden_size = hidden_size
215 + self.ffn_hidden_size = ffn_hidden_size
216 + self.moe_num_experts = moe_num_experts
217 +
218 + self.w1 = nn.Parameter(
219 + torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))
220 + self.w2 = nn.Parameter(
221 + torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))
222 + self.activation_fn = get_act_fn(ffn_act_fn)
223 +
224 + def forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor:
225 + expert_w1 = self.w1.view(self.moe_num_experts, self.ffn_hidden_size,
226 + self.hidden_size)[expert_idx]
227 + expert_w2 = self.w2.view(self.moe_num_experts, self.ffn_hidden_size,
228 + self.hidden_size)[expert_idx]
229 +
230 + x1 = x.matmul(expert_w1.t())
231 + act_out = self.activation_fn(x1)
232 + x2 = act_out.matmul(expert_w2)
233 + return x2
234 +
235 +
236 +class NomicExperts(nn.Module):
237 +
238 + def __init__(self, config, hidden_size: int, ffn_hidden_size: int,
239 + moe_num_experts: int):
240 + super().__init__()
241 + self.moe_num_experts = moe_num_experts
242 +
243 + self.mlp = NomicExpertMLP(hidden_size=config.n_embd,
244 + ffn_hidden_size=config.n_inner,
245 + moe_num_experts=moe_num_experts,
246 + ffn_act_fn=config.hidden_act)
247 + self.bias = nn.Parameter(torch.zeros(config.n_embd))
248 +
249 + def forward(self, x: torch.Tensor, weights: torch.Tensor,
250 + top_weights: torch.Tensor,
251 + top_experts: torch.LongTensor) -> torch.Tensor:
252 + q_len, hidden_size = x.shape
253 + x = x.view(-1, hidden_size)
254 + out = torch.zeros_like(x)
255 +
256 + expert_mask = nn.functional.one_hot(
257 + top_experts, num_classes=self.moe_num_experts).permute(2, 1, 0)
258 + for expert_idx in range(0, self.moe_num_experts):
259 + topk_idx, token_idx = torch.where(expert_mask[expert_idx])
260 + if token_idx.shape[0] == 0:
261 + continue
262 +
263 + token_list = token_idx.tolist()
264 + topk_list = topk_idx.tolist()
265 +
266 + expert_tokens = x[None, token_list].reshape(-1, hidden_size)
267 + expert_out = self.mlp(
268 + expert_tokens, expert_idx) * top_weights[token_list, topk_list,
269 + None]
270 +
271 + out.index_add_(0, token_idx, expert_out)
272 +
273 + out = out.reshape(q_len, hidden_size)
274 + return out + self.bias
275 +
276 +
277 +class NomicMoELayer(nn.Module):
278 +
279 + def __init__(self, config: PretrainedConfig):
280 + super().__init__()
281 +
282 + self.router = NomicRouter(
283 + config.n_embd,
284 + moe_num_experts=config.num_experts,
285 + moe_top_k=config.moe_top_k,
286 + )
287 +
288 + self.experts = NomicExperts(
289 + config,
290 + hidden_size=config.n_embd,
291 + ffn_hidden_size=config.n_inner,
292 + moe_num_experts=config.num_experts,
293 + )
294 +
295 + def forward(self, x: torch.Tensor):
296 + weights, top_weights, top_experts = self.router(x)
297 + out = self.experts(x, weights, top_weights, top_experts)
298 + return out
299 +
300 +
301 +class BertWithRopeBlock(nn.Module):
302 +
303 + def __init__(self,
304 + config: PretrainedConfig,
305 + cache_config: Optional[CacheConfig] = None,
306 + quant_config: Optional[QuantizationConfig] = None,
307 + moe: bool = False,
308 + bias: bool = True,
309 + rotary_kwargs: Optional[dict] = None,
310 + prefix: str = ""):
311 + super().__init__()
312 + self.attn = BertWithRopeAttention(
313 + hidden_size=config.hidden_size,
314 + num_attention_heads=config.num_attention_heads,
315 + cache_config=cache_config,
316 + quant_config=quant_config,
317 + bias=bias,
318 + rotary_kwargs=rotary_kwargs,
319 + prefix=f"{prefix}.attention")
320 +
321 + if moe:
322 + self.mlp = NomicMoELayer(config=config, )
323 + else:
324 + if config.hidden_act in ["silu", "gelu_and_mul"]:
325 + self.mlp = BertWithRopeGatedMLP(
326 + hidden_size=config.hidden_size,
327 + intermediate_size=config.intermediate_size,
328 + hidden_act=config.hidden_act,
329 + bias=bias,
330 + quant_config=quant_config,
331 + prefix=f"{prefix}.mlp")
332 + else:
333 + self.mlp = BertWithRopeMLP(
334 + hidden_size=config.hidden_size,
335 + intermediate_size=config.intermediate_size,
336 + hidden_act=config.hidden_act,
337 + bias=bias,
338 + quant_config=quant_config,
339 + prefix=f"{prefix}.mlp")
340 +
341 + self.attn_ln = nn.LayerNorm(config.hidden_size,
342 + eps=config.layer_norm_eps)
343 + self.mlp_ln = nn.LayerNorm(config.hidden_size,
344 + eps=config.layer_norm_eps)
345 +
346 + def forward(self, positions: torch.Tensor, hidden_states: torch.Tensor):
347 + attn_output = self.attn(positions, hidden_states)
348 + hidden_states = self.attn_ln(hidden_states + attn_output)
349 + mlp_out = self.mlp(hidden_states)
350 + hidden_states = self.mlp_ln(hidden_states + mlp_out)
351 + return hidden_states
352 +
353 +
354 +@support_torch_compile
355 +class BertWithRopeEncoder(nn.Module):
356 +
357 + def __init__(self,
358 + vllm_config: VllmConfig,
359 + bias: bool = True,
360 + rotary_kwargs: Optional[dict] = None,
361 + prefix: str = ""):
362 + super().__init__()
363 + config = vllm_config.model_config.hf_config
364 + cache_config = vllm_config.cache_config
365 + quant_config = vllm_config.quant_config
366 + every_n = getattr(config, "moe_every_n_layers", 0)
367 + self.layers = nn.ModuleList([
368 + BertWithRopeBlock(config=config,
369 + cache_config=cache_config,
370 + quant_config=quant_config,
371 + bias=bias,
372 + moe=every_n > 0 and (layer_idx % every_n == 1),
373 + rotary_kwargs=rotary_kwargs,
374 + prefix=f"{prefix}.layer.{layer_idx}")
375 + for layer_idx in range(config.num_hidden_layers)
376 + ])
377 +
378 + def forward(
379 + self,
380 + positions: torch.Tensor,
381 + hidden_states: torch.Tensor,
382 + ) -> torch.Tensor:
383 + for layer in self.layers:
384 + hidden_states = layer(positions, hidden_states)
385 + return hidden_states
386 +
387 +
388 +class BertWithRope(nn.Module, SupportsV0Only, SupportsQuant):
389 + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
390 +
391 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
392 + super().__init__()
393 + self.config = self.config_verify(vllm_config)
394 + self.embeddings = BertWithRopeEmbedding(self.config)
395 + self.encoder = BertWithRopeEncoder(
396 + vllm_config=vllm_config,
397 + bias=getattr(self.config, "bias", True),
398 + rotary_kwargs=self.config.rotary_kwargs,
399 + prefix=f"{prefix}.encoder")
400 +
401 + def config_verify(self, vllm_config):
402 + raise NotImplementedError
403 +
404 + def forward(
405 + self,
406 + input_ids: Optional[torch.Tensor],
407 + positions: torch.Tensor,
408 + intermediate_tensors: Optional[IntermediateTensors] = None,
409 + inputs_embeds: Optional[torch.Tensor] = None,
410 + token_type_ids: Optional[torch.Tensor] = None,
411 + ) -> torch.Tensor:
412 + if inputs_embeds is not None:
413 + hidden_states = inputs_embeds
414 + else:
415 + hidden_states = self.embeddings(input_ids=input_ids,
416 + token_type_ids=token_type_ids)
417 + return self.encoder(positions, hidden_states)
418 +
419 + def load_weights(self, weights: Iterable[Tuple[str,
420 + torch.Tensor]]) -> Set[str]:
421 + weights = self.hf_to_vllm_mapper.apply(weights)
422 +
423 + if self.config.hidden_act in ["silu", "gelu_and_mul"]:
424 + stacked_params_mapping = [
425 + # (param_name, shard_name, shard_id)
426 + ("gate_up_proj", "gate_proj", 0),
427 + ("gate_up_proj", "up_proj", 1),
428 + ]
429 + else:
430 + stacked_params_mapping = []
431 +
432 + params_dict = dict(self.named_parameters())
433 + loaded_params: Set[str] = set()
434 + for name, loaded_weight in weights:
435 + if "pooler" in name:
436 + continue
437 + for (param_name, weight_name, shard_id) in stacked_params_mapping:
438 + if weight_name not in name:
439 + continue
440 + name = name.replace(weight_name, param_name)
441 + # Skip loading extra bias for GPTQ models.
442 + if name.endswith(".bias") and name not in params_dict:
443 + continue
444 + param = params_dict[name]
445 + weight_loader = param.weight_loader
446 + weight_loader(param, loaded_weight, shard_id)
447 + break
448 + else:
449 + # Skip loading extra bias for GPTQ models.
450 + if name.endswith(".bias") and name not in params_dict:
451 + continue
452 + param = params_dict[name]
453 + weight_loader = getattr(param, "weight_loader",
454 + default_weight_loader)
455 + weight_loader(param, loaded_weight)
456 + loaded_params.add(name)
457 + return loaded_params
458 +
459 +
460 +class NomicBertModel(BertWithRope):
461 + hf_to_vllm_mapper = WeightsMapper(
462 + orig_to_new_substr={
463 + "emb_ln": "embeddings.LayerNorm",
464 + "attn.Wqkv": "attn.qkv_proj",
465 + "norm1": "attn_ln",
466 + "mlp.fc1.": "mlp.up_proj.",
467 + "mlp.fc11": "mlp.up_proj",
468 + "mlp.fc12": "mlp.gate_proj",
469 + "mlp.fc2": "mlp.down_proj",
470 + "norm2": "mlp_ln",
471 + })
472 +
473 + def config_verify(self, vllm_config):
474 + config = vllm_config.model_config.hf_config
475 +
476 + assert config.__class__.__name__ == "NomicBertConfig"
477 + assert config.activation_function in ["swiglu", "gelu"]
478 +
479 + if config.activation_function == "swiglu":
480 + config.hidden_act = "silu"
481 + else:
482 + config.hidden_act = config.activation_function
483 +
484 + assert (config.mlp_fc1_bias == config.mlp_fc2_bias ==
485 + config.qkv_proj_bias)
486 + config.bias = config.qkv_proj_bias
487 +
488 + assert config.rotary_emb_scale_base is None
489 + assert not config.rotary_emb_interleaved
490 +
491 + config.layer_norm_eps = config.layer_norm_epsilon
492 + config.intermediate_size = config.n_inner
493 + config.hidden_size = config.n_embd
494 + config.num_hidden_layers = config.n_layer
495 +
496 + head_dim = config.hidden_size // config.num_attention_heads
497 + rotary_emb_dim = head_dim * config.rotary_emb_fraction
498 + config.rotary_kwargs = {
499 + "head_size": head_dim,
500 + "rotary_dim": rotary_emb_dim,
501 + "max_position": config.max_trained_positions,
502 + "base": getattr(config, "rope_theta", config.rotary_emb_base),
503 + "rope_scaling": getattr(config, "rope_scaling", None)
504 + }
505 +
506 + # we ignore config.rotary_scaling_factor so that for datasets shorter
507 + # than max_trained_positions 2048, the results are consistent
508 + # with SentenceTransformer.
509 + # The context extension uses vllm style rope_theta and rope_scaling.
510 + # See #17785
511 +
512 + return config
513 +
514 +
515 +class GteModel(BertWithRope):
516 + hf_to_vllm_mapper = WeightsMapper(
517 + orig_to_new_substr={
518 + "layer": 'layers',
519 + "attention.qkv_proj": "attn.qkv_proj",
520 + "attention.o_proj": "attn.out_proj",
521 + })
522 +
523 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
524 + super().__init__(vllm_config=vllm_config, prefix=prefix)
525 +
526 + # GteModel only gate_up_proj does not have bias.
527 + # Hack method learned from vllm/model_executor/models/glm.py
528 + for layer in self.encoder.layers:
529 + layer.mlp.gate_up_proj.bias = None
530 + layer.mlp.gate_up_proj.skip_bias_add = True
531 +
532 + def config_verify(self, vllm_config):
533 + config = vllm_config.model_config.hf_config
534 +
535 + assert config.__class__.__name__ == "GteConfig"
536 + assert config.position_embedding_type == "rope"
537 + assert config.hidden_act == "gelu"
538 +
539 + config.position_embedding_type = "rotary"
540 + config.hidden_act = "gelu_and_mul"
541 +
542 + head_dim = config.hidden_size // config.num_attention_heads
543 + config.rotary_kwargs = {
544 + "head_size": head_dim,
545 + "rotary_dim": getattr(config, "rotary_emb_dim", head_dim),
546 + "max_position": config.max_position_embeddings,
547 + "base": config.rope_theta,
548 + "rope_scaling": getattr(config, "rope_scaling", None)
549 + }
550 + return config
551 +
552 + def split_up_gate_proj(self, weights: Iterable[Tuple[str, torch.Tensor]]):
553 + n = "mlp.up_gate_proj"
554 + for name, weight in weights:
555 + if n in name:
556 + up, gate = weight.chunk(2, dim=0)
557 + yield name.replace(n, "mlp.up_proj"), up
558 + yield name.replace(n, "mlp.gate_proj"), gate
559 + else:
560 + yield name, weight
561 +
562 + def load_weights(self, weights: Iterable[Tuple[str,
563 + torch.Tensor]]) -> Set[str]:
564 + weights = self.split_up_gate_proj(weights)
565 + return super().load_weights(weights)
566 +
567 +
568 +class JinaRobertaModel(BertWithRope):
569 + hf_to_vllm_mapper = WeightsMapper(
570 + orig_to_new_substr={
571 + "emb_ln": "embeddings.LayerNorm",
572 + "mixer.Wqkv": "attn.qkv_proj",
573 + "mixer.out_proj": "attn.out_proj",
574 + "norm1": "attn_ln",
575 + "mlp.fc1.": "mlp.up_proj.",
576 + "mlp.fc2": "mlp.down_proj",
577 + "norm2": "mlp_ln",
578 + })
579 +
580 + def config_verify(self, vllm_config):
581 + config = vllm_config.model_config.hf_config
582 + head_dim = config.hidden_size // config.num_attention_heads
583 + config.rotary_kwargs = {
584 + "head_size": head_dim,
585 + "rotary_dim": getattr(config, "rotary_emb_dim", head_dim),
586 + "max_position": config.max_position_embeddings,
587 + "base": getattr(config, "rope_theta", config.rotary_emb_base),
588 + "rope_scaling": getattr(config, "rope_scaling", None)
589 + }
590 + return config
591 +
592 + def forward(
593 + self,
594 + input_ids: torch.Tensor,
595 + position_ids: torch.Tensor,
596 + intermediate_tensors: Optional[IntermediateTensors] = None,
597 + inputs_embeds: Optional[torch.Tensor] = None,
598 + token_type_ids: Optional[torch.Tensor] = None,
599 + ) -> torch.Tensor:
600 + return super().forward(input_ids=input_ids,
601 + positions=position_ids,
602 + intermediate_tensors=intermediate_tensors,
603 + inputs_embeds=inputs_embeds,
604 + token_type_ids=token_type_ids)
605 +
606 + @torch.inference_mode()
607 + def jina_merge_lora_weights(self, weights: Iterable[Tuple[str,
608 + torch.Tensor]]):
609 + # use for jina-embeddings-v3
610 + # Merge Lora weights into a single weight tensor.
611 + # This is a temporary solution until we have a better way to handle
612 +
613 + scaling = self.config.lora_alpha / self.config.lora_rank
614 +
615 + weights = {name: weight for name, weight in weights}
616 +
617 + o = ".original"
618 + a = ".0.lora_A"
619 + b = ".0.lora_B"
620 +
621 + # text-matching
622 + i = -1
623 +
624 + for name in list(weights.keys()):
625 + if o in name:
626 + dtype = weights[name].dtype
627 + shape = weights[name].shape
628 + weight_name = name[:-len(o)]
629 +
630 + if "embeddings" in weight_name:
631 + B = weights[weight_name + a][i].cuda().float()
632 + A = weights[weight_name + b][i].cuda().float()
633 + else:
634 + B = weights[weight_name + b][i].cuda().float()
635 + A = weights[weight_name + a][i].cuda().float()
636 +
637 + weight = (weights[weight_name + o].cuda() +
638 + torch.matmul(B, A).view(shape) * scaling)
639 + weight = weight.cpu().to(dtype)
640 +
641 + weights[weight_name.replace(".parametrizations", "")] = weight
642 +
643 + del weights[weight_name + o], weights[weight_name +
644 + a], weights[weight_name +
645 + b]
646 +
647 + return [(name, weight) for name, weight in weights.items()]
648 +
649 + def load_weights(self, weights: Iterable[Tuple[str,
650 + torch.Tensor]]) -> Set[str]:
651 + weights = self.jina_merge_lora_weights(weights)
652 + return super().load_weights(weights)
vllm/model_executor/models/registry.py
+2 -2
@@ -126,7 +126,7 @@ _EMBEDDING_MODELS = {
126 "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
127 "GlmForCausalLM": ("glm", "GlmForCausalLM"),
128 "GritLM": ("gritlm", "GritLM"),
129 - "GteModel": ("bert", "GteEmbeddingModel"),
129 + "GteModel": ("bert_with_rope", "GteModel"),
130 "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
131 "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"), # noqa: E501
132 "LlamaModel": ("llama", "LlamaForCausalLM"),
@@ -136,7 +136,7 @@ _EMBEDDING_MODELS = {
136 if arch == "LlamaForCausalLM"
137 },
138 "MistralModel": ("llama", "LlamaForCausalLM"),
139 - "NomicBertModel": ("bert", "NomicBertEmbeddingModel"),
139 + "NomicBertModel": ("bert_with_rope", "NomicBertModel"),
140 "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
141 "Qwen2Model": ("qwen2", "Qwen2EmbeddingModel"),
142 "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
vllm/model_executor/models/roberta.py
+17 -77
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 import itertools
4 -from typing import Iterable, Optional, Tuple
4 +from typing import Iterable, Optional, Tuple, Union
5
6 import torch
7 from torch import nn
@@ -19,6 +19,7 @@ from vllm.sequence import IntermediateTensors, PoolerOutput
19 from vllm.transformers_utils.config import (
20 get_cross_encoder_activation_function)
21
22 +from .bert_with_rope import BertWithRope, JinaRobertaModel
23 from .interfaces import SupportsCrossEncoding, SupportsV0Only
24
25
@@ -125,39 +126,20 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
126
127 def _build_model(self,
128 vllm_config: VllmConfig,
128 - prefix: str = "") -> BertModel:
129 + prefix: str = "") -> Union[BertModel, BertWithRope]:
130 if (vllm_config.model_config.hf_config.position_embedding_type ==
131 "rotary"):
131 - config = vllm_config.model_config.hf_config
132 - head_dim = config.hidden_size // config.num_attention_heads
133 -
134 - rotary_kwargs = {
135 - "head_size": head_dim,
136 - "rotary_dim": getattr(config, "rotary_emb_dim", head_dim),
137 - "max_position": config.max_position_embeddings,
138 - "base": config.rotary_emb_base,
139 - "rope_scaling": getattr(config, "rope_scaling", None)
140 - }
141 -
142 - return BertModel(vllm_config=vllm_config,
143 - rotary_kwargs=rotary_kwargs,
144 - prefix=prefix)
132 + return JinaRobertaModel(vllm_config=vllm_config, prefix=prefix)
133 else:
134 return BertModel(vllm_config=vllm_config,
135 prefix=prefix,
136 embedding_class=RobertaEmbedding)
137
138 def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
151 - if getattr(self.config, "lora_rank", 0) > 0:
152 - scaling = self.config.lora_alpha / self.config.lora_rank
153 - weights = jina_merge_lora_weights(weights, scaling)
154 -
139 weights = self.hf_to_vllm_mapper.apply(weights)
140 # Separate weights in "roberta"-prefixed and all else (not in memory).
141 # For use with models like FacebookAI/roberta-base.
142 bert_weights, task_weights = roberta_task_weights_filter(weights)
159 - bert_weights = jina_to_vllm_mapper.apply(bert_weights)
160 -
143 loaded = self.model.load_weights(bert_weights)
144 if not len(loaded):
145 # Fix for models like `sentence-transformers/stsb-roberta-base-v2`
@@ -178,6 +160,18 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
160 _pooler: An instance of Pooler used for pooling operations.
161 """
162
163 + jina_to_vllm_mapper = WeightsMapper(
164 + orig_to_new_substr={
165 + 'emb_ln': "embeddings.LayerNorm",
166 + 'layers': "layer",
167 + 'mixer.Wqkv': "attention.self.qkv_proj",
168 + 'mixer.out_proj': "attention.output.dense",
169 + 'norm1': "attention.output.LayerNorm",
170 + 'mlp.fc1': "intermediate.dense",
171 + 'mlp.fc2': "output.dense",
172 + 'norm2': "output.LayerNorm",
173 + })
174 +
175 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
176 super().__init__()
177 config = vllm_config.model_config.hf_config
@@ -195,7 +189,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
189
190 def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
191 bert_weights, task_weights = roberta_task_weights_filter(weights)
198 - bert_weights = jina_to_vllm_mapper.apply(bert_weights)
192 + bert_weights = self.jina_to_vllm_mapper.apply(bert_weights)
193
194 self.roberta.load_weights(bert_weights)
195
@@ -276,57 +270,3 @@ def roberta_task_weights_filter(
270
271 return encoder_decoder_weights(), ((n, w) for n, w in all_weights2
272 if not n.startswith("roberta."))
279 -
280 -
281 -jina_to_vllm_mapper = WeightsMapper(
282 - orig_to_new_substr={
283 - 'emb_ln': "embeddings.LayerNorm",
284 - 'layers': "layer",
285 - 'mixer.Wqkv': "attention.self.qkv_proj",
286 - 'mixer.out_proj': "attention.output.dense",
287 - 'norm1': "attention.output.LayerNorm",
288 - 'mlp.fc1': "intermediate.dense",
289 - 'mlp.fc2': "output.dense",
290 - 'norm2': "output.LayerNorm",
291 - })
292 -
293 -
294 -@torch.inference_mode()
295 -def jina_merge_lora_weights(weights: Iterable[Tuple[str, torch.Tensor]],
296 - scaling: float = 1.0):
297 - # use for jina-embeddings-v3
298 - # Merge Lora weights into a single weight tensor.
299 - # This is a temporary solution until we have a better way to handle
300 -
301 - weights = {name: weight for name, weight in weights}
302 -
303 - o = ".original"
304 - a = ".0.lora_A"
305 - b = ".0.lora_B"
306 -
307 - # text-matching
308 - i = -1
309 -
310 - for name in list(weights.keys()):
311 - if o in name:
312 - dtype = weights[name].dtype
313 - shape = weights[name].shape
314 - weight_name = name[:-len(o)]
315 -
316 - if "embeddings" in weight_name:
317 - B = weights[weight_name + a][i].cuda().float()
318 - A = weights[weight_name + b][i].cuda().float()
319 - else:
320 - B = weights[weight_name + b][i].cuda().float()
321 - A = weights[weight_name + a][i].cuda().float()
322 -
323 - weight = (weights[weight_name + o].cuda() +
324 - torch.matmul(B, A).view(shape) * scaling)
325 - weight = weight.cpu().to(dtype)
326 -
327 - weights[weight_name.replace(".parametrizations", "")] = weight
328 -
329 - del weights[weight_name + o], weights[weight_name +
330 - a], weights[weight_name + b]
331 -
332 - return [(name, weight) for name, weight in weights.items()]