llama.py 26.6 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
from collections.abc import Iterable
27
from itertools import islice
28
from typing import Any, Optional, Union
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 import Attention, AttentionType
35
from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
36
from vllm.compilation.decorators import support_torch_compile
37
from vllm.config import CacheConfig, VllmConfig
38
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
Woosuk Kwon's avatar
Woosuk Kwon committed
39
from vllm.model_executor.layers.activation import SiluAndMul
40
from vllm.model_executor.layers.layernorm import RMSNorm
41
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
42
43
                                               QKVParallelLinear,
                                               RowParallelLinear)
44
from vllm.model_executor.layers.logits_processor import LogitsProcessor
45
from vllm.model_executor.layers.quantization import QuantizationConfig
46
from vllm.model_executor.layers.rotary_embedding import get_rope
47
from vllm.model_executor.layers.vocab_parallel_embedding import (
48
    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
49
from vllm.model_executor.model_loader.weight_utils import (
50
    default_weight_loader, maybe_remap_kv_scale_name)
51
from vllm.sequence import IntermediateTensors
Woosuk Kwon's avatar
Woosuk Kwon committed
52

53
from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
54
55
from .utils import (AutoWeightsLoader, PPMissingLayer, extract_layer_index,
                    is_pp_missing_parameter,
56
57
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
58

Woosuk Kwon's avatar
Woosuk Kwon committed
59
60

class LlamaMLP(nn.Module):
61

Woosuk Kwon's avatar
Woosuk Kwon committed
62
63
    def __init__(
        self,
64
65
66
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
67
        quant_config: Optional[QuantizationConfig] = None,
68
        bias: bool = False,
69
        prefix: str = "",
70
        reduce_results: bool = True,
71
        disable_tp: bool = False,
72
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
73
        super().__init__()
74
        self.gate_up_proj = MergedColumnParallelLinear(
75
76
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
77
            bias=bias,
78
            quant_config=quant_config,
79
            disable_tp=disable_tp,
80
81
82
83
84
85
86
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=bias,
            quant_config=quant_config,
87
            reduce_results=reduce_results,
88
            disable_tp=disable_tp,
89
90
            prefix=f"{prefix}.down_proj",
        )
91
92
93
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
Woosuk Kwon's avatar
Woosuk Kwon committed
94
        self.act_fn = SiluAndMul()
Woosuk Kwon's avatar
Woosuk Kwon committed
95
96

    def forward(self, x):
97
98
        x, _ = self.gate_up_proj(x)
        x = self.act_fn(x)
Woosuk Kwon's avatar
Woosuk Kwon committed
99
100
101
102
103
104
        x, _ = self.down_proj(x)
        return x


class LlamaAttention(nn.Module):

105
106
107
108
109
110
111
    def __init__(
        self,
        config: LlamaConfig,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        rope_theta: float = 10000,
112
        rope_scaling: Optional[dict[str, Any]] = None,
113
114
115
116
117
118
119
120
        max_position_embeddings: int = 8192,
        quant_config: Optional[QuantizationConfig] = None,
        bias: bool = False,
        bias_o_proj: bool = False,
        cache_config: Optional[CacheConfig] = None,
        prefix: str = "",
        attn_type: str = AttentionType.DECODER,
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
121
        super().__init__()
122
        layer_idx = extract_layer_index(prefix)
123
        self.hidden_size = hidden_size
Zhuohan Li's avatar
Zhuohan Li committed
124
        tp_size = get_tensor_model_parallel_world_size()
125
        self.total_num_heads = num_heads
Zhuohan Li's avatar
Zhuohan Li committed
126
127
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
128
        self.total_num_kv_heads = num_kv_heads
129
130
131
132
133
134
135
136
137
        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)
138
        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
139
140
141
142
        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
Amit Garg's avatar
Amit Garg committed
143
        # Phi models introduced a partial_rotary_factor parameter in the config
144
145
        self.partial_rotary_factor = getattr(config, "partial_rotary_factor",
                                             1)
Zhuohan Li's avatar
Zhuohan Li committed
146
147
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
148
        self.scaling = self.head_dim**-0.5
149
150
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
Woosuk Kwon's avatar
Woosuk Kwon committed
151

152
        self.qkv_proj = QKVParallelLinear(
153
154
155
156
            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,
157
            bias=bias,
158
            quant_config=quant_config,
159
            prefix=f"{prefix}.qkv_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
160
        )
161

162
        self.o_proj = RowParallelLinear(
163
164
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
165
            bias=bias_o_proj,
166
            quant_config=quant_config,
167
            prefix=f"{prefix}.o_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
168
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
169

170
171
172
        self._init_rotary_emb(config,
                              rope_scaling=rope_scaling,
                              quant_config=quant_config)
173

174
175
        sliding_window = None
        if layer_types := getattr(config, "layer_types", None):
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
            # Fix for Eagle3 compatibility:
            # for draft models, subtract target layer count
            # to get draft-relative layer index starting from 0
            if hasattr(config, 'target_layer_count'):
                # 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
            assert effective_layer_idx < len(layer_types), \
                f"effective_layer_idx: {effective_layer_idx} \
                is out of bounds for layer_types: {layer_types}"

            is_sliding = layer_types[
                effective_layer_idx] == "sliding_attention"
192
193
            if is_sliding:
                sliding_window = config.sliding_window
194

195
196
197
198
        attn_cls = (EncoderOnlyAttention
                    if attn_type == AttentionType.ENCODER_ONLY else Attention)

        self.attn = attn_cls(
199
200
201
202
203
204
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            quant_config=quant_config,
205
            per_layer_sliding_window=sliding_window,
206
            attn_type=attn_type,
207
            prefix=f"{prefix}.attn",
208
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
209
210
211

    def forward(
        self,
212
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
213
214
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
215
        qkv, _ = self.qkv_proj(hidden_states)
Zhuohan Li's avatar
Zhuohan Li committed
216
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
217
        q, k = self.rotary_emb(positions, q, k)
218
        attn_output = self.attn(q, k, v)
Woosuk Kwon's avatar
Woosuk Kwon committed
219
220
221
        output, _ = self.o_proj(attn_output)
        return output

222
223
224
225
226
    def _init_rotary_emb(self, config: LlamaConfig,
                         rope_scaling: Optional[dict[str, Any]],
                         quant_config: Optional[QuantizationConfig]) -> None:
        is_neox_style = True
        is_gguf = quant_config and quant_config.get_name() == "gguf"
227
        if is_gguf and config.model_type == "llama":
228
229
230
231
232
233
234
235
236
237
238
239
            is_neox_style = False

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=self.max_position_embeddings,
            base=self.rope_theta,
            rope_scaling=rope_scaling,
            is_neox_style=is_neox_style,
            partial_rotary_factor=self.partial_rotary_factor,
        )

Woosuk Kwon's avatar
Woosuk Kwon committed
240
241
242

class LlamaDecoderLayer(nn.Module):

243
244
245
246
    def __init__(self,
                 vllm_config: VllmConfig,
                 prefix: str = "",
                 config: Optional[LlamaConfig] = None) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
247
        super().__init__()
248
249
250
251
252

        config = config or vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

Woosuk Kwon's avatar
Woosuk Kwon committed
253
        self.hidden_size = config.hidden_size
254
255
        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
256
257
258
259
        if rope_scaling is not None and getattr(
                config, "original_max_position_embeddings", None):
            rope_scaling["original_max_position_embeddings"] = (
                config.original_max_position_embeddings)
260
261
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
262
263
264
265
        # Support abacusai/Smaug-72B-v0.1 with attention_bias
        # Support internlm/internlm-7b with bias
        attention_bias = getattr(config, "attention_bias", False) or getattr(
            config, "bias", False)
266
267
268
269
270
        bias_o_proj = attention_bias
        # support internlm/internlm3-8b with qkv_bias
        if hasattr(config, 'qkv_bias'):
            attention_bias = config.qkv_bias

271
272
273
274
275
276
277
278
279
        # 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
280
        self.self_attn = LlamaAttention(
281
            config=config,
282
283
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
284
285
            num_kv_heads=getattr(config, "num_key_value_heads",
                                 config.num_attention_heads),
286
287
288
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
289
            quant_config=quant_config,
290
            bias=attention_bias,
291
            bias_o_proj=bias_o_proj,
292
            cache_config=cache_config,
293
            prefix=f"{prefix}.self_attn",
294
            attn_type=attn_type,
Woosuk Kwon's avatar
Woosuk Kwon committed
295
296
        )
        self.mlp = LlamaMLP(
297
298
299
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
300
            quant_config=quant_config,
301
            bias=getattr(config, "mlp_bias", False),
302
            prefix=f"{prefix}.mlp",
Woosuk Kwon's avatar
Woosuk Kwon committed
303
        )
304
305
306
307
        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
308
309
310

    def forward(
        self,
311
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
312
        hidden_states: torch.Tensor,
313
        residual: Optional[torch.Tensor],
314
    ) -> tuple[torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
315
        # Self Attention
316
317
318
319
320
321
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
322
        hidden_states = self.self_attn(positions=positions,
323
                                       hidden_states=hidden_states)
Woosuk Kwon's avatar
Woosuk Kwon committed
324
325

        # Fully Connected
326
327
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
328
        hidden_states = self.mlp(hidden_states)
329
        return hidden_states, residual
Woosuk Kwon's avatar
Woosuk Kwon committed
330
331


332
@support_torch_compile
Woosuk Kwon's avatar
Woosuk Kwon committed
333
334
class LlamaModel(nn.Module):

335
336
337
338
    def __init__(self,
                 *,
                 vllm_config: VllmConfig,
                 prefix: str = "",
339
                 layer_type: type[nn.Module] = LlamaDecoderLayer):
Woosuk Kwon's avatar
Woosuk Kwon committed
340
        super().__init__()
341
342
343
344
345

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

Woosuk Kwon's avatar
Woosuk Kwon committed
346
        self.config = config
347
        self.quant_config = quant_config
348
349
350
351
        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
352
353
354
355
356
357
        if get_pp_group().is_first_rank or (config.tie_word_embeddings
                                            and get_pp_group().is_last_rank):
            self.embed_tokens = VocabParallelEmbedding(
                self.vocab_size,
                config.hidden_size,
                org_num_embeddings=config.vocab_size,
358
                quant_config=quant_config,
359
360
361
            )
        else:
            self.embed_tokens = PPMissingLayer()
362
        self.start_layer, self.end_layer, self.layers = make_layers(
363
            config.num_hidden_layers,
364
            lambda prefix: layer_type(vllm_config=vllm_config, prefix=prefix),
365
366
            prefix=f"{prefix}.layers",
        )
367
368
369
370
        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
371

372
        self.aux_hidden_state_layers = tuple[int, ...]()
373

374
375
376
377
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))

378
379
380
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Woosuk Kwon's avatar
Woosuk Kwon committed
381
382
    def forward(
        self,
383
        input_ids: Optional[torch.Tensor],
384
        positions: torch.Tensor,
385
        intermediate_tensors: Optional[IntermediateTensors],
386
        inputs_embeds: Optional[torch.Tensor] = None,
387
388
    ) -> Union[torch.Tensor, IntermediateTensors, tuple[torch.Tensor,
                                                        list[torch.Tensor]]]:
389
390
391
392
393
394
        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
395
        else:
396
397
398
399
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

400
401
        aux_hidden_states = []
        for idx, layer in enumerate(
402
                islice(self.layers, self.start_layer, self.end_layer)):
403
404
            if idx in self.aux_hidden_state_layers:
                aux_hidden_states.append(hidden_states + residual)
405
            hidden_states, residual = layer(positions, hidden_states, residual)
406
407
408
409
410
411
412

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

413
        hidden_states, _ = self.norm(hidden_states, residual)
414
415
416

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

419
420
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
421
422
423
424
425
426
427
428
429
        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())
430
        loaded_params: set[str] = set()
431
432
433
434
435
436
437
438
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
            if ("rotary_emb.cos_cached" in name
                    or "rotary_emb.sin_cached" in name):
                # Models trained using ColossalAI may include these tensors in
                # the checkpoint. Skip them.
                continue
439
440
            if (self.quant_config is not None and
                (scale_name := self.quant_config.get_cache_scale(name))):
441
                # Loading kv cache quantization scales
442
443
444
                param = params_dict[scale_name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
445
446
                loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
                                 loaded_weight[0])
447
                weight_loader(param, loaded_weight)
448
                loaded_params.add(scale_name)
449
                continue
450
451
452
453
454
            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
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
            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]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
482
483
            loaded_params.add(name)
        return loaded_params
484

Woosuk Kwon's avatar
Woosuk Kwon committed
485

486
class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
Terry's avatar
Terry committed
487
    packed_modules_mapping = {
488
489
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"]
Terry's avatar
Terry committed
490
491
492
493
494
    }

    # LoRA specific attributes
    embedding_modules = {
        "embed_tokens": "input_embeddings",
495
        "lm_head": "output_embeddings"
Terry's avatar
Terry committed
496
497
    }
    embedding_padding_modules = ["lm_head"]
498

499
500
501
502
503
    # Mistral/Llama models can also be loaded with --load-format mistral
    # from consolidated.safetensors checkpoints
    mistral_mapping = {
        "layers": "model.layers",
        "attention": "self_attn",
504
505
506
        "qscale_act": "input_scale",
        "qscale_weight": "weight_scale",
        "kv_fake_quantizer.qscale_act": "kv_scale",
507
508
509
        "q_fake_quantizer.qscale_act": "attn.q_scale",
        "k_fake_quantizer.qscale_act": "k_scale",
        "v_fake_quantizer.qscale_act": "v_scale",
510
511
512
513
514
515
516
517
518
519
520
521
        "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",
522
        "norm": "model.norm",
523
    }
524

525
526
527
528
529
    def __init__(self,
                 *,
                 vllm_config: VllmConfig,
                 prefix: str = "",
                 layer_type: type[nn.Module] = LlamaDecoderLayer):
Woosuk Kwon's avatar
Woosuk Kwon committed
530
        super().__init__()
531
532
533
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config
Woosuk Kwon's avatar
Woosuk Kwon committed
534
        self.config = config
535
536
        self.lora_config = lora_config

537
        self.model = self._init_model(vllm_config=vllm_config,
538
539
                                      prefix=maybe_prefix(prefix, "model"),
                                      layer_type=layer_type)
540

541
542
543
544
545
546
547
548
        if get_pp_group().is_last_rank:
            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,
549
550
551
552
553
554
                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),
555
                quant_config=quant_config,
556
                prefix=maybe_prefix(prefix, "lm_head"),
557
558
            )
            if config.tie_word_embeddings:
559
560
                self.lm_head = self.lm_head.tie_weights(
                    self.model.embed_tokens)
561
562
563
564
565
566
567

            logit_scale = getattr(config, "logit_scale", 1.0)
            self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                    config.vocab_size,
                                                    logit_scale)
        else:
            self.lm_head = PPMissingLayer()
568

569
570
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Woosuk Kwon's avatar
Woosuk Kwon committed
571

572
    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
573
574
        self.model.aux_hidden_state_layers = layers

575
    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
576
577
578
        num_layers = len(self.model.layers)
        return (2, num_layers // 2, num_layers - 3)

579
580
581
582
583
584
585
    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)
586

587
588
589
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

Woosuk Kwon's avatar
Woosuk Kwon committed
590
591
    def forward(
        self,
592
593
        input_ids: torch.Tensor,
        positions: torch.Tensor,
594
        intermediate_tensors: Optional[IntermediateTensors] = None,
595
        inputs_embeds: Optional[torch.Tensor] = None,
596
    ) -> Union[torch.Tensor, IntermediateTensors]:
597
        model_output = self.model(input_ids, positions, intermediate_tensors,
598
                                  inputs_embeds)
599
        return model_output
600

601
602
603
604
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> Optional[torch.Tensor]:
605
        logits = self.logits_processor(self.lm_head, hidden_states)
606
607
        return logits

608
609
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
610
611
612
613
614
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=(["lm_head."]
                           if self.config.tie_word_embeddings else None),
        )
615
        return loader.load_weights(
616
            self.maybe_remap_mistral(name, loaded_weight)
617
            for name, loaded_weight in weights)
618

619
620
621
    # This function is used to remap the mistral format as
    # used by Mistral and Llama <=2
    def maybe_remap_mistral(
622
623
624
        self,
        name: str,
        loaded_weight: torch.Tensor,
625
    ) -> tuple[str, torch.Tensor]:
626

627
        def permute(w: torch.Tensor, n_heads: int, attn_out: int):
628
629
630
631
632
633
634
635
636
            attn_in = self.config.head_dim * n_heads

            return w.view(n_heads, attn_in // n_heads // 2, 2,
                          attn_out).transpose(1, 2).reshape(attn_in, attn_out)

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

        # rotary embeds should be sliced
637
638
        # If using quantized model in mistral format,
        # quantization scales (qscale_weight) also need to be sliced
639
        if "wk" in modules and modules[-1] == "weight":
640
            loaded_weight = permute(loaded_weight,
641
642
643
644
645
646
                                    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)
647
        elif "wq" in modules and modules[-1] == "weight":
648
            loaded_weight = permute(loaded_weight,
649
650
651
652
653
654
                                    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)
655

656
657
658
659
660
661
662
663
664
665
666
        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

            combined_item = (f"{item}.{next_item}"
                             if next_item is not None else None)

            if combined_item in mapping:
                name = name.replace(combined_item, mapping[combined_item])
            elif item in mapping and mapping[item] not in name:
667
668
669
                name = name.replace(item, mapping[item])

        return name, loaded_weight