jamba.py 21 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

import torch
from torch import nn
from transformers import JambaConfig

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

45
46
47
48
49
50
51
from .interfaces import (
    HasInnerState,
    IsHybrid,
    SupportsLoRA,
    SupportsMambaPrefixCaching,
    SupportsPP,
)
52
53
54
55
56
57
58
59
from .utils import (
    AutoWeightsLoader,
    WeightsMapper,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
60

Mor Zusman's avatar
Mor Zusman committed
61
62

class JambaMoE(nn.Module):
63
64
65
    def __init__(
        self,
        config: JambaConfig,
66
67
68
69
70
        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,
71
72
        prefix: str = "",
    ):
Mor Zusman's avatar
Mor Zusman committed
73
        super().__init__()
74
75
        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
76
        self.hidden_size = config.hidden_size
77
        self.intermediate_size = config.intermediate_size
Mor Zusman's avatar
Mor Zusman committed
78

79
        if self.num_total_experts > 1:
80
81
82
83
84
85
            self.router = ReplicatedLinear(
                self.hidden_size,
                self.num_total_experts,
                bias=False,
                quant_config=None,
                params_dtype=params_dtype,
86
                prefix=f"{prefix}.router",
87
88
89
90
91
92
93
94
95
96
97
98
99
100
            )

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

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
103
        orig_shape = hidden_states.shape
Mor Zusman's avatar
Mor Zusman committed
104
105
        hidden_states = hidden_states.view(-1, self.hidden_size)
        # router_logits: (batch * sequence_length, n_experts)
106
107
108
        if self.num_total_experts > 1:
            router_logits, _ = self.router(hidden_states)
        else:
109
110
111
112
113
            router_logits = torch.ones(
                (hidden_states.shape[0], 1),
                device=hidden_states.device,
                dtype=hidden_states.dtype,
            )
114
115
        hidden_states = self.experts(hidden_states, router_logits)
        return hidden_states.view(orig_shape)
Mor Zusman's avatar
Mor Zusman committed
116
117
118


class JambaMambaDecoderLayer(nn.Module):
119
120
121
122
    def __init__(
        self,
        config: JambaConfig,
        layer_idx: int,
123
124
125
126
        model_config: ModelConfig | None = None,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        is_lora_enabled: bool | None = False,
127
128
129
        prefix: str = "",
        **kwargs,
    ) -> None:
Mor Zusman's avatar
Mor Zusman committed
130
131
        super().__init__()
        self.config = config
132
        self.is_lora_enabled = is_lora_enabled
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
        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
149
150

        num_experts = config.layers_num_experts[layer_idx]
151
152
153
154
155
156
157
158
159
160
161
162
163
164
        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",
            )
165
166
        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
167
168
169
170

    def forward(
        self,
        hidden_states: torch.Tensor,
171
        residual: torch.Tensor | None,
Mor Zusman's avatar
Mor Zusman committed
172
173
174
175
176
177
        **kwargs,
    ):
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
178
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
Mor Zusman's avatar
Mor Zusman committed
179

180
        output = torch.empty_like(hidden_states)
181
        self.mamba(hidden_states, output)
Mor Zusman's avatar
Mor Zusman committed
182
        # Fully Connected
183
        hidden_states, residual = self.pre_ff_layernorm(output, residual)
Mor Zusman's avatar
Mor Zusman committed
184
185
186
187
188
        hidden_states = self.feed_forward(hidden_states)
        return hidden_states, residual


class JambaAttentionDecoderLayer(nn.Module):
189
190
191
192
    def __init__(
        self,
        config: JambaConfig,
        layer_idx: int,
193
194
195
        model_config: ModelConfig | None = None,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
196
197
198
        prefix: str = "",
        **kwargs,
    ) -> None:
Mor Zusman's avatar
Mor Zusman committed
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
        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,
227
            prefix=f"{prefix}.qkv_proj",
Mor Zusman's avatar
Mor Zusman committed
228
        )
229
230
231
232
233
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
234
            prefix=f"{prefix}.o_proj",
235
        )
Mor Zusman's avatar
Mor Zusman committed
236
237
238
239
240
241
242

        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
243
            prefix=f"{prefix}.attn",
Mor Zusman's avatar
Mor Zusman committed
244
245
246
        )

        num_experts = config.layers_num_experts[layer_idx]
247
248
249
250
251
252
253
254
255
256
257
258
259
260
        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",
            )
261
262
        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
263
264
265
266
267
268
269
270
271

    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)
272
        attn_output = self.attn(q, k, v)
Mor Zusman's avatar
Mor Zusman committed
273
274
275
276
277
278
279
        output, _ = self.o_proj(attn_output)
        return output

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
280
        residual: torch.Tensor | None,
Mor Zusman's avatar
Mor Zusman committed
281
282
283
284
285
286
        **kwargs,
    ):
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
287
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
Mor Zusman's avatar
Mor Zusman committed
288
289
290
291
292
293

        hidden_states = self.self_attention(
            positions=positions,
            hidden_states=hidden_states,
        )
        # Fully Connected
294
        hidden_states, residual = self.pre_ff_layernorm(hidden_states, residual)
Mor Zusman's avatar
Mor Zusman committed
295
296
297
298
299
300
        hidden_states = self.feed_forward(hidden_states)
        return hidden_states, residual


ALL_DECODER_LAYER_TYPES = {
    "attention": JambaAttentionDecoderLayer,
301
    "mamba": JambaMambaDecoderLayer,
Mor Zusman's avatar
Mor Zusman committed
302
303
304
}


305
@support_torch_compile
Mor Zusman's avatar
Mor Zusman committed
306
class JambaModel(nn.Module):
307
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Mor Zusman's avatar
Mor Zusman committed
308
        super().__init__()
309
310

        config = vllm_config.model_config.hf_config
311
        model_config = vllm_config.model_config
312
313
314
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

Mor Zusman's avatar
Mor Zusman committed
315
        self.config = config
316
317

        self.vocab_size = config.vocab_size
Mor Zusman's avatar
Mor Zusman committed
318
319
320
321
322
323

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

324
325
        extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)}

326
327
        def get_layer(prefix: str):
            layer_idx = int(prefix.rsplit(".", 1)[1])
328
329
330
331
332
333
334
335
336
337
            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,
            )
338
339

        self.start_layer, self.end_layer, self.layers = make_layers(
340
341
342
343
344
            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
        )
345

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

348
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
349
350
        return self.embed_tokens(input_ids)

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

369
        for layer in islice(self.layers, self.start_layer, self.end_layer):
370
371
372
            hidden_states, residual = layer(
                positions=positions, hidden_states=hidden_states, residual=residual
            )
373

374
        if not get_pp_group().is_last_rank:
375
376
377
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
Mor Zusman's avatar
Mor Zusman committed
378
379
380
        hidden_states, _ = self.final_layernorm(hidden_states, residual)
        return hidden_states

381
382
383
    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)
384
        return fused_moe_make_expert_params_mapping(
385
            self,
386
387
388
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
389
390
            num_experts=self.config.num_experts,
        )
391

392
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
        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
411
                if "experts" in name:
412
413
414
415
416
417
418
419
420
421
422
423
424
425
                    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 (
426
427
428
429
                    param_name,
                    weight_name,
                    expert_id,
                    shard_id,
430
431
432
433
434
435
436
437
438
                ) 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
439
440
441
442
443
444
445
                    weight_loader(
                        param,
                        loaded_weight,
                        name,
                        shard_id=shard_id,
                        expert_id=expert_id,
                    )
446
447
448
449
450
451
452
453
454
                    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]
455
456
457
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
458
459
460
461
                    weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

Mor Zusman's avatar
Mor Zusman committed
462

463
464
465
466
467
468
469
470
class JambaForCausalLM(
    nn.Module,
    HasInnerState,
    SupportsLoRA,
    SupportsPP,
    IsHybrid,
    SupportsMambaPrefixCaching,
):
471
472
473
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_substr={".self_attn.": ".", ".A_log": ".A"},
    )
Mor Zusman's avatar
Mor Zusman committed
474
475
476
477
478
479
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
480
        "gate_up_proj": ["gate_proj", "up_proj"],
481
        "in_proj": ["in_proj"],
Mor Zusman's avatar
Mor Zusman committed
482
483
484
485
486
487
488
489
    }

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

490
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
491
        config = vllm_config.model_config.hf_config
492

493
        scheduler_config = vllm_config.scheduler_config
494

Mor Zusman's avatar
Mor Zusman committed
495
496
        super().__init__()
        self.config = config
497
498
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
499
        self.scheduler_config = scheduler_config
500
501
502
        self.model = JambaModel(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
503

Mor Zusman's avatar
Mor Zusman committed
504
        self.lm_head = ParallelLMHead(
505
            config.vocab_size,
Mor Zusman's avatar
Mor Zusman committed
506
            config.hidden_size,
507
            prefix=maybe_prefix(prefix, "lm_head"),
Mor Zusman's avatar
Mor Zusman committed
508
        )
509

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

512
        self.make_empty_intermediate_tensors = (
513
514
            self.model.make_empty_intermediate_tensors
        )
515

516
517
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
518

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

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

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

538
539
540
541
542
543
544
545
546
547
548
    @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,
        )

549
550
551
552
553
554
555
556
557
558
559
560
561
562
    @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
563
564
        )

565
566
567
568
    @classmethod
    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
        return MambaStateCopyFuncCalculator.mamba1_state_copy_func()

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

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

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


class JambaForSequenceClassification(JambaForCausalLM):
585
586
    is_pooling_model = True

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

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

        # 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,
601
            dtype=vllm_config.model_config.head_dtype,
602
        )
603
604

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

607
        self.pooler = DispatchPooler.for_seq_cls(pooler_config, classifier=self.score)