llama.py 26.1 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
Woosuk Kwon's avatar
Woosuk Kwon committed
6
# Copyright 2023 The vLLM team.
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Woosuk Kwon's avatar
Woosuk Kwon committed
25
"""Inference-only LLaMA model compatible with HuggingFace weights."""
26

27
from collections.abc import Iterable
28
from itertools import islice
Woosuk Kwon's avatar
Woosuk Kwon committed
29
30
31
32
33

import torch
from torch import nn
from transformers import LlamaConfig

34
from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig
37
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
Woosuk Kwon's avatar
Woosuk Kwon committed
38
from vllm.model_executor.layers.activation import SiluAndMul
39
40
41
from vllm.model_executor.layers.attention.encoder_only_attention import (
    EncoderOnlyAttention,
)
42
from vllm.model_executor.layers.layernorm import RMSNorm
43
44
45
46
47
from vllm.model_executor.layers.linear import (
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
48
from vllm.model_executor.layers.logits_processor import LogitsProcessor
49
from vllm.model_executor.layers.quantization import QuantizationConfig
50
from vllm.model_executor.layers.rotary_embedding import get_rope
51
from vllm.model_executor.layers.vocab_parallel_embedding import (
52
53
54
    ParallelLMHead,
    VocabParallelEmbedding,
)
55
from vllm.model_executor.model_loader.weight_utils import (
56
57
58
    default_weight_loader,
    maybe_remap_kv_scale_name,
)
59
from vllm.sequence import IntermediateTensors
60
from vllm.v1.attention.backend import AttentionType
Woosuk Kwon's avatar
Woosuk Kwon committed
61

62
63
64
65
66
67
68
from .adapters import as_embedding_model, as_seq_cls_model
from .interfaces import (
    SupportsEagle,
    SupportsEagle3,
    SupportsLoRA,
    SupportsPP,
)
69
70
71
72
73
74
75
76
77
from .utils import (
    AutoWeightsLoader,
    PPMissingLayer,
    extract_layer_index,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
78

Woosuk Kwon's avatar
Woosuk Kwon committed
79
80
81
82

class LlamaMLP(nn.Module):
    def __init__(
        self,
83
84
85
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
86
        quant_config: QuantizationConfig | None = None,
87
        bias: bool = False,
88
        prefix: str = "",
89
        reduce_results: bool = True,
90
        disable_tp: bool = False,
91
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
92
        super().__init__()
93
        self.gate_up_proj = MergedColumnParallelLinear(
94
95
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
96
            bias=bias,
97
            quant_config=quant_config,
98
            disable_tp=disable_tp,
99
100
101
102
103
104
105
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=bias,
            quant_config=quant_config,
106
            reduce_results=reduce_results,
107
            disable_tp=disable_tp,
108
109
            prefix=f"{prefix}.down_proj",
        )
110
        if hidden_act != "silu":
111
112
113
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
114
        self.act_fn = SiluAndMul()
Woosuk Kwon's avatar
Woosuk Kwon committed
115
116

    def forward(self, x):
117
118
        x, _ = self.gate_up_proj(x)
        x = self.act_fn(x)
Woosuk Kwon's avatar
Woosuk Kwon committed
119
120
121
122
123
        x, _ = self.down_proj(x)
        return x


class LlamaAttention(nn.Module):
124
125
126
127
128
129
130
    def __init__(
        self,
        config: LlamaConfig,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position_embeddings: int = 8192,
131
        quant_config: QuantizationConfig | None = None,
132
133
        bias: bool = False,
        bias_o_proj: bool = False,
134
        cache_config: CacheConfig | None = None,
135
136
137
        prefix: str = "",
        attn_type: str = AttentionType.DECODER,
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
138
        super().__init__()
139
        layer_idx = extract_layer_index(prefix)
140
        self.hidden_size = hidden_size
Zhuohan Li's avatar
Zhuohan Li committed
141
        tp_size = get_tensor_model_parallel_world_size()
142
        self.total_num_heads = num_heads
Zhuohan Li's avatar
Zhuohan Li committed
143
144
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
145
        self.total_num_kv_heads = num_kv_heads
146
147
148
149
150
151
152
153
154
        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)
155
        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
156
157
158
159
        head_dim = getattr(config, "head_dim", None)
        if head_dim is None:
            head_dim = self.hidden_size // self.total_num_heads
        self.head_dim = head_dim
Zhuohan Li's avatar
Zhuohan Li committed
160
161
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
162
        self.scaling = self.head_dim**-0.5
163
        self.max_position_embeddings = max_position_embeddings
Woosuk Kwon's avatar
Woosuk Kwon committed
164

165
166
167
168
169
170
171
172
        llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
        self.do_llama_4_scaling = llama_4_scaling_config is not None
        if self.do_llama_4_scaling:
            self.llama_4_scaling_original_max_position_embeddings = (
                llama_4_scaling_config["original_max_position_embeddings"]
            )
            self.llama_4_scaling_beta = llama_4_scaling_config["beta"]

173
        self.qkv_proj = QKVParallelLinear(
174
175
176
177
            hidden_size=hidden_size,
            head_size=self.head_dim,
            total_num_heads=self.total_num_heads,
            total_num_kv_heads=self.total_num_kv_heads,
178
            bias=bias,
179
            quant_config=quant_config,
180
            prefix=f"{prefix}.qkv_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
181
        )
182

183
        self.o_proj = RowParallelLinear(
184
185
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
186
            bias=bias_o_proj,
187
            quant_config=quant_config,
188
            prefix=f"{prefix}.o_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
189
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
190

191
        self._init_rotary_emb(config, quant_config=quant_config)
192

193
194
        sliding_window = None
        if layer_types := getattr(config, "layer_types", None):
195
196
197
            # Fix for Eagle3 compatibility:
            # for draft models, subtract target layer count
            # to get draft-relative layer index starting from 0
198
            if hasattr(config, "target_layer_count"):
199
200
201
202
203
204
                # This is a draft model,
                # adjust layer_idx to be relative to draft layers
                effective_layer_idx = layer_idx - config.target_layer_count
            else:
                # This is a target model, use layer_idx directly
                effective_layer_idx = layer_idx
205
            assert effective_layer_idx < len(layer_types), (
206
207
                f"effective_layer_idx: {effective_layer_idx} "
                f"is out of bounds for layer_types: {layer_types}"
208
            )
209

210
            is_sliding = layer_types[effective_layer_idx] == "sliding_attention"
211
212
            if is_sliding:
                sliding_window = config.sliding_window
213

214
215
216
217
218
        attn_cls = (
            EncoderOnlyAttention
            if attn_type == AttentionType.ENCODER_ONLY
            else Attention
        )
219
220

        self.attn = attn_cls(
221
222
223
224
225
226
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            quant_config=quant_config,
227
            per_layer_sliding_window=sliding_window,
228
            attn_type=attn_type,
229
            prefix=f"{prefix}.attn",
230
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
231

232
233
234
235
236
237
238
239
240
241
242
    def _get_llama_4_attn_scale(self, positions: torch.Tensor) -> torch.Tensor:
        # Llama4 scaling
        scaling = 1 + self.llama_4_scaling_beta * torch.log(
            1
            + torch.floor(
                positions / self.llama_4_scaling_original_max_position_embeddings
            )
        )
        # Broadcast over head_dim
        return scaling.unsqueeze(-1)

Woosuk Kwon's avatar
Woosuk Kwon committed
243
244
    def forward(
        self,
245
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
246
247
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
248
        qkv, _ = self.qkv_proj(hidden_states)
Zhuohan Li's avatar
Zhuohan Li committed
249
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
250
        q, k = self.rotary_emb(positions, q, k)
251
252
253
        if self.do_llama_4_scaling:
            attn_scale = self._get_llama_4_attn_scale(positions)
            q = (q * attn_scale).to(q.dtype)
254
        attn_output = self.attn(q, k, v)
Woosuk Kwon's avatar
Woosuk Kwon committed
255
256
257
        output, _ = self.o_proj(attn_output)
        return output

258
259
260
    def _init_rotary_emb(
        self,
        config: LlamaConfig,
261
        quant_config: QuantizationConfig | None,
262
    ) -> None:
263
264
        is_neox_style = True
        is_gguf = quant_config and quant_config.get_name() == "gguf"
265
        if is_gguf and config.model_type == "llama":
266
267
268
269
270
            is_neox_style = False

        self.rotary_emb = get_rope(
            self.head_dim,
            max_position=self.max_position_embeddings,
271
            rope_parameters=getattr(config, "rope_parameters", None),
272
273
274
            is_neox_style=is_neox_style,
        )

Woosuk Kwon's avatar
Woosuk Kwon committed
275
276

class LlamaDecoderLayer(nn.Module):
277
278
279
280
    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
281
        config: LlamaConfig | None = None,
282
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
283
        super().__init__()
284
285
286

        config = config or vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
287
        quant_config = self.get_quant_config(vllm_config)
288

Woosuk Kwon's avatar
Woosuk Kwon committed
289
        self.hidden_size = config.hidden_size
290
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
291
292
293
        # Support abacusai/Smaug-72B-v0.1 with attention_bias
        # Support internlm/internlm-7b with bias
        attention_bias = getattr(config, "attention_bias", False) or getattr(
294
295
            config, "bias", False
        )
296
297
        bias_o_proj = attention_bias
        # support internlm/internlm3-8b with qkv_bias
298
        if hasattr(config, "qkv_bias"):
299
300
            attention_bias = config.qkv_bias

301
302
303
304
305
306
307
308
309
        # By default, Llama uses causal attention as it is a decoder-only model.
        # You can override the HF config with `is_causal=False` to enable
        # bidirectional attention, which is used in some embedding models
        # (e.g. parasail-ai/GritLM-7B-vllm)
        if getattr(config, "is_causal", True):
            attn_type = AttentionType.DECODER
        else:
            attn_type = AttentionType.ENCODER_ONLY

Woosuk Kwon's avatar
Woosuk Kwon committed
310
        self.self_attn = LlamaAttention(
311
            config=config,
312
313
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
314
315
316
            num_kv_heads=getattr(
                config, "num_key_value_heads", config.num_attention_heads
            ),
317
            max_position_embeddings=max_position_embeddings,
318
            quant_config=quant_config,
319
            bias=attention_bias,
320
            bias_o_proj=bias_o_proj,
321
            cache_config=cache_config,
322
            prefix=f"{prefix}.self_attn",
323
            attn_type=attn_type,
Woosuk Kwon's avatar
Woosuk Kwon committed
324
325
        )
        self.mlp = LlamaMLP(
326
327
328
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
329
            quant_config=quant_config,
330
            bias=getattr(config, "mlp_bias", False),
331
            prefix=f"{prefix}.mlp",
Woosuk Kwon's avatar
Woosuk Kwon committed
332
        )
333
334
335
336
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
337
338
339

    def forward(
        self,
340
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
341
        hidden_states: torch.Tensor,
342
        residual: torch.Tensor | None,
343
    ) -> tuple[torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
344
        # Self Attention
345
346
347
348
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
349
350
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
Woosuk Kwon's avatar
Woosuk Kwon committed
351
352

        # Fully Connected
353
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
354
        hidden_states = self.mlp(hidden_states)
355
        return hidden_states, residual
Woosuk Kwon's avatar
Woosuk Kwon committed
356

357
    def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None:
358
359
360
        """Get quantization config for this layer. Override in subclasses."""
        return vllm_config.quant_config

Woosuk Kwon's avatar
Woosuk Kwon committed
361

362
363
364
365
366
367
368
369
370
371
def llama_model_invariants(
    input_ids, positions, intermediate_tensors=None, inputs_embeds=None
):
    """Shape invariants for Llama model compilation, those are translated to
    runtime assertions for unbacked dynamic shapes and are compiled away for
    backed"""
    if input_ids is not None:
        torch._check(positions.size()[0] == input_ids.size()[0])


372
373
374
375
376
@support_torch_compile(
    # TODO[#32068]: Investigate recompilation
    # mark_unbacked_dims={"input_ids": 0},
    shape_invariants=llama_model_invariants
)
Woosuk Kwon's avatar
Woosuk Kwon committed
377
class LlamaModel(nn.Module):
378
379
380
381
382
383
384
    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        layer_type: type[nn.Module] = LlamaDecoderLayer,
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
385
        super().__init__()
386
387
388
389

        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config

Woosuk Kwon's avatar
Woosuk Kwon committed
390
        self.config = config
391
        self.quant_config = quant_config
392
393
394

        self.vocab_size = config.vocab_size

395
396
397
        if get_pp_group().is_first_rank or (
            config.tie_word_embeddings and get_pp_group().is_last_rank
        ):
398
399
400
            self.embed_tokens = VocabParallelEmbedding(
                self.vocab_size,
                config.hidden_size,
401
                quant_config=quant_config,
402
403
404
            )
        else:
            self.embed_tokens = PPMissingLayer()
405
        self.start_layer, self.end_layer, self.layers = make_layers(
406
            config.num_hidden_layers,
407
            lambda prefix: layer_type(vllm_config=vllm_config, prefix=prefix),
408
409
            prefix=f"{prefix}.layers",
        )
410
411
412
413
        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()
Woosuk Kwon's avatar
Woosuk Kwon committed
414

415
        self.aux_hidden_state_layers = tuple[int, ...]()
416

417
418
419
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
420

421
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
422
423
        return self.embed_tokens(input_ids)

Woosuk Kwon's avatar
Woosuk Kwon committed
424
425
    def forward(
        self,
426
        input_ids: torch.Tensor | None,
427
        positions: torch.Tensor,
428
429
430
        intermediate_tensors: IntermediateTensors | None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
431
432
433
434
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
435
                hidden_states = self.embed_input_ids(input_ids)
436
            residual = None
437
        else:
438
439
440
441
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

442
443
        aux_hidden_states = []
        for idx, layer in enumerate(
444
445
            islice(self.layers, self.start_layer, self.end_layer)
        ):
446
447
            if idx in self.aux_hidden_state_layers:
                aux_hidden_states.append(hidden_states + residual)
448
            hidden_states, residual = layer(positions, hidden_states, residual)
449
450

        if not get_pp_group().is_last_rank:
451
452
453
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
454

455
        hidden_states, _ = self.norm(hidden_states, residual)
456
457
458

        if len(aux_hidden_states) > 0:
            return hidden_states, aux_hidden_states
Woosuk Kwon's avatar
Woosuk Kwon committed
459
460
        return hidden_states

461
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
462
463
464
465
466
467
468
469
470
        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())
471
        loaded_params: set[str] = set()
472
473
474
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
475
            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
476
477
478
                # Models trained using ColossalAI may include these tensors in
                # the checkpoint. Skip them.
                continue
479
480
481
            if self.quant_config is not None and (
                scale_name := self.quant_config.get_cache_scale(name)
            ):
482
                # Loading kv cache quantization scales
483
                param = params_dict[scale_name]
484
485
486
487
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                loaded_weight = (
                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
                )
488
                weight_loader(param, loaded_weight)
489
                loaded_params.add(scale_name)
490
                continue
491
492
493
494
495
            if "scale" in name:
                # Remapping the name of FP8 kv-scale.
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not 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

                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:
                # 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]
520
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
521
                weight_loader(param, loaded_weight)
522
523
            loaded_params.add(name)
        return loaded_params
524

Woosuk Kwon's avatar
Woosuk Kwon committed
525

526
527
528
class LlamaForCausalLM(
    nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
):
Terry's avatar
Terry committed
529
    packed_modules_mapping = {
530
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
531
        "gate_up_proj": ["gate_proj", "up_proj"],
Terry's avatar
Terry committed
532
533
534
535
536
    }

    # LoRA specific attributes
    embedding_modules = {
        "embed_tokens": "input_embeddings",
537
        "lm_head": "output_embeddings",
Terry's avatar
Terry committed
538
    }
539

540
541
542
543
544
    # Mistral/Llama models can also be loaded with --load-format mistral
    # from consolidated.safetensors checkpoints
    mistral_mapping = {
        "layers": "model.layers",
        "attention": "self_attn",
545
546
547
        "qscale_act": "input_scale",
        "qscale_weight": "weight_scale",
        "kv_fake_quantizer.qscale_act": "kv_scale",
548
549
550
        "q_fake_quantizer.qscale_act": "attn.q_scale",
        "k_fake_quantizer.qscale_act": "k_scale",
        "v_fake_quantizer.qscale_act": "v_scale",
551
552
553
554
555
556
557
558
559
560
561
562
        "wq": "q_proj",
        "wk": "k_proj",
        "wv": "v_proj",
        "wo": "o_proj",
        "attention_norm": "input_layernorm",
        "feed_forward": "mlp",
        "w1": "gate_proj",
        "w2": "down_proj",
        "w3": "up_proj",
        "ffn_norm": "post_attention_layernorm",
        "tok_embeddings": "model.embed_tokens",
        "output": "lm_head",
563
        "norm": "model.norm",
564
    }
565

566
567
568
569
570
571
572
    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        layer_type: type[nn.Module] = LlamaDecoderLayer,
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
573
        super().__init__()
574
575
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
Woosuk Kwon's avatar
Woosuk Kwon committed
576
        self.config = config
577

578
579
580
581
582
        self.model = self._init_model(
            vllm_config=vllm_config,
            prefix=maybe_prefix(prefix, "model"),
            layer_type=layer_type,
        )
583

584
585
        if get_pp_group().is_last_rank:
            self.lm_head = ParallelLMHead(
586
                config.vocab_size,
587
588
                config.hidden_size,
                quant_config=quant_config,
589
                prefix=maybe_prefix(prefix, "lm_head"),
590
591
            )
            if config.tie_word_embeddings:
592
                self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
593
594

            logit_scale = getattr(config, "logit_scale", 1.0)
595
            self.logits_processor = LogitsProcessor(
596
                config.vocab_size, scale=logit_scale
597
            )
598
599
        else:
            self.lm_head = PPMissingLayer()
600

601
        self.make_empty_intermediate_tensors = (
602
603
            self.model.make_empty_intermediate_tensors
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
604

605
    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
606
607
        self.model.aux_hidden_state_layers = layers

608
    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
609
610
611
612
613
        """Override to return default layers for Llama

        Note: The GPU model runner will override this with layers from
        the speculative config if available, providing dynamic configuration.
        """
614
615
616
        num_layers = len(self.model.layers)
        return (2, num_layers // 2, num_layers - 3)

617
618
619
620
621
622
623
    def _init_model(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
        layer_type: type[nn.Module] = LlamaDecoderLayer,
    ):
        return LlamaModel(vllm_config=vllm_config, prefix=prefix, layer_type=layer_type)
624

625
626
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
627

Woosuk Kwon's avatar
Woosuk Kwon committed
628
629
    def forward(
        self,
630
631
        input_ids: torch.Tensor,
        positions: torch.Tensor,
632
633
634
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
635
636
637
        model_output = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
638
        return model_output
639

640
641
642
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
643
    ) -> torch.Tensor | None:
644
        logits = self.logits_processor(self.lm_head, hidden_states)
645
646
        return logits

647
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
648
649
        loader = AutoWeightsLoader(
            self,
650
            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
651
        )
652
        return loader.load_weights(
653
            self.maybe_remap_mistral(name, loaded_weight)
654
655
            for name, loaded_weight in weights
        )
656

657
658
659
    # This function is used to remap the mistral format as
    # used by Mistral and Llama <=2
    def maybe_remap_mistral(
660
661
662
        self,
        name: str,
        loaded_weight: torch.Tensor,
663
    ) -> tuple[str, torch.Tensor]:
664
        def permute(w: torch.Tensor, n_heads: int, attn_out: int):
665
666
            attn_in = self.config.head_dim * n_heads

667
668
669
670
671
            return (
                w.view(n_heads, attn_in // n_heads // 2, 2, attn_out)
                .transpose(1, 2)
                .reshape(attn_in, attn_out)
            )
672
673
674
675
676

        mapping = self.mistral_mapping
        modules = name.split(".")

        # rotary embeds should be sliced
677
678
        # If using quantized model in mistral format,
        # quantization scales (qscale_weight) also need to be sliced
679
        if "wk" in modules and modules[-1] == "weight":
680
681
682
683
684
685
686
687
688
            loaded_weight = permute(
                loaded_weight, self.config.num_key_value_heads, self.config.hidden_size
            )
        elif (
            "wk" in modules
            and modules[-1] == "qscale_weight"
            and loaded_weight.numel() > 1
        ):
            loaded_weight = permute(loaded_weight, self.config.num_key_value_heads, 1)
689
        elif "wq" in modules and modules[-1] == "weight":
690
691
692
693
694
695
696
697
698
            loaded_weight = permute(
                loaded_weight, self.config.num_attention_heads, self.config.hidden_size
            )
        elif (
            "wq" in modules
            and modules[-1] == "qscale_weight"
            and loaded_weight.numel() > 1
        ):
            loaded_weight = permute(loaded_weight, self.config.num_attention_heads, 1)
699

700
701
702
703
704
        num_modules = len(modules)
        for i in range(num_modules):
            item = modules[i]
            next_item = modules[i + 1] if i < num_modules - 1 else None

705
            combined_item = f"{item}.{next_item}" if next_item is not None else None
706
707
708
709

            if combined_item in mapping:
                name = name.replace(combined_item, mapping[combined_item])
            elif item in mapping and mapping[item] not in name:
710
711
712
                name = name.replace(item, mapping[item])

        return name, loaded_weight
713
714
715
716
717
718
719
720
721
722
723
724


class LlamaBidirectionalForSequenceClassification(as_seq_cls_model(LlamaForCausalLM)):
    # This class sets the correct attention type and pooling type
    # through LlamaBidirectionalConfig.
    pass


class LlamaBidirectionalModel(as_embedding_model(LlamaForCausalLM)):
    # This class sets the correct attention type and pooling type
    # through LlamaBidirectionalConfig.
    pass