jamba.py 25.1 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
from collections.abc import Iterable
5
from itertools import islice
6
from typing import Optional
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 import envs
Mor Zusman's avatar
Mor Zusman committed
13
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
from vllm.model_executor.layers.linear import (QKVParallelLinear,
Mor Zusman's avatar
Mor Zusman committed
21
22
23
                                               ReplicatedLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
24
from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
25
from vllm.model_executor.layers.mamba.mamba_utils import (
26
    MambaStateDtypeCalculator, MambaStateShapeCalculator)
27
from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
28
from vllm.model_executor.layers.quantization import QuantizationConfig
Mor Zusman's avatar
Mor Zusman committed
29
30
from vllm.model_executor.layers.vocab_parallel_embedding import (
    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
31
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
32
from vllm.model_executor.models.llama import LlamaMLP as JambaMLP
33
34
from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
                                                    MambaCacheParams)
35
from vllm.sequence import IntermediateTensors
36
from vllm.utils import LayerBlockType
Mor Zusman's avatar
Mor Zusman committed
37

38
from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
39
from .utils import (AutoWeightsLoader, WeightsMapper, is_pp_missing_parameter,
40
41
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
42

Mor Zusman's avatar
Mor Zusman committed
43
44
45

class JambaMoE(nn.Module):

46
47
48
49
50
51
    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,
52
53
                 quant_config: Optional[QuantizationConfig] = None,
                 prefix: str = ""):
Mor Zusman's avatar
Mor Zusman committed
54
        super().__init__()
55
56
        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
57
        self.hidden_size = config.hidden_size
58
        self.intermediate_size = config.intermediate_size
Mor Zusman's avatar
Mor Zusman committed
59

60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
        if self.num_total_experts > 1:
            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,
76
77
                                quant_config=quant_config,
                                prefix=f"{prefix}.experts")
Mor Zusman's avatar
Mor Zusman committed
78
79

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
80
        orig_shape = hidden_states.shape
Mor Zusman's avatar
Mor Zusman committed
81
82
        hidden_states = hidden_states.view(-1, self.hidden_size)
        # router_logits: (batch * sequence_length, n_experts)
83
84
85
86
87
88
89
90
        if self.num_total_experts > 1:
            router_logits, _ = self.router(hidden_states)
        else:
            router_logits = torch.ones((hidden_states.shape[0], 1),
                                       device=hidden_states.device,
                                       dtype=hidden_states.dtype)
        hidden_states = self.experts(hidden_states, router_logits)
        return hidden_states.view(orig_shape)
Mor Zusman's avatar
Mor Zusman committed
91
92
93
94
95
96
97


class JambaMambaDecoderLayer(nn.Module):

    def __init__(self,
                 config: JambaConfig,
                 layer_idx: int,
98
                 model_config: Optional[ModelConfig] = None,
Mor Zusman's avatar
Mor Zusman committed
99
                 cache_config: Optional[CacheConfig] = None,
100
                 quant_config: Optional[QuantizationConfig] = None,
101
                 is_lora_enabled: Optional[bool] = False,
102
                 prefix: str = "",
103
                 **kwargs) -> None:
Mor Zusman's avatar
Mor Zusman committed
104
105
        super().__init__()
        self.config = config
106
        self.is_lora_enabled = is_lora_enabled
107
108
109
110
111
112
113
114
115
116
        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,
117
                                activation=config.hidden_act,
118
                                is_lora_enabled = self.is_lora_enabled,
119
120
                                model_config=model_config,
                                cache_config=cache_config,
121
                                prefix=f"{prefix}.mixer",
122
                                )
Mor Zusman's avatar
Mor Zusman committed
123
124

        num_experts = config.layers_num_experts[layer_idx]
125
126
127
128
129
130
131
132
133
134
135
136
137
138
        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",
            )
Mor Zusman's avatar
Mor Zusman committed
139
140
141
142
143
144
145
146
147
        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)

    def forward(
        self,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor],
148
        mamba_cache_params: MambaCacheParams,
Mor Zusman's avatar
Mor Zusman committed
149
150
151
152
153
154
155
156
157
        **kwargs,
    ):
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)

158
159
        output = torch.empty_like(hidden_states)
        self.mamba(hidden_states, output, mamba_cache_params)
Mor Zusman's avatar
Mor Zusman committed
160
        # Fully Connected
161
        hidden_states, residual = self.pre_ff_layernorm(output, residual)
Mor Zusman's avatar
Mor Zusman committed
162
163
164
165
166
167
        hidden_states = self.feed_forward(hidden_states)
        return hidden_states, residual


class JambaAttentionDecoderLayer(nn.Module):

168
169
170
    def __init__(self,
                 config: JambaConfig,
                 layer_idx: int,
171
                 model_config: Optional[ModelConfig] = None,
172
173
174
175
                 cache_config: Optional[CacheConfig] = None,
                 quant_config: Optional[QuantizationConfig] = None,
                 prefix: str = "",
                 **kwargs) -> None:
Mor Zusman's avatar
Mor Zusman committed
176
177
178
179
180
181
182
183
184
185
186
187
188
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
        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,
        )
        self.o_proj = RowParallelLinear(self.total_num_heads * self.head_dim,
                                        config.hidden_size,
                                        bias=False,
                                        quant_config=quant_config)

        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
216
            prefix=f"{prefix}.attn",
Mor Zusman's avatar
Mor Zusman committed
217
218
219
        )

        num_experts = config.layers_num_experts[layer_idx]
220
221
222
223
224
225
226
227
228
229
230
231
232
233
        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",
            )
Mor Zusman's avatar
Mor Zusman committed
234
235
236
237
238
239
240
241
242
243
244
245
246
        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)

    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)
247
        attn_output = self.attn(q, k, v)
Mor Zusman's avatar
Mor Zusman committed
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
        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:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)

        hidden_states = self.self_attention(
            positions=positions,
            hidden_states=hidden_states,
        )
        # Fully Connected
        hidden_states, residual = self.pre_ff_layernorm(
            hidden_states, residual)
        hidden_states = self.feed_forward(hidden_states)
        return hidden_states, residual


ALL_DECODER_LAYER_TYPES = {
    "attention": JambaAttentionDecoderLayer,
    "mamba": JambaMambaDecoderLayer
}


282
@support_torch_compile
Mor Zusman's avatar
Mor Zusman committed
283
284
class JambaModel(nn.Module):

285
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Mor Zusman's avatar
Mor Zusman committed
286
        super().__init__()
287
288

        config = vllm_config.model_config.hf_config
289
        model_config = vllm_config.model_config
290
291
292
293
        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
294
295
296
297
298
299
300
301
302
303
304
305
        self.config = config
        lora_vocab = ((lora_config.lora_extra_vocab_size *
                       (lora_config.max_loras or 1)) if lora_config else 0)
        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,
        )

306
307
        extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)}

308
309
310
311
        def get_layer(prefix: str):
            layer_idx = int(prefix.rsplit(".", 1)[1])
            layer_class = ALL_DECODER_LAYER_TYPES[
                config.layers_block_type[layer_idx]]
312
313
            return layer_class(config,
                               layer_idx,
314
                               model_config,
315
316
317
318
                               cache_config,
                               quant_config=quant_config,
                               prefix=prefix,
                               **extra_kwargs)
319
320
321
322
323
324
325

        self.start_layer, self.end_layer, self.layers = make_layers(
            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))

Mor Zusman's avatar
Mor Zusman committed
326
327
328
        self.final_layernorm = RMSNorm(config.hidden_size,
                                       eps=config.rms_norm_eps)

329
330
331
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Mor Zusman's avatar
Mor Zusman committed
332
333
334
335
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
336
        mamba_cache_params: MambaCacheParams,
337
        intermediate_tensors: Optional[IntermediateTensors] = None,
338
        inputs_embeds: Optional[torch.Tensor] = None,
Mor Zusman's avatar
Mor Zusman committed
339
    ) -> torch.Tensor:
340
341
342
343
344
345
        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
346
        else:
347
348
349
350
351
352
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

        kv_cache_index = 0
        mamba_cache_index = 0
353
        for layer in islice(self.layers, self.start_layer, self.end_layer):
354
            layer_mamba_cache_params = None
Mor Zusman's avatar
Mor Zusman committed
355
            if isinstance(layer, JambaAttentionDecoderLayer):
356
                kv_cache_index += 1
357
358
            if isinstance(layer,
                          JambaMambaDecoderLayer) and mamba_cache_params:
359
                current_state_layer = mamba_cache_index
360
361
                layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
                    current_state_layer)
362
                mamba_cache_index += 1
Mor Zusman's avatar
Mor Zusman committed
363
364
365
366
367

            hidden_states, residual = layer(
                positions=positions,
                hidden_states=hidden_states,
                residual=residual,
368
                mamba_cache_params=layer_mamba_cache_params)
369
370
371
372
373
        if not get_pp_group().is_last_rank:
            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
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
    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",
            num_experts=self.config.num_experts)

    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
        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
                if 'experts' in name:
                    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 (
                        param_name,
                        weight_name,
                        expert_id,
                        shard_id,
                ) 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
                    weight_loader(param,
                                  loaded_weight,
                                  name,
                                  shard_id=shard_id,
                                  expert_id=expert_id)
                    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]
                    weight_loader = getattr(param, "weight_loader",
                                            default_weight_loader)
                    weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

Mor Zusman's avatar
Mor Zusman committed
454

455
class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
456
                       IsHybrid):
457
458
459
460
    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
484
485
        assert not cache_config.enable_prefix_caching, \
            "Jamba currently does not support prefix caching"

Mor Zusman's avatar
Mor Zusman committed
486
487
        super().__init__()
        self.config = config
488
489
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
490
        self.scheduler_config = scheduler_config
491
492
        self.model = JambaModel(vllm_config=vllm_config,
                                prefix=maybe_prefix(prefix, "model"))
Mor Zusman's avatar
Mor Zusman committed
493
494
495
496
497
498
499
500
501
502
503
        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
            if not lora_config else lora_config.lora_vocab_padding_size,
504
            prefix=maybe_prefix(prefix, "lm_head"),
Mor Zusman's avatar
Mor Zusman committed
505
506
        )
        # Used to track and store by the Mamba cache between steps.
507
508
        self.mamba_cache: Optional[MambaCacheManager] = None

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

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

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

Mor Zusman's avatar
Mor Zusman committed
518
519
520
521
    def forward(self,
                input_ids: torch.Tensor,
                positions: torch.Tensor,
                intermediate_tensors: Optional[IntermediateTensors] = None,
522
                inputs_embeds: Optional[torch.Tensor] = None,
Mor Zusman's avatar
Mor Zusman committed
523
                **kwargs):
524
525
526
527
528
529
530
531
        # NOTE: mamba_cache_params is not needed for v1
        mamba_cache_params = None
        if not envs.VLLM_USE_V1:
            if self.mamba_cache is None:
                num_layers = self.model_config.get_num_layers_by_block_type(
                    self.vllm_config.parallel_config, LayerBlockType.mamba)
                state_shape = self.get_mamba_state_shape_from_config(
                    self.vllm_config)
532
533
                state_dtype = self.get_mamba_state_dtype_from_config(
                    self.vllm_config)
534
                self.mamba_cache = MambaCacheManager(self.vllm_config,
535
536
                                                     num_layers, *state_shape,
                                                     *state_dtype)
537
538

            mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
539

540
        hidden_states = self.model(input_ids, positions, mamba_cache_params,
541
                                   intermediate_tensors, inputs_embeds)
Mor Zusman's avatar
Mor Zusman committed
542
543
544
        return hidden_states

    def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
545
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
562
    @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,
        )

563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
    @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,
            use_v1=envs.VLLM_USE_V1,
Mor Zusman's avatar
Mor Zusman committed
578
579
        )

580
581
582
583
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> Optional[torch.Tensor]:
584
        logits = self.logits_processor(self.lm_head, hidden_states)
Mor Zusman's avatar
Mor Zusman committed
585
586
        return logits

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

592
593
    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
        return self.model.get_expert_mapping()
594
595
596
597


class JambaForSequenceClassification(JambaForCausalLM):

598
599
    is_pooling_model = True

600
601
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__(vllm_config=vllm_config, prefix=prefix)
602

603
604
605
        config = vllm_config.model_config.hf_config
        num_labels: int = config.num_labels
        score_bias: bool = getattr(config, 'score_bias', False)
606
607
608
609
610
611
612
613

        # 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,
614
            dtype=vllm_config.model_config.head_dtype,
615
        )
616
617

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

620
621
622
623
624
625
626
627
628
        self.pooler = DispatchPooler({
            "encode":
            Pooler.for_encode(pooler_config),
            "classify":
            Pooler.for_classify(
                pooler_config,
                classifier=self.score,
            ),
        })