jamba.py 22 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
42
43
44
45
46
47
from .interfaces import (
    HasInnerState,
    IsHybrid,
    SupportsLoRA,
    SupportsMambaPrefixCaching,
    SupportsPP,
)
48
49
50
51
52
53
54
55
from .utils import (
    AutoWeightsLoader,
    WeightsMapper,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
56

Mor Zusman's avatar
Mor Zusman committed
57
58

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

75
        if self.num_total_experts > 1:
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
            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
97
98

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


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

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

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

176
        output = torch.empty_like(hidden_states)
177
        self.mamba(hidden_states, output)
Mor Zusman's avatar
Mor Zusman committed
178
        # Fully Connected
179
        hidden_states, residual = self.pre_ff_layernorm(output, residual)
Mor Zusman's avatar
Mor Zusman committed
180
181
182
183
184
        hidden_states = self.feed_forward(hidden_states)
        return hidden_states, residual


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

        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
239
            prefix=f"{prefix}.attn",
Mor Zusman's avatar
Mor Zusman committed
240
241
242
        )

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

    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)
268
        attn_output = self.attn(q, k, v)
Mor Zusman's avatar
Mor Zusman committed
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,
276
        residual: torch.Tensor | None,
Mor Zusman's avatar
Mor Zusman committed
277
278
279
280
281
282
        **kwargs,
    ):
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
283
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
Mor Zusman's avatar
Mor Zusman committed
284
285
286
287
288
289

        hidden_states = self.self_attention(
            positions=positions,
            hidden_states=hidden_states,
        )
        # Fully Connected
290
        hidden_states, residual = self.pre_ff_layernorm(hidden_states, residual)
Mor Zusman's avatar
Mor Zusman committed
291
292
293
294
295
296
        hidden_states = self.feed_forward(hidden_states)
        return hidden_states, residual


ALL_DECODER_LAYER_TYPES = {
    "attention": JambaAttentionDecoderLayer,
297
    "mamba": JambaMambaDecoderLayer,
Mor Zusman's avatar
Mor Zusman committed
298
299
300
}


301
@support_torch_compile
Mor Zusman's avatar
Mor Zusman committed
302
class JambaModel(nn.Module):
303
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Mor Zusman's avatar
Mor Zusman committed
304
        super().__init__()
305
306

        config = vllm_config.model_config.hf_config
307
        model_config = vllm_config.model_config
308
309
310
311
        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
312
        self.config = config
313
314
315
316
317
        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
318
319
320
321
322
323
324
325
326
        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,
        )

327
328
        extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)}

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

        self.start_layer, self.end_layer, self.layers = make_layers(
343
344
345
346
347
            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
        )
348

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

351
352
353
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Mor Zusman's avatar
Mor Zusman committed
354
355
356
357
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
358
359
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
Mor Zusman's avatar
Mor Zusman committed
360
    ) -> torch.Tensor:
361
362
363
364
365
366
        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
367
        else:
368
369
370
371
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

372
        for layer in islice(self.layers, self.start_layer, self.end_layer):
373
374
375
            hidden_states, residual = layer(
                positions=positions, hidden_states=hidden_states, residual=residual
            )
376

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

384
385
386
387
388
389
390
    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",
391
392
            num_experts=self.config.num_experts,
        )
393

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

Mor Zusman's avatar
Mor Zusman committed
464

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

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

493
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
494
495
496
        config = vllm_config.model_config.hf_config
        lora_config = vllm_config.lora_config
        scheduler_config = vllm_config.scheduler_config
497

Mor Zusman's avatar
Mor Zusman committed
498
499
        super().__init__()
        self.config = config
500
501
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
502
        self.scheduler_config = scheduler_config
503
504
505
        self.model = JambaModel(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Mor Zusman's avatar
Mor Zusman committed
506
507
508
509
510
511
512
513
514
515
        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
516
517
            if not lora_config
            else lora_config.lora_vocab_padding_size,
518
            prefix=maybe_prefix(prefix, "lm_head"),
Mor Zusman's avatar
Mor Zusman committed
519
        )
520

521
522
523
        self.logits_processor = LogitsProcessor(
            self.unpadded_vocab_size, config.vocab_size
        )
Mor Zusman's avatar
Mor Zusman committed
524

525
        self.make_empty_intermediate_tensors = (
526
527
            self.model.make_empty_intermediate_tensors
        )
528

529
530
531
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

532
533
534
535
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
536
537
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
538
539
540
541
542
        **kwargs,
    ):
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
Mor Zusman's avatar
Mor Zusman committed
543
544
545
        return hidden_states

    def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
546
        return self.mamba_cache.copy_inputs_before_cuda_graphs(input_buffers, **kwargs)
Mor Zusman's avatar
Mor Zusman committed
547
548

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

551
552
553
554
555
556
557
558
559
560
561
    @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,
        )

562
563
564
565
566
567
568
569
570
571
572
573
574
575
    @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
576
577
        )

578
579
580
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
581
    ) -> torch.Tensor | None:
582
        logits = self.logits_processor(self.lm_head, hidden_states)
Mor Zusman's avatar
Mor Zusman committed
583
584
        return logits

585
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
586
587
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
588

589
590
    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
        return self.model.get_expert_mapping()
591
592
593


class JambaForSequenceClassification(JambaForCausalLM):
594
595
    is_pooling_model = True

596
597
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__(vllm_config=vllm_config, prefix=prefix)
598

599
600
        config = vllm_config.model_config.hf_config
        num_labels: int = config.num_labels
601
        score_bias: bool = getattr(config, "score_bias", False)
602
603
604
605
606
607
608
609

        # 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,
610
            dtype=vllm_config.model_config.head_dtype,
611
        )
612
613

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

616
617
        self.pooler = DispatchPooler(
            {
618
619
620
                "token_classify": Pooler.for_token_classify(
                    pooler_config, classifier=self.score
                ),
621
                "classify": Pooler.for_classify(
622
623
624
625
                    pooler_config, classifier=self.score, act_fn="classify"
                ),
                "score": Pooler.for_classify(
                    pooler_config, classifier=self.score, act_fn="score"
626
627
628
                ),
            }
        )