jamba.py 21.9 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
7
from typing import Optional
Mor Zusman's avatar
Mor Zusman committed
8
9
10
11
12
13

import torch
from torch import nn
from transformers import JambaConfig

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

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

Mor Zusman's avatar
Mor Zusman committed
52
53

class JambaMoE(nn.Module):
54
55
56
57
58
59
60
61
62
63
    def __init__(
        self,
        config: JambaConfig,
        num_experts: Optional[int] = None,
        top_k: Optional[int] = None,
        params_dtype: Optional[torch.dtype] = None,
        tp_size: Optional[int] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
Mor Zusman's avatar
Mor Zusman committed
64
        super().__init__()
65
66
        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
67
        self.hidden_size = config.hidden_size
68
        self.intermediate_size = config.intermediate_size
Mor Zusman's avatar
Mor Zusman committed
69

70
        if self.num_total_experts > 1:
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
            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
92
93

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


class JambaMambaDecoderLayer(nn.Module):
110
111
112
113
114
115
116
117
118
119
120
    def __init__(
        self,
        config: JambaConfig,
        layer_idx: int,
        model_config: Optional[ModelConfig] = None,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        is_lora_enabled: Optional[bool] = False,
        prefix: str = "",
        **kwargs,
    ) -> None:
Mor Zusman's avatar
Mor Zusman committed
121
122
        super().__init__()
        self.config = config
123
        self.is_lora_enabled = is_lora_enabled
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
        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
140
141

        num_experts = config.layers_num_experts[layer_idx]
142
143
144
145
146
147
148
149
150
151
152
153
154
155
        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",
            )
156
157
        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
158
159
160
161
162
163
164
165
166
167
168

    def forward(
        self,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor],
        **kwargs,
    ):
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
169
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
Mor Zusman's avatar
Mor Zusman committed
170

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


class JambaAttentionDecoderLayer(nn.Module):
180
181
182
183
184
185
186
187
188
189
    def __init__(
        self,
        config: JambaConfig,
        layer_idx: int,
        model_config: Optional[ModelConfig] = None,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        **kwargs,
    ) -> None:
Mor Zusman's avatar
Mor Zusman committed
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
218
        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,
        )
219
220
221
222
223
224
        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
225
226
227
228
229
230
231

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

        num_experts = config.layers_num_experts[layer_idx]
236
237
238
239
240
241
242
243
244
245
246
247
248
249
        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",
            )
250
251
        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
252
253
254
255
256
257
258
259
260

    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)
261
        attn_output = self.attn(q, k, v)
Mor Zusman's avatar
Mor Zusman committed
262
263
264
265
266
267
268
269
270
271
272
273
274
275
        output, _ = self.o_proj(attn_output)
        return output

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor],
        **kwargs,
    ):
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
276
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
Mor Zusman's avatar
Mor Zusman committed
277
278
279
280
281
282

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


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


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

        config = vllm_config.model_config.hf_config
300
        model_config = vllm_config.model_config
301
302
303
304
        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
305
        self.config = config
306
307
308
309
310
        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
311
312
313
314
315
316
317
318
319
        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,
        )

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

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

        self.start_layer, self.end_layer, self.layers = make_layers(
336
337
338
339
340
            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
        )
341

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

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

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

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

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

377
378
379
380
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)
        return FusedMoE.make_expert_params_mapping(
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
384
385
            num_experts=self.config.num_experts,
        )
386

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

Mor Zusman's avatar
Mor Zusman committed
457

458
459
460
461
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
462
463
464
465
466
467
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
468
        "gate_up_proj": ["gate_proj", "up_proj"],
469
        "in_proj": ["in_proj"],
Mor Zusman's avatar
Mor Zusman committed
470
471
472
473
474
475
476
477
478
    }

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

479
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
480
481
482
483
        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
484
        assert not cache_config.enable_prefix_caching, (
485
            "Jamba currently does not support prefix caching"
486
        )
487

Mor Zusman's avatar
Mor Zusman committed
488
489
        super().__init__()
        self.config = config
490
491
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
492
        self.scheduler_config = scheduler_config
493
494
495
        self.model = JambaModel(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Mor Zusman's avatar
Mor Zusman committed
496
497
498
499
500
501
502
503
504
505
        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
506
507
            if not lora_config
            else lora_config.lora_vocab_padding_size,
508
            prefix=maybe_prefix(prefix, "lm_head"),
Mor Zusman's avatar
Mor Zusman committed
509
        )
510

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

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

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

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

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

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

541
542
543
544
545
546
547
548
549
550
551
    @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,
        )

552
553
554
555
556
557
558
559
560
561
562
563
564
565
    @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
566
567
        )

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

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

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


class JambaForSequenceClassification(JambaForCausalLM):
584
585
    is_pooling_model = True

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

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

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

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

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