llama.py 24.9 KB
Newer Older
1
2
# 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
3
# Copyright 2023 The vLLM team.
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# 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
22
"""Inference-only LLaMA model compatible with HuggingFace weights."""
23
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Type, Union
Woosuk Kwon's avatar
Woosuk Kwon committed
24
25
26
27
28

import torch
from torch import nn
from transformers import LlamaConfig

29
from vllm.attention import Attention, AttentionMetadata
30
from vllm.compilation.decorators import support_torch_compile
31
from vllm.config import CacheConfig, VllmConfig
32
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
33
                              get_tensor_model_parallel_world_size)
Woosuk Kwon's avatar
Woosuk Kwon committed
34
from vllm.model_executor.layers.activation import SiluAndMul
35
from vllm.model_executor.layers.layernorm import RMSNorm
36
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
37
38
                                               QKVParallelLinear,
                                               RowParallelLinear)
39
from vllm.model_executor.layers.logits_processor import LogitsProcessor
40
from vllm.model_executor.layers.quantization import QuantizationConfig
41
42
from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
    get_compressed_tensors_cache_scale)
43
from vllm.model_executor.layers.rotary_embedding import get_rope
Joe Runde's avatar
Joe Runde committed
44
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
45
from vllm.model_executor.layers.vocab_parallel_embedding import (
46
    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
47
from vllm.model_executor.model_loader.weight_utils import (
48
    default_weight_loader, kv_cache_scales_loader, maybe_remap_kv_scale_name)
49
from vllm.model_executor.sampling_metadata import SamplingMetadata
50
from vllm.platforms import current_platform
51
from vllm.sequence import IntermediateTensors
Woosuk Kwon's avatar
Woosuk Kwon committed
52

53
from .interfaces import 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
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
71
        super().__init__()
72
        self.gate_up_proj = MergedColumnParallelLinear(
73
74
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
75
            bias=bias,
76
            quant_config=quant_config,
77
78
79
80
81
82
83
84
85
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
86
87
88
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
Woosuk Kwon's avatar
Woosuk Kwon committed
89
        self.act_fn = SiluAndMul()
Woosuk Kwon's avatar
Woosuk Kwon committed
90
91

    def forward(self, x):
92
93
        x, _ = self.gate_up_proj(x)
        x = self.act_fn(x)
Woosuk Kwon's avatar
Woosuk Kwon committed
94
95
96
97
98
99
100
101
        x, _ = self.down_proj(x)
        return x


class LlamaAttention(nn.Module):

    def __init__(
        self,
102
        config: LlamaConfig,
103
104
105
106
107
108
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        rope_theta: float = 10000,
        rope_scaling: Optional[Dict[str, Any]] = None,
        max_position_embeddings: int = 8192,
109
        quant_config: Optional[QuantizationConfig] = None,
110
        bias: bool = False,
111
        cache_config: Optional[CacheConfig] = None,
112
        prefix: str = "",
113
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
114
        super().__init__()
115
        layer_idx = extract_layer_index(prefix)
116
        self.hidden_size = hidden_size
Zhuohan Li's avatar
Zhuohan Li committed
117
        tp_size = get_tensor_model_parallel_world_size()
118
        self.total_num_heads = num_heads
Zhuohan Li's avatar
Zhuohan Li committed
119
120
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
121
        self.total_num_kv_heads = num_kv_heads
122
123
124
125
126
127
128
129
130
        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)
131
132
133
        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
        self.head_dim = getattr(config, "head_dim",
                                self.hidden_size // self.total_num_heads)
Zhuohan Li's avatar
Zhuohan Li committed
134
135
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
136
        self.scaling = self.head_dim**-0.5
137
138
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
Woosuk Kwon's avatar
Woosuk Kwon committed
139

140
        self.qkv_proj = QKVParallelLinear(
141
142
143
144
            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,
145
            bias=bias,
146
            quant_config=quant_config,
147
            prefix=f"{prefix}.qkv_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
148
        )
149

150
        self.o_proj = RowParallelLinear(
151
152
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
153
            bias=bias,
154
            quant_config=quant_config,
155
            prefix=f"{prefix}.o_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
156
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
157

158
        is_neox_style = True
159
160
        is_gguf = quant_config and quant_config.get_name() == "gguf"
        if is_gguf and config.model_type == "llama":
161
162
            is_neox_style = False

163
164
165
166
167
168
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position_embeddings,
            base=rope_theta,
            rope_scaling=rope_scaling,
169
            is_neox_style=is_neox_style,
170
        )
171
172

        if hasattr(config, "interleaved_sliding_window"):
173
174
175
176
177
178
            interleaved_sliding_window = config.interleaved_sliding_window
            if isinstance(interleaved_sliding_window, int):
                sliding_window = interleaved_sliding_window
            elif isinstance(interleaved_sliding_window, list):
                sw_idx = layer_idx % len(interleaved_sliding_window)
                sliding_window = interleaved_sliding_window[sw_idx]
179
            else:
180
181
                raise ValueError(
                    f"{type(interleaved_sliding_window)} is not supported.")
182
183
184
        else:
            sliding_window = None

185
186
187
188
189
190
191
        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            quant_config=quant_config,
192
            per_layer_sliding_window=sliding_window,
193
            prefix=f"{prefix}.attn",
194
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
195
196
197

    def forward(
        self,
198
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
199
        hidden_states: torch.Tensor,
200
201
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
202
    ) -> torch.Tensor:
203
        qkv, _ = self.qkv_proj(hidden_states)
Zhuohan Li's avatar
Zhuohan Li committed
204
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
205
        q, k = self.rotary_emb(positions, q, k)
206
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
207
208
209
210
211
212
        output, _ = self.o_proj(attn_output)
        return output


class LlamaDecoderLayer(nn.Module):

213
214
215
    def __init__(
        self,
        config: LlamaConfig,
216
        cache_config: Optional[CacheConfig] = None,
217
        quant_config: Optional[QuantizationConfig] = None,
218
        prefix: str = "",
219
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
220
221
        super().__init__()
        self.hidden_size = config.hidden_size
222
223
        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
224
225
226
227
        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)
228
229
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
230
231
232
233
        # 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)
Woosuk Kwon's avatar
Woosuk Kwon committed
234
        self.self_attn = LlamaAttention(
235
            config=config,
236
237
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
238
239
            num_kv_heads=getattr(config, "num_key_value_heads",
                                 config.num_attention_heads),
240
241
242
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
243
            quant_config=quant_config,
244
            bias=attention_bias,
245
            cache_config=cache_config,
246
            prefix=f"{prefix}.self_attn",
Woosuk Kwon's avatar
Woosuk Kwon committed
247
248
        )
        self.mlp = LlamaMLP(
249
250
251
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
252
            quant_config=quant_config,
253
            bias=getattr(config, "mlp_bias", False),
254
            prefix=f"{prefix}.mlp",
Woosuk Kwon's avatar
Woosuk Kwon committed
255
        )
256
257
258
259
        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
260
261
262

    def forward(
        self,
263
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
264
        hidden_states: torch.Tensor,
265
266
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
267
268
        residual: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
269
        # Self Attention
270
271
272
273
274
275
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
276
277
278
279
        hidden_states = self.self_attn(positions=positions,
                                       hidden_states=hidden_states,
                                       kv_cache=kv_cache,
                                       attn_metadata=attn_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
280
281

        # Fully Connected
282
283
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
284
        hidden_states = self.mlp(hidden_states)
285
        return hidden_states, residual
Woosuk Kwon's avatar
Woosuk Kwon committed
286
287


288
@support_torch_compile
Woosuk Kwon's avatar
Woosuk Kwon committed
289
290
class LlamaModel(nn.Module):

291
292
293
294
295
    def __init__(self,
                 *,
                 vllm_config: VllmConfig,
                 prefix: str = "",
                 layer_type: Type[LlamaDecoderLayer] = LlamaDecoderLayer):
Woosuk Kwon's avatar
Woosuk Kwon committed
296
        super().__init__()
297
298
299
300
301
302

        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

Woosuk Kwon's avatar
Woosuk Kwon committed
303
304
        self.config = config
        self.padding_idx = config.pad_token_id
305
306
307
308
        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
309
310
311
312
313
314
        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,
315
                quant_config=quant_config,
316
317
318
            )
        else:
            self.embed_tokens = PPMissingLayer()
319
        self.start_layer, self.end_layer, self.layers = make_layers(
320
            config.num_hidden_layers,
321
322
323
324
            lambda prefix: layer_type(config=config,
                                      cache_config=cache_config,
                                      quant_config=quant_config,
                                      prefix=prefix),
325
326
            prefix=f"{prefix}.layers",
        )
327
328
329
330
        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
331

332
333
334
335
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))

336
337
338
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Woosuk Kwon's avatar
Woosuk Kwon committed
339
340
    def forward(
        self,
341
        input_ids: Optional[torch.Tensor],
342
        positions: torch.Tensor,
343
344
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
345
        intermediate_tensors: Optional[IntermediateTensors],
346
        inputs_embeds: Optional[torch.Tensor] = None,
347
348
349
350
351
352
353
    ) -> Union[torch.Tensor, IntermediateTensors]:
        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
354
        else:
355
356
357
358
359
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

        for i in range(self.start_layer, self.end_layer):
Woosuk Kwon's avatar
Woosuk Kwon committed
360
            layer = self.layers[i]
361
362
363
            hidden_states, residual = layer(positions, hidden_states,
                                            kv_caches[i - self.start_layer],
                                            attn_metadata, residual)
364
365
366
367
368
369
370

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

371
        hidden_states, _ = self.norm(hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
372
373
        return hidden_states

374
375
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
376
377
378
379
380
381
382
383
384
        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())
385
        loaded_params: Set[str] = set()
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
        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
            if scale_name := get_compressed_tensors_cache_scale(name):
                # Loading kv cache scales for compressed-tensors quantization
                param = params_dict[scale_name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                loaded_weight = loaded_weight[0]
                weight_loader(param, loaded_weight)
401
                loaded_params.add(scale_name)
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
                continue
            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
                # Remapping the name of FP8 kv-scale.
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    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)
434
435
            loaded_params.add(name)
        return loaded_params
436
437
438
439
440
441
442
443
444
445
446
447
448
449

    # If this function is called, it should always initialize KV cache scale
    # factors (or else raise an exception). Thus, handled exceptions should
    # make sure to leave KV cache scale factors in a known good (dummy) state
    def load_kv_cache_scales(self, quantization_param_path: str) -> None:
        tp_size = get_tensor_model_parallel_world_size()
        tp_rank = get_tensor_model_parallel_rank()
        for layer_idx, scaling_factor in kv_cache_scales_loader(
                quantization_param_path, tp_rank, tp_size,
                self.config.num_hidden_layers,
                self.config.__class__.model_type):
            if not isinstance(self.layers[layer_idx], nn.Identity):
                layer_self_attn = self.layers[layer_idx].self_attn

450
            if current_platform.is_rocm():
451
452
453
454
455
                # The scaling factor convention we are assuming is
                # quantized_value * scaling_factor ~= true_value
                # which is consistent with the practice of setting
                # scaling_factor = tensor_amax / FPtype_max
                scaling_factor *= 2
456
457
458
            if hasattr(layer_self_attn.attn, "_k_scale"):
                layer_self_attn.attn._k_scale = scaling_factor
                layer_self_attn.attn._v_scale = scaling_factor
459
460
461
462
            else:
                raise RuntimeError("Self attention has no KV cache scaling "
                                   "factor attribute!")

Woosuk Kwon's avatar
Woosuk Kwon committed
463

464
class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
Terry's avatar
Terry committed
465
    packed_modules_mapping = {
466
467
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"]
Terry's avatar
Terry committed
468
469
470
471
    }

    # LoRA specific attributes
    supported_lora_modules = [
472
473
        "qkv_proj", "o_proj", "gate_up_proj", "down_proj", "embed_tokens",
        "lm_head"
Terry's avatar
Terry committed
474
475
476
    ]
    embedding_modules = {
        "embed_tokens": "input_embeddings",
477
        "lm_head": "output_embeddings"
Terry's avatar
Terry committed
478
479
    }
    embedding_padding_modules = ["lm_head"]
480

481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
    # Mistral/Llama models can also be loaded with --load-format mistral
    # from consolidated.safetensors checkpoints
    mistral_mapping = {
        "layers": "model.layers",
        "attention": "self_attn",
        "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",
        "norm": "model.norm"
    }
500

501
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Woosuk Kwon's avatar
Woosuk Kwon committed
502
        super().__init__()
503
504
505
        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
506
        self.config = config
507
508
        self.lora_config = lora_config

509
510
511
        self.model = self._init_model(vllm_config=vllm_config,
                                      prefix=maybe_prefix(prefix, "model"))

512
513
514
515
516
517
518
519
        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,
520
521
522
523
524
525
                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),
526
                quant_config=quant_config,
527
                prefix=maybe_prefix(prefix, "lm_head"),
528
529
            )
            if config.tie_word_embeddings:
530
531
                self.lm_head = self.lm_head.tie_weights(
                    self.model.embed_tokens)
532
533
534
535
536
537
538

            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()
539

540
541
        self.sampler = get_sampler()

542
543
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Woosuk Kwon's avatar
Woosuk Kwon committed
544

545
546
547
    def _init_model(self, vllm_config: VllmConfig, prefix: str = ""):
        return LlamaModel(vllm_config=vllm_config, prefix=prefix)

548
549
550
    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
551
552
    def forward(
        self,
553
554
        input_ids: torch.Tensor,
        positions: torch.Tensor,
555
556
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
557
        intermediate_tensors: Optional[IntermediateTensors] = None,
558
        inputs_embeds: Optional[torch.Tensor] = None,
559
560
    ) -> Union[torch.Tensor, IntermediateTensors]:
        model_output = self.model(input_ids, positions, kv_caches,
561
562
                                  attn_metadata, intermediate_tensors,
                                  inputs_embeds)
563
        return model_output
564

565
566
567
568
569
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
570
        logits = self.logits_processor(self.lm_head, hidden_states,
571
572
573
                                       sampling_metadata)
        return logits

574
575
    def sample(self, logits: torch.Tensor,
               sampling_metadata: SamplingMetadata) -> Optional[SamplerOutput]:
576
        next_tokens = self.sampler(logits, sampling_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
577
578
        return next_tokens

579
580
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
581
582
583
584
585
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=(["lm_head."]
                           if self.config.tie_word_embeddings else None),
        )
586
        return loader.load_weights(
587
            self.maybe_remap_mistral(name, loaded_weight)
588
            for name, loaded_weight in weights)
589
590

    def load_kv_cache_scales(self, quantization_param_path: str) -> None:
591
        self.model.load_kv_cache_scales(quantization_param_path)
592
593
594
595

    # This function is used to remap the mistral format as
    # used by Mistral and Llama <=2
    def maybe_remap_mistral(
596
597
598
599
        self,
        name: str,
        loaded_weight: torch.Tensor,
    ) -> Tuple[str, torch.Tensor]:
600

601
        def permute(w: torch.Tensor, n_heads: int):
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
            attn_in = self.config.head_dim * n_heads
            attn_out = self.config.hidden_size

            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
        if "wk" in modules:
            loaded_weight = permute(loaded_weight,
                                    self.config.num_key_value_heads)
        elif "wq" in modules:
            loaded_weight = permute(loaded_weight,
                                    self.config.num_attention_heads)

        for item in modules:
            if item in mapping and mapping[item] not in name:
                name = name.replace(item, mapping[item])

        return name, loaded_weight