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]: support GTE NewModel (#17986)
wang.yuqi committed
May 14, 2025 at 16:31 UTC
63ad6222332c8091c1229098a3b0dc5b7f686fa1
11 files changed
+279
-32
docs/source/models/supported_models.md
+16
-2
@@ -701,12 +701,22 @@ Specified using `--task embed`.
701
* ✅︎
702
* ✅︎
703
- * `GteModel`
704
- * GteModel
704
+ * Arctic-Embed-2.0-M
705
* `Snowflake/snowflake-arctic-embed-m-v2.0`.
706
*
707
* ︎
708
+- * `GteNewModel`
709
+ * mGTE-TRM (see note)
710
+ * `Alibaba-NLP/gte-multilingual-base`, etc.
711
+ * ︎
712
+ * ︎
713
+- * `ModernBertModel`
714
+ * ModernBERT-based
715
+ * `Alibaba-NLP/gte-modernbert-base`, etc.
716
+ * ︎
717
+ * ︎
718
- * `NomicBertModel`
709
- * NomicBertModel
719
+ * Nomic BERT
720
* `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc.
721
* ︎
722
* ︎
@@ -749,6 +759,10 @@ See [relevant issue on HF Transformers](https://github.com/huggingface/transform
759
`jinaai/jina-embeddings-v3` supports multiple tasks through lora, while vllm temporarily only supports text-matching tasks by merging lora weights.
760
:::
761
762
+:::{note}
763
+The second-generation GTE model (mGTE-TRM) is named `NewModel`. The name `NewModel` is too generic, you should set `--hf-overrides '{"architectures": ["GteNewModel"]}'` to specify the use of the `GteNewModel` architecture.
764
+:::
765
+
766
If your model is not in the above list, we will try to automatically convert the model using
767
{func}`~vllm.model_executor.models.adapters.as_embedding_model`. By default, the embeddings
768
of the whole prompt are extracted from the normalized hidden state corresponding to the last token.
tests/models/language/pooling/mteb_utils.py
+12
-5
@@ -7,6 +7,7 @@ import numpy as np
7
import pytest
8
9
from tests.models.utils import EmbedModelInfo
10
+from vllm.model_executor.model_loader.utils import set_default_torch_dtype
11
12
# Most models on the STS12 task (See #17175):
13
# - Model implementation and minor changes in tensor dtype
@@ -77,16 +78,22 @@ def run_mteb_embed_task_st(model_name, tasks):
78
return run_mteb_embed_task(model, tasks)
79
80
80
-def mteb_test_embed_models(hf_runner, vllm_runner, model_info: EmbedModelInfo):
81
+def mteb_test_embed_models(hf_runner,
82
+ vllm_runner,
83
+ model_info: EmbedModelInfo,
84
+ vllm_extra_kwargs=None):
85
if not model_info.enable_test:
86
# A model family has many models with the same architecture,
87
# and we don't need to test each one.
88
pytest.skip("Skipping test.")
89
90
+ vllm_extra_kwargs = vllm_extra_kwargs or {}
91
+
92
with vllm_runner(model_info.name,
93
task="embed",
94
max_model_len=None,
89
- dtype=model_info.dtype) as vllm_model:
95
+ dtype=model_info.dtype,
96
+ **vllm_extra_kwargs) as vllm_model:
97
98
if model_info.architecture:
99
assert (model_info.architecture
@@ -99,9 +106,9 @@ def mteb_test_embed_models(hf_runner, vllm_runner, model_info: EmbedModelInfo):
106
vllm_model.model.llm_engine.model_config.hf_config, "torch_dtype",
107
vllm_dtype)
108
102
- with hf_runner(model_info.name,
103
- is_sentence_transformer=True,
104
- dtype=model_dtype) as hf_model:
109
+ with set_default_torch_dtype(model_dtype) and hf_runner(
110
+ model_info.name, is_sentence_transformer=True,
111
+ dtype=model_dtype) as hf_model:
112
st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
113
114
print("VLLM:", vllm_dtype, vllm_main_score)
tests/models/language/pooling/test_gte.py
new
+104
@@ -0,0 +1,104 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import Any
3
+
4
+import pytest
5
+
6
+from ...utils import EmbedModelInfo, run_embedding_correctness_test
7
+
8
+MODELS = [
9
+ ########## BertModel
10
+ EmbedModelInfo("thenlper/gte-large",
11
+ architecture="BertModel",
12
+ dtype="float32",
13
+ enable_test=True),
14
+ EmbedModelInfo("thenlper/gte-base",
15
+ architecture="BertModel",
16
+ dtype="float32",
17
+ enable_test=False),
18
+ EmbedModelInfo("thenlper/gte-small",
19
+ architecture="BertModel",
20
+ dtype="float32",
21
+ enable_test=False),
22
+ EmbedModelInfo("thenlper/gte-large-zh",
23
+ architecture="BertModel",
24
+ dtype="float32",
25
+ enable_test=False),
26
+ EmbedModelInfo("thenlper/gte-base-zh",
27
+ architecture="BertModel",
28
+ dtype="float32",
29
+ enable_test=False),
30
+ EmbedModelInfo("thenlper/gte-small-zh",
31
+ architecture="BertModel",
32
+ dtype="float32",
33
+ enable_test=False),
34
+ ########### NewModel
35
+ EmbedModelInfo("Alibaba-NLP/gte-multilingual-base",
36
+ architecture="GteNewModel",
37
+ enable_test=True),
38
+ EmbedModelInfo("Alibaba-NLP/gte-base-en-v1.5",
39
+ architecture="GteNewModel",
40
+ enable_test=True),
41
+ EmbedModelInfo("Alibaba-NLP/gte-large-en-v1.5",
42
+ architecture="GteNewModel",
43
+ enable_test=True),
44
+ ########### Qwen2ForCausalLM
45
+ EmbedModelInfo("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
46
+ architecture="Qwen2ForCausalLM",
47
+ enable_test=True),
48
+ EmbedModelInfo("Alibaba-NLP/gte-Qwen2-7B-instruct",
49
+ architecture="Qwen2ForCausalLM",
50
+ enable_test=False),
51
+ ########## ModernBertModel
52
+ EmbedModelInfo("Alibaba-NLP/gte-modernbert-base",
53
+ architecture="ModernBertModel",
54
+ enable_test=True),
55
+]
56
+
57
+
58
+@pytest.mark.parametrize("model_info", MODELS)
59
+def test_models_mteb(hf_runner, vllm_runner,
60
+ model_info: EmbedModelInfo) -> None:
61
+ pytest.skip("Skipping mteb test.")
62
+
63
+ from .mteb_utils import mteb_test_embed_models
64
+
65
+ vllm_extra_kwargs: dict[str, Any] = {}
66
+ if model_info.name == "Alibaba-NLP/gte-Qwen2-1.5B-instruct":
67
+ vllm_extra_kwargs["hf_overrides"] = {"is_causal": True}
68
+
69
+ if model_info.architecture == "GteNewModel":
70
+ vllm_extra_kwargs["hf_overrides"] = {"architectures": ["GteNewModel"]}
71
+
72
+ mteb_test_embed_models(hf_runner, vllm_runner, model_info,
73
+ vllm_extra_kwargs)
74
+
75
+
76
+@pytest.mark.parametrize("model_info", MODELS)
77
+def test_models_correctness(hf_runner, vllm_runner, model_info: EmbedModelInfo,
78
+ example_prompts) -> None:
79
+ if not model_info.enable_test:
80
+ pytest.skip("Skipping test.")
81
+
82
+ # ST will strip the input texts, see test_embedding.py
83
+ example_prompts = [str(s).strip() for s in example_prompts]
84
+
85
+ vllm_extra_kwargs: dict[str, Any] = {}
86
+ if model_info.name == "Alibaba-NLP/gte-Qwen2-1.5B-instruct":
87
+ vllm_extra_kwargs["hf_overrides"] = {"is_causal": True}
88
+
89
+ if model_info.architecture == "GteNewModel":
90
+ vllm_extra_kwargs["hf_overrides"] = {"architectures": ["GteNewModel"]}
91
+
92
+ with vllm_runner(model_info.name,
93
+ task="embed",
94
+ dtype=model_info.dtype,
95
+ max_model_len=None,
96
+ **vllm_extra_kwargs) as vllm_model:
97
+ vllm_outputs = vllm_model.encode(example_prompts)
98
+
99
+ with hf_runner(
100
+ model_info.name,
101
+ dtype=model_info.dtype,
102
+ is_sentence_transformer=True,
103
+ ) as hf_model:
104
+ run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs)
tests/models/language/pooling/test_nomic.py
+4
@@ -23,6 +23,7 @@ MODELS = [
23
@pytest.mark.parametrize("model_info", MODELS)
24
def test_models_mteb(hf_runner, vllm_runner,
25
model_info: EmbedModelInfo) -> None:
26
+ pytest.skip("Skipping mteb test.")
27
from .mteb_utils import mteb_test_embed_models
28
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
29
@@ -33,6 +34,9 @@ def test_models_correctness(hf_runner, vllm_runner, model_info: EmbedModelInfo,
34
if not model_info.enable_test:
35
pytest.skip("Skipping test.")
36
37
+ # ST will strip the input texts, see test_embedding.py
38
+ example_prompts = [str(s).strip() for s in example_prompts]
39
+
40
with vllm_runner(model_info.name,
41
task="embed",
42
dtype=model_info.dtype,
tests/models/language/pooling/test_snowflake_arctic_embed.py
+4
@@ -46,6 +46,7 @@ def test_models_mteb(
46
vllm_runner,
47
model_info: EmbedModelInfo,
48
) -> None:
49
+ pytest.skip("Skipping mteb test.")
50
from .mteb_utils import mteb_test_embed_models
51
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
52
@@ -60,6 +61,9 @@ def test_models_correctness(
61
if not model_info.enable_test:
62
pytest.skip("Skipping test.")
63
64
+ # ST will strip the input texts, see test_embedding.py
65
+ example_prompts = [str(s).strip() for s in example_prompts]
66
+
67
with vllm_runner(model_info.name,
68
task="embed",
69
dtype=model_info.dtype,
tests/models/registry.py
+6
@@ -256,11 +256,17 @@ _EMBEDDING_EXAMPLE_MODELS = {
256
"GritLM": _HfExamplesInfo("parasail-ai/GritLM-7B-vllm"),
257
"GteModel": _HfExamplesInfo("Snowflake/snowflake-arctic-embed-m-v2.0",
258
trust_remote_code=True),
259
+ "GteNewModel": _HfExamplesInfo("Alibaba-NLP/gte-base-en-v1.5",
260
+ trust_remote_code=True,
261
+ hf_overrides={"architectures":
262
+ ["GteNewModel"]}),
263
"InternLM2ForRewardModel": _HfExamplesInfo("internlm/internlm2-1_8b-reward",
264
trust_remote_code=True),
265
"JambaForSequenceClassification": _HfExamplesInfo("ai21labs/Jamba-tiny-reward-dev"), # noqa: E501
266
"LlamaModel": _HfExamplesInfo("llama", is_available_online=False),
267
"MistralModel": _HfExamplesInfo("intfloat/e5-mistral-7b-instruct"),
268
+ "ModernBertModel": _HfExamplesInfo("Alibaba-NLP/gte-modernbert-base",
269
+ trust_remote_code=True),
270
"NomicBertModel": _HfExamplesInfo("Snowflake/snowflake-arctic-embed-m-long", # noqa: E501
271
trust_remote_code=True),
272
"Qwen2Model": _HfExamplesInfo("ssmits/Qwen2-7B-Instruct-embed-base"),
vllm/model_executor/layers/activation.py
+1
-1
@@ -354,7 +354,7 @@ def get_act_fn(act_fn_name: str) -> nn.Module:
354
_ACTIVATION_AND_MUL_REGISTRY = LazyDict({
355
"gelu": lambda: GeluAndMul(),
356
"silu": lambda: SiluAndMul(),
357
- "gelu_and_mul": lambda: GeluAndMul(),
357
+ "geglu": lambda: GeluAndMul(),
358
})
359
360
vllm/model_executor/layers/rotary_embedding.py
+42
@@ -456,6 +456,40 @@ class LinearScalingRotaryEmbedding(RotaryEmbedding):
456
return self._scaling_factor_to_offset
457
458
459
+class NTKScalingRotaryEmbedding(RotaryEmbedding):
460
+ """RotaryEmbedding extended with fixed and mixed NTK scaling.
461
+ https://kexue.fm/archives/9706 """
462
+
463
+ def __init__(self,
464
+ head_size: int,
465
+ rotary_dim: int,
466
+ max_position_embeddings: int,
467
+ base: int,
468
+ is_neox_style: bool,
469
+ scaling_factor: float,
470
+ dtype: torch.dtype,
471
+ mixed_b: Optional[float] = None) -> None:
472
+ self.scaling_factor = scaling_factor
473
+ self.mixed_b = mixed_b
474
+ super().__init__(head_size, rotary_dim, max_position_embeddings, base,
475
+ is_neox_style, dtype)
476
+
477
+ def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
478
+ base = self.base * (self.scaling_factor if self.mixed_b is None else 1)
479
+ inv_freq = super()._compute_inv_freq(base)
480
+
481
+ if self.mixed_b is None:
482
+ inv_freq = inv_freq / self.scaling_factor**(2 / self.rotary_dim)
483
+ else:
484
+ a = torch.tensor(self.scaling_factor).log() / (self.rotary_dim /
485
+ 2)**self.mixed_b
486
+ lambda_1_m = (a * torch.arange(
487
+ 1, self.rotary_dim // 2 + 1).float()**self.mixed_b).exp()
488
+ inv_freq = inv_freq / lambda_1_m
489
+
490
+ return inv_freq
491
+
492
+
493
class DynamicNTKScalingRotaryEmbedding(RotaryEmbedding):
494
"""RotaryEmbedding extended with Dynamic NTK scaling.
495
@@ -1765,6 +1799,14 @@ def get_rope(
1799
max_position, base,
1800
is_neox_style,
1801
scaling_factor, dtype)
1802
+ elif scaling_type == "ntk":
1803
+ scaling_factor = rope_scaling["factor"]
1804
+ mixed_b = rope_scaling.get('mixed_b', None)
1805
+ rotary_emb = NTKScalingRotaryEmbedding(head_size, rotary_dim,
1806
+ max_position, base,
1807
+ is_neox_style,
1808
+ scaling_factor, dtype,
1809
+ mixed_b)
1810
elif scaling_type == "dynamic":
1811
scaling_factor = rope_scaling["factor"]
1812
rotary_emb = DynamicNTKScalingRotaryEmbedding(
vllm/model_executor/models/bert_with_rope.py
+84
-23
@@ -32,11 +32,18 @@ class BertWithRopeEmbedding(nn.Module):
32
def __init__(self, config: PretrainedConfig):
33
34
super().__init__()
35
- assert config.type_vocab_size > 0
35
+ if config.position_embedding_type not in ["rope", "rotary"]:
36
+ raise ValueError("Only 'rotary'('rope') position_embedding_type" +
37
+ " is supported")
38
+
39
self.word_embeddings = VocabParallelEmbedding(config.vocab_size,
40
config.hidden_size)
38
- self.token_type_embeddings = VocabParallelEmbedding(
39
- config.type_vocab_size, config.hidden_size)
41
+ if config.type_vocab_size > 0:
42
+ self.token_type_embeddings = VocabParallelEmbedding(
43
+ config.type_vocab_size, config.hidden_size)
44
+ else:
45
+ self.token_type_embeddings = None
46
+
47
self.LayerNorm = nn.LayerNorm(config.hidden_size,
48
eps=config.layer_norm_eps)
49
@@ -47,13 +54,17 @@ class BertWithRopeEmbedding(nn.Module):
54
) -> torch.Tensor:
55
input_shape = input_ids.size()
56
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)
57
55
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
56
- embeddings = inputs_embeds + token_type_embeddings
58
+ embeddings = inputs_embeds
59
+ if self.token_type_embeddings is not None:
60
+ if token_type_ids is None:
61
+ token_type_ids = torch.zeros(input_shape,
62
+ dtype=torch.long,
63
+ device=inputs_embeds.device)
64
+
65
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
66
+ embeddings += token_type_embeddings
67
+
68
embeddings = self.LayerNorm(embeddings)
69
return embeddings
70
@@ -321,7 +332,7 @@ class BertWithRopeBlock(nn.Module):
332
if moe:
333
self.mlp = NomicMoELayer(config=config, )
334
else:
324
- if config.hidden_act in ["silu", "gelu_and_mul"]:
335
+ if config.hidden_act in ["silu", "geglu"]:
336
self.mlp = BertWithRopeGatedMLP(
337
hidden_size=config.hidden_size,
338
intermediate_size=config.intermediate_size,
@@ -390,6 +401,7 @@ class BertWithRope(nn.Module, SupportsV0Only, SupportsQuant):
401
402
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
403
super().__init__()
404
+ self.vllm_config = vllm_config
405
self.config = self.config_verify(vllm_config)
406
self.embeddings = BertWithRopeEmbedding(self.config)
407
self.encoder = BertWithRopeEncoder(
@@ -420,7 +432,7 @@ class BertWithRope(nn.Module, SupportsV0Only, SupportsQuant):
432
torch.Tensor]]) -> Set[str]:
433
weights = self.hf_to_vllm_mapper.apply(weights)
434
423
- if self.config.hidden_act in ["silu", "gelu_and_mul"]:
435
+ if self.config.hidden_act in ["silu", "geglu"]:
436
stacked_params_mapping = [
437
# (param_name, shard_name, shard_id)
438
("gate_up_proj", "gate_proj", 0),
@@ -458,6 +470,8 @@ class BertWithRope(nn.Module, SupportsV0Only, SupportsQuant):
470
471
472
class NomicBertModel(BertWithRope):
473
+ # for https://huggingface.co/nomic-ai/nomic-bert-2048
474
+
475
hf_to_vllm_mapper = WeightsMapper(
476
orig_to_new_substr={
477
"emb_ln": "embeddings.LayerNorm",
@@ -475,6 +489,9 @@ class NomicBertModel(BertWithRope):
489
490
assert config.__class__.__name__ == "NomicBertConfig"
491
assert config.activation_function in ["swiglu", "gelu"]
492
+ config.position_embedding_type = getattr(config,
493
+ "position_embedding_type",
494
+ "rope")
495
496
if config.activation_function == "swiglu":
497
config.hidden_act = "silu"
@@ -512,10 +529,13 @@ class NomicBertModel(BertWithRope):
529
return config
530
531
515
-class GteModel(BertWithRope):
532
+class GteNewModel(BertWithRope):
533
+ # for https://huggingface.co/Alibaba-NLP/new-impl
534
+
535
hf_to_vllm_mapper = WeightsMapper(
536
orig_to_new_substr={
518
- "layer": 'layers',
537
+ "new.": "",
538
+ "layer": "layers",
539
"attention.qkv_proj": "attn.qkv_proj",
540
"attention.o_proj": "attn.out_proj",
541
})
@@ -523,7 +543,7 @@ class GteModel(BertWithRope):
543
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
544
super().__init__(vllm_config=vllm_config, prefix=prefix)
545
526
- # GteModel only gate_up_proj does not have bias.
546
+ # GteNewModel only gate_up_proj does not have bias.
547
# Hack method learned from vllm/model_executor/models/glm.py
548
for layer in self.encoder.layers:
549
layer.mlp.gate_up_proj.bias = None
@@ -532,12 +552,10 @@ class GteModel(BertWithRope):
552
def config_verify(self, vllm_config):
553
config = vllm_config.model_config.hf_config
554
535
- assert config.__class__.__name__ == "GteConfig"
536
- assert config.position_embedding_type == "rope"
555
+ assert config.__class__.__name__ == "NewConfig"
556
assert config.hidden_act == "gelu"
557
539
- config.position_embedding_type = "rotary"
540
- config.hidden_act = "gelu_and_mul"
558
+ config.hidden_act = "geglu"
559
560
head_dim = config.hidden_size // config.num_attention_heads
561
config.rotary_kwargs = {
@@ -559,13 +577,52 @@ class GteModel(BertWithRope):
577
else:
578
yield name, weight
579
580
+ def ignore_unnecessary_layers(self,
581
+ weights: Iterable[Tuple[str, torch.Tensor]]):
582
+ for name, weight in weights:
583
+ if name.startswith("classifier"):
584
+ continue
585
+ yield name, weight
586
+
587
def load_weights(self, weights: Iterable[Tuple[str,
588
torch.Tensor]]) -> Set[str]:
589
+ weights = self.ignore_unnecessary_layers(weights)
590
weights = self.split_up_gate_proj(weights)
591
return super().load_weights(weights)
592
593
594
+class SnowflakeGteNewModel(GteNewModel):
595
+ # for Snowflake/snowflake-arctic-embed-m-v2.0
596
+
597
+ hf_to_vllm_mapper = WeightsMapper(
598
+ orig_to_new_substr={
599
+ "layer": "layers",
600
+ "attention.qkv_proj": "attn.qkv_proj",
601
+ "attention.o_proj": "attn.out_proj",
602
+ })
603
+
604
+ def config_verify(self, vllm_config):
605
+ config = vllm_config.model_config.hf_config
606
+
607
+ assert config.__class__.__name__ == "GteConfig"
608
+ assert config.hidden_act == "gelu"
609
+
610
+ config.hidden_act = "geglu"
611
+
612
+ head_dim = config.hidden_size // config.num_attention_heads
613
+ config.rotary_kwargs = {
614
+ "head_size": head_dim,
615
+ "rotary_dim": getattr(config, "rotary_emb_dim", head_dim),
616
+ "max_position": config.max_position_embeddings,
617
+ "base": config.rope_theta,
618
+ "rope_scaling": getattr(config, "rope_scaling", None)
619
+ }
620
+ return config
621
+
622
+
623
class JinaRobertaModel(BertWithRope):
624
+ # for https://huggingface.co/jinaai/jina-embeddings-v3
625
+
626
hf_to_vllm_mapper = WeightsMapper(
627
orig_to_new_substr={
628
"emb_ln": "embeddings.LayerNorm",
@@ -579,6 +636,9 @@ class JinaRobertaModel(BertWithRope):
636
637
def config_verify(self, vllm_config):
638
config = vllm_config.model_config.hf_config
639
+
640
+ assert config.__class__.__name__ == "XLMRobertaFlashConfig"
641
+
642
head_dim = config.hidden_size // config.num_attention_heads
643
config.rotary_kwargs = {
644
"head_size": head_dim,
@@ -611,6 +671,7 @@ class JinaRobertaModel(BertWithRope):
671
# This is a temporary solution until we have a better way to handle
672
673
scaling = self.config.lora_alpha / self.config.lora_rank
674
+ device = self.vllm_config.device_config.device
675
676
weights = {name: weight for name, weight in weights}
677
@@ -628,13 +689,13 @@ class JinaRobertaModel(BertWithRope):
689
weight_name = name[:-len(o)]
690
691
if "embeddings" in weight_name:
631
- B = weights[weight_name + a][i].cuda().float()
632
- A = weights[weight_name + b][i].cuda().float()
692
+ B = weights[weight_name + a][i].to(device).float()
693
+ A = weights[weight_name + b][i].to(device).float()
694
else:
634
- B = weights[weight_name + b][i].cuda().float()
635
- A = weights[weight_name + a][i].cuda().float()
695
+ B = weights[weight_name + b][i].to(device).float()
696
+ A = weights[weight_name + a][i].to(device).float()
697
637
- weight = (weights[weight_name + o].cuda() +
698
+ weight = (weights[weight_name + o].to(device) +
699
torch.matmul(B, A).view(shape) * scaling)
700
weight = weight.cpu().to(dtype)
701
vllm/model_executor/models/modernbert.py
+3
@@ -230,9 +230,12 @@ class ModernBertModel(nn.Module):
230
def forward(
231
self,
232
input_ids: Optional[torch.LongTensor] = None,
233
+ positions: Optional[torch.Tensor] = None,
234
+ intermediate_tensors: Optional[IntermediateTensors] = None,
235
inputs_embeds: Optional[torch.Tensor] = None,
236
position_ids: Optional[torch.LongTensor] = None,
237
) -> torch.Tensor:
238
+ position_ids = positions if positions is not None else position_ids
239
if inputs_embeds is not None:
240
hidden_states = inputs_embeds
241
else:
vllm/model_executor/models/registry.py
+3
-1
@@ -127,7 +127,8 @@ _EMBEDDING_MODELS = {
127
"Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
128
"GlmForCausalLM": ("glm", "GlmForCausalLM"),
129
"GritLM": ("gritlm", "GritLM"),
130
- "GteModel": ("bert_with_rope", "GteModel"),
130
+ "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
131
+ "GteNewModel": ("bert_with_rope", "GteNewModel"),
132
"InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
133
"JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"), # noqa: E501
134
"LlamaModel": ("llama", "LlamaForCausalLM"),
@@ -137,6 +138,7 @@ _EMBEDDING_MODELS = {
138
if arch == "LlamaForCausalLM"
139
},
140
"MistralModel": ("llama", "LlamaForCausalLM"),
141
+ "ModernBertModel": ("modernbert", "ModernBertModel"),
142
"NomicBertModel": ("bert_with_rope", "NomicBertModel"),
143
"Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
144
"Qwen2Model": ("qwen2", "Qwen2EmbeddingModel"),