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, 13:30 UTC
[New Model]: Snowflake Arctic Embed (Family) (#16649)
wang.yuqi committed
Apr 18, 2025 at 23:11 UTC
3d3ab3689f19a3296995002395f9f756edd27957
7 files changed
+312
-26
tests/entrypoints/openai/test_embedding_dimensions.py
+4
-11
@@ -3,24 +3,17 @@
3
Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`.
4
"""
5
6
-from typing import NamedTuple
7
-
6
import openai
7
import pytest
8
9
from vllm.entrypoints.openai.protocol import EmbeddingResponse
10
11
+from ...models.embedding.utils import EmbedModelInfo
12
from ...utils import RemoteOpenAIServer
13
15
-
16
-class ModelInfo(NamedTuple):
17
- name: str
18
- is_matryoshka: bool
19
-
20
-
14
MODELS = [
22
- ModelInfo(name="BAAI/bge-m3", is_matryoshka=False),
23
- ModelInfo(name="jinaai/jina-embeddings-v3", is_matryoshka=True),
15
+ EmbedModelInfo(name="BAAI/bge-m3", is_matryoshka=False),
16
+ EmbedModelInfo(name="jinaai/jina-embeddings-v3", is_matryoshka=True),
17
]
18
19
input_texts = [
@@ -30,7 +23,7 @@ input_texts = [
23
24
@pytest.mark.asyncio
25
@pytest.mark.parametrize("model", MODELS)
33
-async def test_validating_dimensions(model: ModelInfo):
26
+async def test_validating_dimensions(model: EmbedModelInfo):
27
args = [
28
"--task",
29
"embed",
tests/models/embedding/language/test_snowflake_arctic_embed.py
new
+101
@@ -0,0 +1,101 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""Compare the embedding outputs of HF and vLLM models.
3
+
4
+Run `pytest tests/models/embedding/language/test_snowflake_arctic_embed.py`.
5
+"""
6
+import pytest
7
+
8
+from tests.models.embedding.utils import EmbedModelInfo
9
+
10
+from ..utils import check_embeddings_close
11
+
12
+EMBEDDING_PROMPTS = [
13
+ 'what is snowflake?', 'Where can I get the best tacos?', 'The Data Cloud!',
14
+ 'Mexico City of Course!'
15
+]
16
+
17
+MODELS = [
18
+ EmbedModelInfo("Snowflake/snowflake-arctic-embed-xs",
19
+ is_matryoshka=False,
20
+ architecture="BertModel",
21
+ enable_test=True),
22
+ EmbedModelInfo("Snowflake/snowflake-arctic-embed-s",
23
+ is_matryoshka=False,
24
+ architecture="BertModel",
25
+ enable_test=False),
26
+ EmbedModelInfo("Snowflake/snowflake-arctic-embed-m",
27
+ is_matryoshka=False,
28
+ architecture="BertModel",
29
+ enable_test=False),
30
+ EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-long",
31
+ is_matryoshka=False,
32
+ architecture="NomicBertModel",
33
+ enable_test=True),
34
+ EmbedModelInfo("Snowflake/snowflake-arctic-embed-l",
35
+ is_matryoshka=False,
36
+ architecture="BertModel",
37
+ enable_test=False),
38
+ EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-v1.5",
39
+ is_matryoshka=True,
40
+ architecture="BertModel",
41
+ enable_test=True),
42
+ EmbedModelInfo("Snowflake/snowflake-arctic-embed-l-v2.0",
43
+ is_matryoshka=True,
44
+ architecture="XLMRobertaModel",
45
+ enable_test=True),
46
+ EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-v2.0",
47
+ is_matryoshka=True,
48
+ architecture="GteModel",
49
+ enable_test=True),
50
+]
51
+
52
+
53
+@pytest.mark.parametrize("model_info", MODELS)
54
+@pytest.mark.parametrize("dtype", ["half"])
55
+def test_models(
56
+ hf_runner,
57
+ vllm_runner,
58
+ example_prompts,
59
+ model_info: EmbedModelInfo,
60
+ dtype: str,
61
+ monkeypatch,
62
+) -> None:
63
+ if not model_info.enable_test:
64
+ # A model family has many models with the same architecture,
65
+ # and we don't need to test each one.
66
+ pytest.skip("Skipping test.")
67
+
68
+ example_prompts = example_prompts + EMBEDDING_PROMPTS
69
+
70
+ vllm_extra_kwargs = {
71
+ "hf_overrides": {
72
+ "is_matryoshka": model_info.is_matryoshka
73
+ }
74
+ }
75
+
76
+ with hf_runner(model_info.name, dtype=dtype,
77
+ is_sentence_transformer=True) as hf_model:
78
+ hf_outputs = hf_model.encode(example_prompts)
79
+
80
+ with vllm_runner(model_info.name,
81
+ task="embed",
82
+ dtype=dtype,
83
+ max_model_len=None,
84
+ **vllm_extra_kwargs) as vllm_model:
85
+
86
+ assert (vllm_model.model.llm_engine.model_config.is_matryoshka ==
87
+ model_info.is_matryoshka)
88
+
89
+ if model_info.architecture:
90
+ assert (model_info.architecture
91
+ in vllm_model.model.llm_engine.model_config.architectures)
92
+
93
+ vllm_outputs = vllm_model.encode(example_prompts)
94
+
95
+ check_embeddings_close(
96
+ embeddings_0_lst=hf_outputs,
97
+ embeddings_1_lst=vllm_outputs,
98
+ name_0="hf",
99
+ name_1="vllm",
100
+ tol=1e-2,
101
+ )
tests/models/embedding/utils.py
+8
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
from collections.abc import Sequence
4
+from typing import NamedTuple
5
6
import torch
7
import torch.nn.functional as F
@@ -37,3 +38,10 @@ def matryoshka_fy(tensor, dimensions):
38
tensor = tensor[..., :dimensions]
39
tensor = F.normalize(tensor, p=2, dim=1)
40
return tensor
41
+
42
+
43
+class EmbedModelInfo(NamedTuple):
44
+ name: str
45
+ is_matryoshka: bool
46
+ architecture: str = ""
47
+ enable_test: bool = True
tests/models/registry.py
+4
@@ -247,11 +247,15 @@ _EMBEDDING_EXAMPLE_MODELS = {
247
"BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
248
"Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"),
249
"GritLM": _HfExamplesInfo("parasail-ai/GritLM-7B-vllm"),
250
+ "GteModel": _HfExamplesInfo("Snowflake/snowflake-arctic-embed-m-v2.0",
251
+ trust_remote_code=True),
252
"InternLM2ForRewardModel": _HfExamplesInfo("internlm/internlm2-1_8b-reward",
253
trust_remote_code=True),
254
"JambaForSequenceClassification": _HfExamplesInfo("ai21labs/Jamba-tiny-reward-dev"), # noqa: E501
255
"LlamaModel": _HfExamplesInfo("llama", is_available_online=False),
256
"MistralModel": _HfExamplesInfo("intfloat/e5-mistral-7b-instruct"),
257
+ "NomicBertModel": _HfExamplesInfo("Snowflake/snowflake-arctic-embed-m-long", # noqa: E501
258
+ trust_remote_code=True),
259
"Qwen2Model": _HfExamplesInfo("ssmits/Qwen2-7B-Instruct-embed-base"),
260
"Qwen2ForRewardModel": _HfExamplesInfo("Qwen/Qwen2.5-Math-RM-72B"),
261
"Qwen2ForProcessRewardModel": _HfExamplesInfo("Qwen/Qwen2.5-Math-PRM-7B"),
vllm/model_executor/layers/activation.py
+1
@@ -354,6 +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(),
358
})
359
360
vllm/model_executor/models/bert.py
+189
-12
@@ -11,8 +11,10 @@ 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_fn
14
+from vllm.model_executor.layers.activation import (get_act_and_mul_fn,
15
+ get_act_fn)
16
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
17
+ MergedColumnParallelLinear,
18
QKVParallelLinear,
19
RowParallelLinear)
20
from vllm.model_executor.layers.pooler import (CrossEncodingPooler, Pooler,
@@ -108,6 +110,7 @@ class BertEncoder(nn.Module):
110
111
def __init__(self,
112
vllm_config: VllmConfig,
113
+ bias: bool = True,
114
rotary_kwargs: Optional[dict] = None,
115
prefix: str = ""):
116
super().__init__()
@@ -118,6 +121,7 @@ class BertEncoder(nn.Module):
121
BertLayer(config=config,
122
cache_config=cache_config,
123
quant_config=quant_config,
124
+ bias=bias,
125
rotary_kwargs=rotary_kwargs,
126
prefix=f"{prefix}.layer.{layer_idx}")
127
for layer_idx in range(config.num_hidden_layers)
@@ -139,6 +143,7 @@ class BertLayer(nn.Module):
143
config: BertConfig,
144
cache_config: Optional[CacheConfig] = None,
145
quant_config: Optional[QuantizationConfig] = None,
146
+ bias: bool = True,
147
rotary_kwargs: Optional[dict] = None,
148
prefix: str = ""):
149
super().__init__()
@@ -149,19 +154,31 @@ class BertLayer(nn.Module):
154
layer_norm_eps=config.layer_norm_eps,
155
cache_config=cache_config,
156
quant_config=quant_config,
157
+ bias=bias,
158
rotary_kwargs=rotary_kwargs,
159
prefix=f"{prefix}.attention")
160
155
- self.intermediate = BertIntermediate(
156
- hidden_size=config.hidden_size,
157
- intermediate_size=config.intermediate_size,
158
- hidden_act=config.hidden_act,
159
- quant_config=quant_config,
160
- prefix=f"{prefix}.intermediate")
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")
177
178
self.output = BertOutput(hidden_size=config.hidden_size,
179
intermediate_size=config.intermediate_size,
180
layer_norm_eps=config.layer_norm_eps,
181
+ bias=bias,
182
quant_config=quant_config,
183
prefix=f"{prefix}.output")
184
@@ -181,6 +198,7 @@ class BertAttention(nn.Module):
198
layer_norm_eps: float,
199
cache_config: Optional[CacheConfig] = None,
200
quant_config: Optional[QuantizationConfig] = None,
201
+ bias: bool = True,
202
rotary_kwargs: Optional[dict] = None,
203
prefix: str = "",
204
):
@@ -190,11 +208,13 @@ class BertAttention(nn.Module):
208
num_attention_heads=num_attention_heads,
209
cache_config=cache_config,
210
quant_config=quant_config,
211
+ bias=bias,
212
rotary_kwargs=rotary_kwargs,
213
prefix=f"{prefix}.output")
214
215
self.output = BertSelfOutput(hidden_size=hidden_size,
216
layer_norm_eps=layer_norm_eps,
217
+ bias=bias,
218
quant_config=quant_config,
219
prefix=f"{prefix}.output")
220
@@ -215,6 +235,7 @@ class BertSelfAttention(nn.Module):
235
num_attention_heads: int,
236
cache_config: Optional[CacheConfig] = None,
237
quant_config: Optional[QuantizationConfig] = None,
238
+ bias: bool = True,
239
rotary_kwargs: Optional[dict] = None,
240
prefix: str = "",
241
):
@@ -240,7 +261,7 @@ class BertSelfAttention(nn.Module):
261
head_size=self.head_dim,
262
total_num_heads=self.total_num_heads,
263
total_num_kv_heads=self.total_num_kv_heads,
243
- bias=True,
264
+ bias=bias,
265
quant_config=quant_config,
266
prefix=f"{prefix}.qkv_proj")
267
@@ -278,12 +299,13 @@ class BertSelfOutput(nn.Module):
299
def __init__(self,
300
hidden_size: int,
301
layer_norm_eps: float,
302
+ bias: bool = True,
303
quant_config: Optional[QuantizationConfig] = None,
304
prefix: str = ""):
305
super().__init__()
306
self.dense = RowParallelLinear(input_size=hidden_size,
307
output_size=hidden_size,
286
- bias=True,
308
+ bias=bias,
309
quant_config=quant_config,
310
prefix=f"{prefix}.dense")
311
self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
@@ -301,12 +323,13 @@ class BertIntermediate(nn.Module):
323
hidden_size: int,
324
intermediate_size: int,
325
hidden_act: str,
326
+ bias: bool = True,
327
quant_config: Optional[QuantizationConfig] = None,
328
prefix: str = ""):
329
super().__init__()
330
self.dense = ColumnParallelLinear(input_size=hidden_size,
331
output_size=intermediate_size,
309
- bias=True,
332
+ bias=bias,
333
quant_config=quant_config,
334
prefix=f"{prefix}.dense")
335
self.intermediate_act_fn = get_act_fn(hidden_act)
@@ -317,19 +340,46 @@ class BertIntermediate(nn.Module):
340
return hidden_states
341
342
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
+
369
class BertOutput(nn.Module):
370
371
def __init__(self,
372
hidden_size: int,
373
intermediate_size: int,
374
layer_norm_eps: float,
375
+ bias: bool = True,
376
quant_config: Optional[QuantizationConfig] = None,
377
prefix: str = ""):
378
super().__init__()
379
380
self.dense = RowParallelLinear(input_size=intermediate_size,
381
output_size=hidden_size,
332
- bias=True,
382
+ bias=bias,
383
quant_config=quant_config,
384
prefix=f"{prefix}.dense")
385
@@ -343,19 +393,32 @@ class BertOutput(nn.Module):
393
394
395
class BertModel(nn.Module, SupportsQuant):
346
- packed_modules_mapping = {"qkv_proj": ["query", "key", "value"]}
396
+ packed_modules_mapping = {
397
+ "qkv_proj": ["query", "key", "value"],
398
+ "gate_up_proj": [
399
+ "gate_proj",
400
+ "up_proj",
401
+ ],
402
+ }
403
404
def __init__(self,
405
*,
406
vllm_config: VllmConfig,
407
prefix: str = "",
408
embedding_class: type = BertEmbedding,
409
+ bias: bool = True,
410
rotary_kwargs: Optional[dict] = None,
411
add_pooling_layer: bool = False):
412
super().__init__()
413
+ """
414
+ For BertModel, all linear layers have bias.
415
+ For NomicBertModel, all linear layers do not have bias.
416
+ """
417
+
418
config = vllm_config.model_config.hf_config
419
self.embeddings = embedding_class(config)
420
self.encoder = BertEncoder(vllm_config=vllm_config,
421
+ bias=bias,
422
rotary_kwargs=rotary_kwargs,
423
prefix=f"{prefix}.encoder")
424
self.pooler = BertPooler(config) if add_pooling_layer else None
@@ -387,6 +450,8 @@ class BertModel(nn.Module, SupportsQuant):
450
("qkv_proj", "query", "q"),
451
("qkv_proj", "key", "k"),
452
("qkv_proj", "value", "v"),
453
+ ("gate_up_proj", "gate_proj", 0),
454
+ ("gate_up_proj", "up_proj", 1),
455
]
456
457
params_dict = dict(self.named_parameters())
@@ -546,3 +611,115 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding,
611
inputs_embeds=inputs_embeds,
612
intermediate_tensors=intermediate_tensors,
613
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/registry.py
+5
-3
@@ -122,13 +122,11 @@ _TEXT_GENERATION_MODELS = {
122
_EMBEDDING_MODELS = {
123
# [Text-only]
124
"BertModel": ("bert", "BertEmbeddingModel"),
125
- "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
126
- "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
127
- "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
125
"DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
126
"Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
127
"GlmForCausalLM": ("glm", "GlmForCausalLM"),
128
"GritLM": ("gritlm", "GritLM"),
129
+ "GteModel": ("bert", "GteEmbeddingModel"),
130
"InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
131
"JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"), # noqa: E501
132
"LlamaModel": ("llama", "LlamaForCausalLM"),
@@ -138,12 +136,16 @@ _EMBEDDING_MODELS = {
136
if arch == "LlamaForCausalLM"
137
},
138
"MistralModel": ("llama", "LlamaForCausalLM"),
139
+ "NomicBertModel": ("bert", "NomicBertEmbeddingModel"),
140
"Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
141
"Qwen2Model": ("qwen2", "Qwen2EmbeddingModel"),
142
"Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
143
"Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
144
"Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
145
+ "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
146
+ "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
147
"TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
148
+ "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
149
# [Multimodal]
150
"LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"), # noqa: E501
151
"Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),