modernbert.py 16.4 KB
Newer Older
xsank's avatar
xsank committed
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable, Set
xsank's avatar
xsank committed
4
5
6
7

import torch
from torch import nn
from transformers import ModernBertConfig
8
from transformers.activations import ACT2FN
xsank's avatar
xsank committed
9

10
from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
11
from vllm.compilation.decorators import support_torch_compile
xsank's avatar
xsank committed
12
13
from vllm.config import VllmConfig
from vllm.distributed import get_tensor_model_parallel_world_size
14
15
16
17
18
19
20
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
from vllm.model_executor.layers.pooler import (
    ClassifierPooler,
    DispatchPooler,
    Pooler,
    PoolingMethod,
    PoolingParamsUpdate,
21
22
    TokenPoolerHeadOutput,
    TokenPoolingMethodOutput,
23
)
24
from vllm.model_executor.layers.rotary_embedding import get_rope
25
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
xsank's avatar
xsank committed
26
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
27
from vllm.sequence import IntermediateTensors
28
from vllm.tasks import PoolingTask
29
from vllm.v1.outputs import TokenPoolerOutput
30
from vllm.v1.pool.metadata import PoolingMetadata
xsank's avatar
xsank committed
31

32
from .interfaces import SupportsCrossEncoding
33
from .interfaces_base import attn_type, default_pooling_type
34
from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
xsank's avatar
xsank committed
35
36
37
38
39
40


class ModernBertEmbeddings(nn.Module):
    def __init__(self, config: ModernBertConfig):
        super().__init__()
        self.config = config
41
42
43
        self.tok_embeddings = VocabParallelEmbedding(
            config.vocab_size, config.hidden_size
        )
44
45
46
47
        eps = (
            getattr(config, "norm_eps", None)
            or getattr(config, "layer_norm_eps", None)
            or 1e-5
48
        )
49
        self.norm = nn.LayerNorm(config.hidden_size, eps=eps, bias=config.norm_bias)
xsank's avatar
xsank committed
50

51
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
52
53
        return self.tok_embeddings(input_ids)

xsank's avatar
xsank committed
54
55
56
    def forward(
        self,
        input_ids: torch.Tensor,
57
        inputs_embeds: torch.Tensor | None = None,
xsank's avatar
xsank committed
58
    ) -> torch.Tensor:
59
        if inputs_embeds is not None:
xsank's avatar
xsank committed
60
61
62
63
64
65
66
67
            return self.norm(inputs_embeds)
        else:
            inputs_embeds = self.tok_embeddings(input_ids)
            embeddings = self.norm(inputs_embeds)
            return embeddings


class ModernBertAttention(nn.Module):
68
69
70
    def __init__(
        self, config: ModernBertConfig, layer_id: int | None = None, prefix: str = ""
    ):
xsank's avatar
xsank committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.layer_id = layer_id
        self.deterministic_flash_attn = config.deterministic_flash_attn
        self.num_heads = config.num_attention_heads
        assert self.num_heads % tp_size == 0
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.all_head_size = self.head_dim * self.num_heads
        self.scaling = self.head_dim**-0.5
        self.Wqkv = QKVParallelLinear(
            config.hidden_size,
            self.head_dim,
            self.num_heads,
            bias=config.attention_bias,
87
            prefix=f"{prefix}.Wqkv",
xsank's avatar
xsank committed
88
89
        )

90
91
92
93
94
95
96
        if layer_types := getattr(config, "layer_types", None):
            # Transformers v5
            layer_type = layer_types[layer_id]
            rope_parameters = config.rope_parameters[layer_type]
            sliding_window: int | None = None
            if layer_type == "sliding_attention":
                sliding_window = config.local_attention // 2
xsank's avatar
xsank committed
97
        else:
98
99
100
101
102
103
104
105
106
107
108
109
            # Transformers v4
            sliding_window = None
            if layer_id % config.global_attn_every_n_layers != 0:
                sliding_window = config.local_attention // 2
                rope_theta = (
                    config.local_rope_theta
                    if config.local_rope_theta is not None
                    else config.global_rope_theta
                )
            else:
                rope_theta = config.global_rope_theta
            rope_parameters = {"rope_type": "default", "rope_theta": rope_theta}
xsank's avatar
xsank committed
110

111
112
113
114
115
        self.rotary_emb = get_rope(
            head_size=self.head_dim,
            max_position=config.max_position_embeddings,
            rope_parameters=rope_parameters,
            dtype=torch.float16,
116
        )
117
118
119
120
121
        self.attn = EncoderOnlyAttention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            prefix=f"{layer_id}.attn",
122
123
124
            per_layer_sliding_window=sliding_window,
        )
        self.Wo = RowParallelLinear(
125
126
127
128
            config.hidden_size,
            config.hidden_size,
            bias=config.attention_bias,
            prefix=f"{prefix}.Wo",
129
        )
xsank's avatar
xsank committed
130
131
132
133

    def forward(
        self,
        hidden_states: torch.Tensor,
134
        position_ids: torch.Tensor,
xsank's avatar
xsank committed
135
136
137
138
139
140
141
142
143
144
145
    ) -> torch.Tensor:
        qkv, _ = self.Wqkv(hidden_states)
        q, k, v = qkv.split([self.all_head_size] * 3, dim=-1)
        q, k = self.rotary_emb(position_ids, q, k)
        attn_outputs = self.attn(q, k, v)
        hidden_states = attn_outputs
        hidden_states, _ = self.Wo(hidden_states)
        return hidden_states


class ModernBertMLP(nn.Module):
146
    def __init__(self, config: ModernBertConfig, prefix: str = ""):
xsank's avatar
xsank committed
147
148
        super().__init__()
        self.config = config
149
150
151
        self.Wi = nn.Linear(
            config.hidden_size, int(config.intermediate_size) * 2, bias=config.mlp_bias
        )
xsank's avatar
xsank committed
152
        self.act = nn.GELU()
153
        self.Wo = RowParallelLinear(
154
155
156
157
            config.intermediate_size,
            config.hidden_size,
            bias=config.mlp_bias,
            prefix=f"{prefix}.Wo",
158
        )
xsank's avatar
xsank committed
159
160
161
162
163
164
165

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input, gate = self.Wi(hidden_states).chunk(2, dim=-1)
        return self.Wo(self.act(input) * gate)[0]


class ModernBertLayer(nn.Module):
166
    def __init__(
167
        self, config: ModernBertConfig, prefix: str = "", layer_id: int | None = None
168
    ):
xsank's avatar
xsank committed
169
170
171
172
173
        super().__init__()
        self.config = config
        if layer_id == 0:
            self.attn_norm = nn.Identity()
        else:
174
175
176
            self.attn_norm = nn.LayerNorm(
                config.hidden_size, eps=config.norm_eps, bias=config.norm_bias
            )
177
178
179
        self.attn = ModernBertAttention(
            config=config, layer_id=layer_id, prefix=f"{prefix}.attn"
        )
180
181
182
        self.mlp_norm = nn.LayerNorm(
            config.hidden_size, eps=config.norm_eps, bias=config.norm_bias
        )
183
        self.mlp = ModernBertMLP(config, prefix=f"{prefix}.mlp")
xsank's avatar
xsank committed
184
185
186
187

    def forward(
        self,
        hidden_states: torch.Tensor,
188
189
        position_ids: torch.Tensor,
    ) -> torch.Tensor:
190
191
192
        attn_outputs = self.attn(
            hidden_states=self.attn_norm(hidden_states), position_ids=position_ids
        )
xsank's avatar
xsank committed
193
194
195
196
197
198
199
200
201
202
        hidden_states = hidden_states + attn_outputs
        mlp_output = self.mlp(self.mlp_norm(hidden_states))
        hidden_states = hidden_states + mlp_output
        return hidden_states


class ModernBertEncoderLayer(nn.Module):
    def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
203
204
        self.layers = nn.ModuleList(
            [
205
206
207
208
209
                ModernBertLayer(
                    config=config,
                    layer_id=layer_id,
                    prefix=f"{prefix}.layers.{layer_id}",
                )
210
211
212
                for layer_id in range(config.num_hidden_layers)
            ]
        )
xsank's avatar
xsank committed
213
214
215
216

    def forward(
        self,
        hidden_states: torch.Tensor,
217
        position_ids: torch.Tensor,
xsank's avatar
xsank committed
218
219
220
221
222
223
    ) -> torch.Tensor:
        for i, layer in enumerate(self.layers):
            hidden_states = layer(hidden_states, position_ids)
        return hidden_states


224
@support_torch_compile
225
@default_pooling_type("CLS")
xsank's avatar
xsank committed
226
227
class ModernBertModel(nn.Module):
    hf_to_vllm_mapper = WeightsMapper(
228
229
        orig_to_new_prefix={"layers.": "encoder_layer.layers."}
    )
xsank's avatar
xsank committed
230
231
232
233
234
235
236
237
238
239

    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
    ):
        super().__init__()
        config = vllm_config.model_config.hf_config
        self.config = config
        self.embeddings = ModernBertEmbeddings(config)
240
241
242
        self.encoder_layer = ModernBertEncoderLayer(
            vllm_config, prefix=f"{prefix}.encoder_layer"
        )
243
244
245
        self.final_norm = nn.LayerNorm(
            config.hidden_size, eps=config.norm_eps, bias=config.norm_bias
        )
xsank's avatar
xsank committed
246

247
248
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embeddings.embed_input_ids(input_ids)
249

250
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
xsank's avatar
xsank committed
251
252
        weights = self.hf_to_vllm_mapper.apply(weights)
        params_dict = dict(self.named_parameters())
253
        loaded_params: set[str] = set()
xsank's avatar
xsank committed
254
255
256
257
        for name, loaded_weight in weights:
            if name.endswith(".bias") and name not in params_dict:
                continue
            param = params_dict[name]
258
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
xsank's avatar
xsank committed
259
260
261
262
263
264
            weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

    def forward(
        self,
265
266
        input_ids: torch.Tensor,
        positions: torch.Tensor,
267
268
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
xsank's avatar
xsank committed
269
270
271
272
    ) -> torch.Tensor:
        if inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
273
274
275
            hidden_states = self.embeddings(
                input_ids=input_ids, inputs_embeds=inputs_embeds
            )
xsank's avatar
xsank committed
276
277
278

        outputs = self.encoder_layer(
            hidden_states=hidden_states,
279
            position_ids=positions,
xsank's avatar
xsank committed
280
281
282
283
284
        )
        norm_outputs = self.final_norm(outputs)
        return norm_outputs


285
class ModernBertPooler(Pooler):
xsank's avatar
xsank committed
286
287
    def __init__(self, config: ModernBertConfig):
        super().__init__()
288

289
        pooling_type = config.classifier_pooling.upper()
290
        self.pooling = PoolingMethod.from_pooling_type(pooling_type)
291
292
293
        self.dense = nn.Linear(
            config.hidden_size, config.hidden_size, config.classifier_bias
        )
xsank's avatar
xsank committed
294
        self.act = nn.GELU()
295
296
297
        self.norm = nn.LayerNorm(
            config.hidden_size, eps=config.norm_eps, bias=config.norm_bias
        )
xsank's avatar
xsank committed
298

299
300
301
302
    def get_supported_tasks(self) -> Set[PoolingTask]:
        return self.pooling.get_supported_tasks()

    def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
303
        return self.pooling.get_pooling_updates(task)
304

305
    def head(
306
        self,
307
        pooled_data: TokenPoolingMethodOutput,
308
        pooling_metadata: PoolingMetadata,
309
310
311
    ) -> TokenPoolerHeadOutput:
        if isinstance(pooled_data, list):
            pooled_data = torch.stack(pooled_data)
312

313
314
        pooled_data = pooled_data.to(self.dense.weight.dtype)
        return self.norm(self.act(self.dense(pooled_data)))
315

316
317
318
319
320
321
322
323
    def forward(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> TokenPoolerOutput:
        pooled_data = self.pooling(hidden_states, pooling_metadata)
        pooled_data = self.head(pooled_data, pooling_metadata)
        return pooled_data
xsank's avatar
xsank committed
324
325


326
@default_pooling_type("CLS")
327
class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):
328
329
    is_pooling_model = True

xsank's avatar
xsank committed
330
331
332
333
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        self.config = config
334
335
336
337
338
339
340
341
        self.model = ModernBertModel(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "modernbert")
        )
        self.classifier = nn.Linear(
            config.hidden_size,
            config.num_labels,
            dtype=vllm_config.model_config.head_dtype,
        )
342
        self.pooling = ModernBertPooler(config)
343
344
345
346

        pooler_config = vllm_config.model_config.pooler_config
        assert pooler_config is not None

347
348
        self.pooler = DispatchPooler(
            {
349
350
351
                "token_classify": Pooler.for_token_classify(
                    pooler_config, classifier=self.classifier
                ),
352
                "classify": ClassifierPooler(
353
                    pooling=self.pooling, classifier=self.classifier, act_fn="classify"
354
355
                ),
                "score": ClassifierPooler(
356
                    pooling=self.pooling, classifier=self.classifier, act_fn="score"
357
358
359
                ),
            }
        )
xsank's avatar
xsank committed
360

361
362
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
363

364
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
xsank's avatar
xsank committed
365
366
367
368
369
        self_weights = []

        def weight_filter():
            for name, weight in weights:
                if name.startswith("model."):
370
                    yield name[len("model.") :], weight
xsank's avatar
xsank committed
371
372
373
374
375
376
377
378
379
380
                else:
                    self_weights.append((name, weight))

        self.model.load_weights(weight_filter())

        params_dict = dict(self.named_parameters())

        for name, loaded_weight in self_weights:
            if name.startswith("classifier"):
                param = params_dict[name]
381
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
xsank's avatar
xsank committed
382
383
                weight_loader(param, loaded_weight)
            if name.startswith("head"):
384
385
                param = params_dict["pooling." + name[len("head") + 1 :]]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
xsank's avatar
xsank committed
386
387
388
389
                weight_loader(param, loaded_weight)

    def forward(
        self,
390
        input_ids: torch.LongTensor | None,
xsank's avatar
xsank committed
391
        positions: torch.Tensor,
392
393
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
xsank's avatar
xsank committed
394
395
396
397
    ) -> torch.Tensor:
        return self.model(
            input_ids=input_ids,
            inputs_embeds=inputs_embeds,
398
            positions=positions,
xsank's avatar
xsank committed
399
        )
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419


class ModernBertPredictionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.dense = nn.Linear(
            config.hidden_size, config.hidden_size, bias=config.classifier_bias
        )
        self.act = ACT2FN[config.classifier_activation]
        self.norm = nn.LayerNorm(
            config.hidden_size,
            eps=getattr(config, "norm_eps", 1e-5),
            bias=getattr(config, "norm_bias", True),
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.norm(self.act(self.dense(hidden_states)))


420
@attn_type("encoder_only")
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
@default_pooling_type("ALL")
class ModernBertForTokenClassification(nn.Module):
    is_pooling_model = True

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        self.head_dtype = vllm_config.model_config.head_dtype
        self.num_labels = config.num_labels
        self.model = ModernBertModel(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "modernbert")
        )
        self.head = ModernBertPredictionHead(config)
        self.classifier = nn.Linear(
            config.hidden_size, config.num_labels, dtype=self.head_dtype
        )

        pooler_config = vllm_config.model_config.pooler_config
        assert pooler_config is not None

        self.pooler = DispatchPooler(
            {
443
444
445
                "token_classify": Pooler.for_token_classify(
                    pooler_config=pooler_config
                ),
446
447
448
            }
        )

449
450
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
451
452
453
454
455
456
457
458

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
        loader = AutoWeightsLoader(self, skip_prefixes=["drop"])
        loaded_params = loader.load_weights(weights)
        return loaded_params

    def forward(
        self,
459
        input_ids: torch.Tensor | None,
460
        positions: torch.Tensor,
461
462
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
463
464
465
466
467
468
469
470
471
472
    ) -> torch.Tensor:
        hidden_states = self.model(
            input_ids=input_ids,
            positions=positions,
            inputs_embeds=inputs_embeds,
            intermediate_tensors=intermediate_tensors,
        )
        hidden_states = self.head(hidden_states)
        hidden_states = hidden_states.to(self.head_dtype)
        return self.classifier(hidden_states)