jamba.py 24.9 KB
Newer Older
1
"""Inference-only Jamba model."""
2
from typing import Iterable, List, Optional, Set, Tuple
Mor Zusman's avatar
Mor Zusman committed
3
4
5
6
7
8
9

import torch
from torch import nn
from transformers import JambaConfig

from vllm.attention.backends.abstract import AttentionMetadata
from vllm.attention.layer import Attention
10
from vllm.config import CacheConfig, VllmConfig
11
from vllm.distributed import get_tensor_model_parallel_world_size
12
from vllm.distributed.parallel_state import get_pp_group
13
from vllm.model_executor.layers.fused_moe import FusedMoE
Mor Zusman's avatar
Mor Zusman committed
14
from vllm.model_executor.layers.layernorm import RMSNorm
15
from vllm.model_executor.layers.linear import (QKVParallelLinear,
Mor Zusman's avatar
Mor Zusman committed
16
17
18
                                               ReplicatedLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
20
from vllm.model_executor.layers.pooler import Pooler, PoolingType
21
from vllm.model_executor.layers.quantization import QuantizationConfig
Joe Runde's avatar
Joe Runde committed
22
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
Mor Zusman's avatar
Mor Zusman committed
23
24
from vllm.model_executor.layers.vocab_parallel_embedding import (
    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
25
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
26
27
from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
                                                    MambaCacheParams)
28
from vllm.model_executor.pooling_metadata import PoolingMetadata
Mor Zusman's avatar
Mor Zusman committed
29
from vllm.model_executor.sampling_metadata import SamplingMetadata
30
from vllm.sequence import IntermediateTensors, PoolerOutput
31
from vllm.utils import LayerBlockType
Mor Zusman's avatar
Mor Zusman committed
32

33
34
35
36
from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
from .utils import (is_pp_missing_parameter,
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
37

Mor Zusman's avatar
Mor Zusman committed
38
39
40
41
42
KVCache = Tuple[torch.Tensor, torch.Tensor]


class JambaMoE(nn.Module):

43
44
45
46
47
48
49
    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):
Mor Zusman's avatar
Mor Zusman committed
50
        super().__init__()
51
52
        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
53
        self.hidden_size = config.hidden_size
54
        self.intermediate_size = config.intermediate_size
Mor Zusman's avatar
Mor Zusman committed
55

56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
        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,
                                quant_config=quant_config)
Mor Zusman's avatar
Mor Zusman committed
73
74

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
75
        orig_shape = hidden_states.shape
Mor Zusman's avatar
Mor Zusman committed
76
77
        hidden_states = hidden_states.view(-1, self.hidden_size)
        # router_logits: (batch * sequence_length, n_experts)
78
79
80
81
82
83
84
85
        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
86
87


88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
class JambaMLP(JambaMoE):

    def __init__(self,
                 config: JambaConfig,
                 params_dtype: Optional[torch.dtype] = None,
                 tp_size: Optional[int] = None,
                 quant_config: Optional[QuantizationConfig] = None):
        super().__init__(config,
                         num_experts=1,
                         top_k=1,
                         params_dtype=params_dtype,
                         tp_size=tp_size,
                         quant_config=quant_config)


Mor Zusman's avatar
Mor Zusman committed
103
104
105
106
107
108
class JambaMambaDecoderLayer(nn.Module):

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

        num_experts = config.layers_num_experts[layer_idx]
        ffn_layer_class = JambaMoE if num_experts > 1 else JambaMLP
        self.feed_forward = ffn_layer_class(config, quant_config=quant_config)
        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,
        attn_metadata: AttentionMetadata,
        residual: Optional[torch.Tensor],
138
        mamba_cache_params: MambaCacheParams,
Mor Zusman's avatar
Mor Zusman committed
139
140
141
142
143
144
145
146
147
        **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)

148
149
        hidden_states = self.mamba(hidden_states, attn_metadata,
                                   mamba_cache_params)
Mor Zusman's avatar
Mor Zusman committed
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
        # Fully Connected
        hidden_states, residual = self.pre_ff_layernorm(
            hidden_states, residual)
        hidden_states = self.feed_forward(hidden_states)
        return hidden_states, residual


class JambaAttentionDecoderLayer(nn.Module):

    def __init__(
        self,
        config: JambaConfig,
        layer_idx: int,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
165
        prefix: str = "",
Mor Zusman's avatar
Mor Zusman committed
166
167
168
169
170
171
172
173
174
175
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
    ) -> None:
        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,
207
            prefix=f"{prefix}.attn",
Mor Zusman's avatar
Mor Zusman committed
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
        )

        num_experts = config.layers_num_experts[layer_idx]
        ffn_layer_class = JambaMoE if num_experts > 1 else JambaMLP
        self.feed_forward = ffn_layer_class(config, quant_config=quant_config)
        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,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
        **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)
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
        output, _ = self.o_proj(attn_output)
        return output

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
        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,
            kv_cache=kv_cache,
            attn_metadata=attn_metadata,
        )
        # 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
}


class JambaModel(nn.Module):

269
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Mor Zusman's avatar
Mor Zusman committed
270
        super().__init__()
271
272
273
274
275
276

        config = vllm_config.model_config.hf_config
        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
277
278
279
280
281
282
283
284
285
286
287
288
289
        self.config = config
        self.padding_idx = config.pad_token_id
        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,
        )

290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
        def get_layer(prefix: str):
            layer_idx = int(prefix.rsplit(".", 1)[1])
            layer_class = ALL_DECODER_LAYER_TYPES[
                config.layers_block_type[layer_idx]]
            return layer_class(
                config,
                layer_idx,
                cache_config,
                quant_config=quant_config,
                prefix=prefix,
            )

        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
308
309
310
        self.final_layernorm = RMSNorm(config.hidden_size,
                                       eps=config.rms_norm_eps)

311
312
313
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Mor Zusman's avatar
Mor Zusman committed
314
315
316
317
318
319
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
320
        mamba_cache_params: MambaCacheParams,
321
        intermediate_tensors: Optional[IntermediateTensors] = None,
322
        inputs_embeds: Optional[torch.Tensor] = None,
Mor Zusman's avatar
Mor Zusman committed
323
    ) -> torch.Tensor:
324
325
326
327
328
329
        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
330
        else:
331
332
333
334
335
336
337
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

        kv_cache_index = 0
        mamba_cache_index = 0
        for i in range(self.start_layer, self.end_layer):
Mor Zusman's avatar
Mor Zusman committed
338
339
            layer = self.layers[i]
            kv_cache = None
340
            layer_mamba_cache_params = None
Mor Zusman's avatar
Mor Zusman committed
341
            if isinstance(layer, JambaAttentionDecoderLayer):
342
343
                kv_cache = kv_caches[kv_cache_index]
                kv_cache_index += 1
Mor Zusman's avatar
Mor Zusman committed
344
            if isinstance(layer, JambaMambaDecoderLayer):
345
                current_state_layer = mamba_cache_index
346
347
                layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
                    current_state_layer)
348
                mamba_cache_index += 1
Mor Zusman's avatar
Mor Zusman committed
349
350
351
352
353
354
355

            hidden_states, residual = layer(
                positions=positions,
                hidden_states=hidden_states,
                kv_cache=kv_cache,
                attn_metadata=attn_metadata,
                residual=residual,
356
                mamba_cache_params=layer_mamba_cache_params)
357
358
359
360
361
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
                "residual": residual
            })
Mor Zusman's avatar
Mor Zusman committed
362
363
364
365
        hidden_states, _ = self.final_layernorm(hidden_states, residual)
        return hidden_states


366
367
class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
                       IsHybrid):
Mor Zusman's avatar
Mor Zusman committed
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
    }

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

389
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
390
391
392
393
        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
394
395
396
        assert not cache_config.enable_prefix_caching, \
            "Jamba currently does not support prefix caching"

Mor Zusman's avatar
Mor Zusman committed
397
398
        super().__init__()
        self.config = config
399
400
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
401
        self.scheduler_config = scheduler_config
402
403
        self.model = JambaModel(vllm_config=vllm_config,
                                prefix=maybe_prefix(prefix, "model"))
Mor Zusman's avatar
Mor Zusman committed
404
405
406
407
408
409
410
411
412
413
414
415
416
        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,
        )
        # Used to track and store by the Mamba cache between steps.
417
418
        self.mamba_cache: Optional[MambaCacheManager] = None

Mor Zusman's avatar
Mor Zusman committed
419
420
        self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                config.vocab_size)
Joe Runde's avatar
Joe Runde committed
421
        self.sampler = get_sampler()
Mor Zusman's avatar
Mor Zusman committed
422

423
424
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
425
426
427
428
429
430
431
432
433
434
435
        if self.scheduler_config is not None and \
            not self.model_config.enforce_eager:
            if self.scheduler_config.max_num_seqs > \
                vllm_config.compilation_config.max_capture_size:
                self.max_batch_size = \
                    vllm_config.compilation_config.max_capture_size
            else:
                self.max_batch_size = vllm_config.pad_for_cudagraph(
                    self.scheduler_config.max_num_seqs)
        else:
            self.max_batch_size = 8192 + 2
436

437
438
439
    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
440
441
442
443
444
445
    def forward(self,
                input_ids: torch.Tensor,
                positions: torch.Tensor,
                kv_caches: List[KVCache],
                attn_metadata: AttentionMetadata,
                intermediate_tensors: Optional[IntermediateTensors] = None,
446
                inputs_embeds: Optional[torch.Tensor] = None,
Mor Zusman's avatar
Mor Zusman committed
447
                **kwargs):
448
449
        if self.mamba_cache is None:

450
451
            num_mamba_layers = self.model_config.get_num_layers_by_block_type(
                self.vllm_config.parallel_config, LayerBlockType.mamba)
452
            self.mamba_cache = MambaCacheManager(
453
454
                self.lm_head.weight.dtype, num_mamba_layers,
                self.max_batch_size, *self._get_mamba_cache_shape())
455
456
457
458
459
460
461
462
        (
            mamba_cache_tensors,
            state_indices_tensor,
        ) = self.mamba_cache.current_run_tensors(input_ids, attn_metadata,
                                                 **kwargs)
        mamba_cache_params = MambaCacheParams(mamba_cache_tensors[0],
                                              mamba_cache_tensors[1],
                                              state_indices_tensor)
Mor Zusman's avatar
Mor Zusman committed
463
        hidden_states = self.model(input_ids, positions, kv_caches,
464
                                   attn_metadata, mamba_cache_params,
465
                                   intermediate_tensors, inputs_embeds)
Mor Zusman's avatar
Mor Zusman committed
466
467
468
        return hidden_states

    def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
469
470
        return self.mamba_cache.copy_inputs_before_cuda_graphs(
            input_buffers, **kwargs)
Mor Zusman's avatar
Mor Zusman committed
471
472

    def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
473
        return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
Mor Zusman's avatar
Mor Zusman committed
474
475

    def _get_mamba_cache_shape(
476
            self) -> Tuple[Tuple[int, int], Tuple[int, int]]:
Mor Zusman's avatar
Mor Zusman committed
477
478
479
480
        world_size = get_tensor_model_parallel_world_size()
        hidden_size = self.config.hidden_size
        conv_state_shape = (
            self.config.mamba_expand * hidden_size // world_size,
481
            self.config.mamba_d_conv - 1,
Mor Zusman's avatar
Mor Zusman committed
482
483
        )
        temporal_state_shape = (
484
            self.config.mamba_expand * hidden_size // world_size,
Mor Zusman's avatar
Mor Zusman committed
485
486
487
488
            self.config.mamba_d_state,
        )
        return conv_state_shape, temporal_state_shape

489
490
491
492
493
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
494
        logits = self.logits_processor(self.lm_head, hidden_states,
Mor Zusman's avatar
Mor Zusman committed
495
496
497
498
499
500
501
502
503
504
505
                                       sampling_metadata)
        return logits

    def sample(
        self,
        logits: Optional[torch.Tensor],
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
        next_tokens = self.sampler(logits, sampling_metadata)
        return next_tokens

506
507
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
Mor Zusman's avatar
Mor Zusman committed
508
509
510
511
512
513
514
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
        ]

515
516
517
518
519
520
521
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
        expert_params_mapping = 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)
Mor Zusman's avatar
Mor Zusman committed
522
523

        params_dict = dict(self.named_parameters())
524
        loaded_params: Set[str] = set()
Mor Zusman's avatar
Mor Zusman committed
525
526
527
528
529
530
531
532
533
534
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue

            if "A_log" in name:
                name = name.replace("A_log", "A")

            if ".self_attn." in name:
                name = name.replace(".self_attn", "")

535
536
537
538
            if "feed_forward" in name and not _is_moe_layer(name):
                ## map MLP layers to expert with ID=0
                name = name.replace("feed_forward", "feed_forward.experts.0")

Mor Zusman's avatar
Mor Zusman committed
539
540
541
542
543
544
545
            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.
546

Mor Zusman's avatar
Mor Zusman committed
547
548
                if name.endswith(".bias") and name not in params_dict:
                    continue
549
550
551
                # Skip layers on other devices.
                if is_pp_missing_parameter(name, self):
                    continue
Mor Zusman's avatar
Mor Zusman committed
552
553
554
555
556
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
557
558
559
560
561
562
                for (
                        param_name,
                        weight_name,
                        expert_id,
                        shard_id,
                ) in expert_params_mapping:
Mor Zusman's avatar
Mor Zusman committed
563
564
                    if weight_name not in name:
                        continue
565

566
567
                    if is_pp_missing_parameter(name, self):
                        continue
Mor Zusman's avatar
Mor Zusman committed
568
569
570
571
572
                    name = name.replace(weight_name, param_name)
                    param = params_dict[name]
                    weight_loader = param.weight_loader
                    weight_loader(param,
                                  loaded_weight,
573
                                  name,
574
                                  shard_id=shard_id,
Mor Zusman's avatar
Mor Zusman committed
575
576
577
578
579
580
                                  expert_id=expert_id)
                    break
                else:
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue
581
582
                    if is_pp_missing_parameter(name, self):
                        continue
Mor Zusman's avatar
Mor Zusman committed
583
584
585
586
587

                    param = params_dict[name]
                    weight_loader = getattr(param, "weight_loader",
                                            default_weight_loader)
                    weight_loader(param, loaded_weight)
588
589
            loaded_params.add(name)
        return loaded_params
590
591
592
593
594
595
596
597


def _is_moe_layer(name: str):
    return any(
        [experts_name in name for experts_name in [
            "experts",
            "router",
        ]])
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629


class JambaForSequenceClassification(JambaForCausalLM):

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__(vllm_config=vllm_config, prefix=prefix)
        config = vllm_config.model_config.hf_config
        num_labels: int = config.num_labels
        score_bias: bool = getattr(config, 'score_bias', False)
        self.score = nn.Linear(config.hidden_size, num_labels, bias=score_bias)

        pooler_config = vllm_config.model_config.pooler_config
        self._pooler = Pooler.from_config_with_defaults(
            pooler_config,
            pooling_type=PoolingType.LAST,
            normalize=False,
            softmax=False)

    def pooler(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> Optional[PoolerOutput]:
        hidden_states = hidden_states.float()
        logits = self.score(hidden_states)
        return self._pooler(logits, pooling_metadata)

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        # TODO: The reward weights themselves have float32 accuracy data, we
        # would like to load them in fp32 to get that extra precision.
        super().load_weights(weights)
        self.score = self.score.float()