jamba.py 21.8 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Inference-only Jamba model."""
4

5
from collections.abc import Iterable
6
from itertools import islice
Mor Zusman's avatar
Mor Zusman committed
7
8
9
10
11
12

import torch
from torch import nn
from transformers import JambaConfig

from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import CacheConfig, ModelConfig, VllmConfig
15
from vllm.distributed import get_tensor_model_parallel_world_size
16
from vllm.distributed.parallel_state import get_pp_group
17
from vllm.model_executor.layers.fused_moe import FusedMoE
Mor Zusman's avatar
Mor Zusman committed
18
from vllm.model_executor.layers.layernorm import RMSNorm
19
20
21
22
23
from vllm.model_executor.layers.linear import (
    QKVParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
Mor Zusman's avatar
Mor Zusman committed
24
from vllm.model_executor.layers.logits_processor import LogitsProcessor
25
from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
26
from vllm.model_executor.layers.mamba.mamba_utils import (
27
28
29
    MambaStateDtypeCalculator,
    MambaStateShapeCalculator,
)
30
from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
31
from vllm.model_executor.layers.quantization import QuantizationConfig
Mor Zusman's avatar
Mor Zusman committed
32
from vllm.model_executor.layers.vocab_parallel_embedding import (
33
34
35
36
    DEFAULT_VOCAB_PADDING_SIZE,
    ParallelLMHead,
    VocabParallelEmbedding,
)
37
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38
from vllm.model_executor.models.llama import LlamaMLP as JambaMLP
39
from vllm.sequence import IntermediateTensors
Mor Zusman's avatar
Mor Zusman committed
40

41
from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
42
43
44
45
46
47
48
49
from .utils import (
    AutoWeightsLoader,
    WeightsMapper,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
50

Mor Zusman's avatar
Mor Zusman committed
51
52

class JambaMoE(nn.Module):
53
54
55
    def __init__(
        self,
        config: JambaConfig,
56
57
58
59
60
        num_experts: int | None = None,
        top_k: int | None = None,
        params_dtype: torch.dtype | None = None,
        tp_size: int | None = None,
        quant_config: QuantizationConfig | None = None,
61
62
        prefix: str = "",
    ):
Mor Zusman's avatar
Mor Zusman committed
63
        super().__init__()
64
65
        self.num_total_experts = num_experts or config.num_experts
        self.top_k = top_k or config.num_experts_per_tok
Mor Zusman's avatar
Mor Zusman committed
66
        self.hidden_size = config.hidden_size
67
        self.intermediate_size = config.intermediate_size
Mor Zusman's avatar
Mor Zusman committed
68

69
        if self.num_total_experts > 1:
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
            self.router = ReplicatedLinear(
                self.hidden_size,
                self.num_total_experts,
                bias=False,
                quant_config=None,
                params_dtype=params_dtype,
            )

        self.experts = FusedMoE(
            self.num_total_experts,
            self.top_k,
            self.hidden_size,
            self.intermediate_size,
            tp_size=tp_size,
            params_dtype=params_dtype,
            reduce_results=True,
            renormalize=False,
            use_grouped_topk=False,
            quant_config=quant_config,
            prefix=f"{prefix}.experts",
        )
Mor Zusman's avatar
Mor Zusman committed
91
92

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
93
        orig_shape = hidden_states.shape
Mor Zusman's avatar
Mor Zusman committed
94
95
        hidden_states = hidden_states.view(-1, self.hidden_size)
        # router_logits: (batch * sequence_length, n_experts)
96
97
98
        if self.num_total_experts > 1:
            router_logits, _ = self.router(hidden_states)
        else:
99
100
101
102
103
            router_logits = torch.ones(
                (hidden_states.shape[0], 1),
                device=hidden_states.device,
                dtype=hidden_states.dtype,
            )
104
105
        hidden_states = self.experts(hidden_states, router_logits)
        return hidden_states.view(orig_shape)
Mor Zusman's avatar
Mor Zusman committed
106
107
108


class JambaMambaDecoderLayer(nn.Module):
109
110
111
112
    def __init__(
        self,
        config: JambaConfig,
        layer_idx: int,
113
114
115
116
        model_config: ModelConfig | None = None,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        is_lora_enabled: bool | None = False,
117
118
119
        prefix: str = "",
        **kwargs,
    ) -> None:
Mor Zusman's avatar
Mor Zusman committed
120
121
        super().__init__()
        self.config = config
122
        self.is_lora_enabled = is_lora_enabled
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
        self.mamba = MambaMixer(
            hidden_size=config.hidden_size,
            ssm_state_size=config.mamba_d_state,
            conv_kernel_size=config.mamba_d_conv,
            intermediate_size=config.mamba_expand * config.hidden_size,
            time_step_rank=config.mamba_dt_rank,
            use_conv_bias=config.mamba_conv_bias,
            use_bias=config.mamba_proj_bias,
            use_rms_norm=True,
            rms_norm_eps=config.rms_norm_eps,
            activation=config.hidden_act,
            is_lora_enabled=self.is_lora_enabled,
            model_config=model_config,
            cache_config=cache_config,
            prefix=f"{prefix}.mixer",
        )
Mor Zusman's avatar
Mor Zusman committed
139
140

        num_experts = config.layers_num_experts[layer_idx]
141
142
143
144
145
146
147
148
149
150
151
152
153
154
        if num_experts > 1:
            self.feed_forward = JambaMoE(
                config,
                quant_config=quant_config,
                prefix=f"{prefix}.feed_forward",
            )
        else:
            self.feed_forward = JambaMLP(
                config.hidden_size,
                config.intermediate_size,
                config.hidden_act,
                quant_config=quant_config,
                prefix=f"{prefix}.feed_forward",
            )
155
156
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Mor Zusman's avatar
Mor Zusman committed
157
158
159
160

    def forward(
        self,
        hidden_states: torch.Tensor,
161
        residual: torch.Tensor | None,
Mor Zusman's avatar
Mor Zusman committed
162
163
164
165
166
167
        **kwargs,
    ):
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
168
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
Mor Zusman's avatar
Mor Zusman committed
169

170
        output = torch.empty_like(hidden_states)
171
        self.mamba(hidden_states, output)
Mor Zusman's avatar
Mor Zusman committed
172
        # Fully Connected
173
        hidden_states, residual = self.pre_ff_layernorm(output, residual)
Mor Zusman's avatar
Mor Zusman committed
174
175
176
177
178
        hidden_states = self.feed_forward(hidden_states)
        return hidden_states, residual


class JambaAttentionDecoderLayer(nn.Module):
179
180
181
182
    def __init__(
        self,
        config: JambaConfig,
        layer_idx: int,
183
184
185
        model_config: ModelConfig | None = None,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
186
187
188
        prefix: str = "",
        **kwargs,
    ) -> None:
Mor Zusman's avatar
Mor Zusman committed
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
        super().__init__()
        self.hidden_size = config.hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = config.num_attention_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = config.num_key_value_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
        self.head_dim = config.hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5

        self.qkv_proj = QKVParallelLinear(
            config.hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
            quant_config=quant_config,
        )
218
219
220
221
222
223
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
        )
Mor Zusman's avatar
Mor Zusman committed
224
225
226
227
228
229
230

        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
231
            prefix=f"{prefix}.attn",
Mor Zusman's avatar
Mor Zusman committed
232
233
234
        )

        num_experts = config.layers_num_experts[layer_idx]
235
236
237
238
239
240
241
242
243
244
245
246
247
248
        if num_experts > 1:
            self.feed_forward = JambaMoE(
                config,
                quant_config=quant_config,
                prefix=f"{prefix}.feed_forward",
            )
        else:
            self.feed_forward = JambaMLP(
                config.hidden_size,
                config.intermediate_size,
                config.hidden_act,
                quant_config=quant_config,
                prefix=f"{prefix}.feed_forward",
            )
249
250
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Mor Zusman's avatar
Mor Zusman committed
251
252
253
254
255
256
257
258
259

    def self_attention(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
260
        attn_output = self.attn(q, k, v)
Mor Zusman's avatar
Mor Zusman committed
261
262
263
264
265
266
267
        output, _ = self.o_proj(attn_output)
        return output

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
268
        residual: torch.Tensor | None,
Mor Zusman's avatar
Mor Zusman committed
269
270
271
272
273
274
        **kwargs,
    ):
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
275
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
Mor Zusman's avatar
Mor Zusman committed
276
277
278
279
280
281

        hidden_states = self.self_attention(
            positions=positions,
            hidden_states=hidden_states,
        )
        # Fully Connected
282
        hidden_states, residual = self.pre_ff_layernorm(hidden_states, residual)
Mor Zusman's avatar
Mor Zusman committed
283
284
285
286
287
288
        hidden_states = self.feed_forward(hidden_states)
        return hidden_states, residual


ALL_DECODER_LAYER_TYPES = {
    "attention": JambaAttentionDecoderLayer,
289
    "mamba": JambaMambaDecoderLayer,
Mor Zusman's avatar
Mor Zusman committed
290
291
292
}


293
@support_torch_compile
Mor Zusman's avatar
Mor Zusman committed
294
class JambaModel(nn.Module):
295
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Mor Zusman's avatar
Mor Zusman committed
296
        super().__init__()
297
298

        config = vllm_config.model_config.hf_config
299
        model_config = vllm_config.model_config
300
301
302
303
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config

Mor Zusman's avatar
Mor Zusman committed
304
        self.config = config
305
306
307
308
309
        lora_vocab = (
            (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
            if lora_config
            else 0
        )
Mor Zusman's avatar
Mor Zusman committed
310
311
312
313
314
315
316
317
318
        self.vocab_size = config.vocab_size + lora_vocab
        self.org_vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
            org_num_embeddings=config.vocab_size,
        )

319
320
        extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)}

321
322
        def get_layer(prefix: str):
            layer_idx = int(prefix.rsplit(".", 1)[1])
323
324
325
326
327
328
329
330
331
332
            layer_class = ALL_DECODER_LAYER_TYPES[config.layers_block_type[layer_idx]]
            return layer_class(
                config,
                layer_idx,
                model_config,
                cache_config,
                quant_config=quant_config,
                prefix=prefix,
                **extra_kwargs,
            )
333
334

        self.start_layer, self.end_layer, self.layers = make_layers(
335
336
337
338
339
            config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
        )
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
340

341
        self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Mor Zusman's avatar
Mor Zusman committed
342

343
344
345
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Mor Zusman's avatar
Mor Zusman committed
346
347
348
349
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
350
351
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
Mor Zusman's avatar
Mor Zusman committed
352
    ) -> torch.Tensor:
353
354
355
356
357
358
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
            residual = None
359
        else:
360
361
362
363
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

364
        for layer in islice(self.layers, self.start_layer, self.end_layer):
365
366
367
            hidden_states, residual = layer(
                positions=positions, hidden_states=hidden_states, residual=residual
            )
368

369
        if not get_pp_group().is_last_rank:
370
371
372
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
Mor Zusman's avatar
Mor Zusman committed
373
374
375
        hidden_states, _ = self.final_layernorm(hidden_states, residual)
        return hidden_states

376
377
378
379
380
381
382
    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
        return FusedMoE.make_expert_params_mapping(
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
383
384
            num_experts=self.config.num_experts,
        )
385

386
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            (".gate_up_proj", ".gate_proj", 0),
            (".gate_up_proj", ".up_proj", 1),
        ]

        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()
        expert_params_mapping = self.get_expert_mapping()
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name:
                    continue
405
                if "experts" in name:
406
407
408
409
410
411
412
413
414
415
416
417
418
419
                    continue
                name = name.replace(weight_name, param_name)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                # Skip layers on other devices.
                if is_pp_missing_parameter(name, self):
                    continue
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                for (
420
421
422
423
                    param_name,
                    weight_name,
                    expert_id,
                    shard_id,
424
425
426
427
428
429
430
431
432
                ) in expert_params_mapping:
                    if weight_name not in name:
                        continue

                    if is_pp_missing_parameter(name, self):
                        continue
                    name = name.replace(weight_name, param_name)
                    param = params_dict[name]
                    weight_loader = param.weight_loader
433
434
435
436
437
438
439
                    weight_loader(
                        param,
                        loaded_weight,
                        name,
                        shard_id=shard_id,
                        expert_id=expert_id,
                    )
440
441
442
443
444
445
446
447
448
                    break
                else:
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue
                    if is_pp_missing_parameter(name, self):
                        continue

                    param = params_dict[name]
449
450
451
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
452
453
454
455
                    weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

Mor Zusman's avatar
Mor Zusman committed
456

457
458
459
460
class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid):
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_substr={".self_attn.": ".", ".A_log": ".A"},
    )
Mor Zusman's avatar
Mor Zusman committed
461
462
463
464
465
466
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
467
        "gate_up_proj": ["gate_proj", "up_proj"],
468
        "in_proj": ["in_proj"],
Mor Zusman's avatar
Mor Zusman committed
469
470
471
472
473
474
475
476
477
    }

    # LoRA specific attributes
    embedding_modules = {
        "embed_tokens": "input_embeddings",
        "lm_head": "output_embeddings",
    }
    embedding_padding_modules = ["lm_head"]

478
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
479
480
481
482
        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        lora_config = vllm_config.lora_config
        scheduler_config = vllm_config.scheduler_config
483
        assert not cache_config.enable_prefix_caching, (
484
            "Jamba currently does not support prefix caching"
485
        )
486

Mor Zusman's avatar
Mor Zusman committed
487
488
        super().__init__()
        self.config = config
489
490
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
491
        self.scheduler_config = scheduler_config
492
493
494
        self.model = JambaModel(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Mor Zusman's avatar
Mor Zusman committed
495
496
497
498
499
500
501
502
503
504
        self.unpadded_vocab_size = config.vocab_size
        if lora_config:
            self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
        self.lm_head = ParallelLMHead(
            self.unpadded_vocab_size,
            config.hidden_size,
            org_num_embeddings=config.vocab_size,
            padding_size=DEFAULT_VOCAB_PADDING_SIZE
            # We need bigger padding if using lora for kernel
            # compatibility
505
506
            if not lora_config
            else lora_config.lora_vocab_padding_size,
507
            prefix=maybe_prefix(prefix, "lm_head"),
Mor Zusman's avatar
Mor Zusman committed
508
        )
509

510
511
512
        self.logits_processor = LogitsProcessor(
            self.unpadded_vocab_size, config.vocab_size
        )
Mor Zusman's avatar
Mor Zusman committed
513

514
        self.make_empty_intermediate_tensors = (
515
516
            self.model.make_empty_intermediate_tensors
        )
517

518
519
520
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

521
522
523
524
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
525
526
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
527
528
529
530
531
        **kwargs,
    ):
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
Mor Zusman's avatar
Mor Zusman committed
532
533
534
        return hidden_states

    def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
535
        return self.mamba_cache.copy_inputs_before_cuda_graphs(input_buffers, **kwargs)
Mor Zusman's avatar
Mor Zusman committed
536
537

    def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
538
        return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
Mor Zusman's avatar
Mor Zusman committed
539

540
541
542
543
544
545
546
547
548
549
550
    @classmethod
    def get_mamba_state_dtype_from_config(
        cls,
        vllm_config: "VllmConfig",
    ) -> tuple[torch.dtype, torch.dtype]:
        return MambaStateDtypeCalculator.mamba1_state_dtype(
            vllm_config.model_config.dtype,
            vllm_config.cache_config.mamba_cache_dtype,
            vllm_config.cache_config.mamba_ssm_cache_dtype,
        )

551
552
553
554
555
556
557
558
559
560
561
562
563
564
    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
        vllm_config: "VllmConfig",
    ) -> tuple[tuple[int, int], tuple[int, int]]:
        parallel_config = vllm_config.parallel_config
        hf_config = vllm_config.model_config.hf_config
        hidden_size = hf_config.hidden_size

        return MambaStateShapeCalculator.mamba1_state_shape(
            tp_world_size=parallel_config.tensor_parallel_size,
            intermediate_size=hf_config.mamba_expand * hidden_size,
            state_size=hf_config.mamba_d_state,
            conv_kernel=hf_config.mamba_d_conv,
Mor Zusman's avatar
Mor Zusman committed
565
566
        )

567
568
569
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
570
    ) -> torch.Tensor | None:
571
        logits = self.logits_processor(self.lm_head, hidden_states)
Mor Zusman's avatar
Mor Zusman committed
572
573
        return logits

574
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
575
576
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
577

578
579
    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
        return self.model.get_expert_mapping()
580
581
582


class JambaForSequenceClassification(JambaForCausalLM):
583
584
    is_pooling_model = True

585
586
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__(vllm_config=vllm_config, prefix=prefix)
587

588
589
        config = vllm_config.model_config.hf_config
        num_labels: int = config.num_labels
590
        score_bias: bool = getattr(config, "score_bias", False)
591
592
593
594
595
596
597
598

        # TODO: The original reward weights have float32 accuracy data, we
        # would like to load them in fp32 to get that extra precision.
        # Currently weight_loader passes the weight which is already in bf16
        self.score = nn.Linear(
            config.hidden_size,
            num_labels,
            bias=score_bias,
599
            dtype=vllm_config.model_config.head_dtype,
600
        )
601
602

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

605
606
607
608
609
610
611
612
613
        self.pooler = DispatchPooler(
            {
                "encode": Pooler.for_encode(pooler_config),
                "classify": Pooler.for_classify(
                    pooler_config,
                    classifier=self.score,
                ),
            }
        )