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

import torch
from torch import nn
from transformers import LlamaConfig

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

54
from .interfaces import SupportsLoRA, SupportsPP
55
from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
56
                    make_empty_intermediate_tensors_factory, make_layers)
57

Woosuk Kwon's avatar
Woosuk Kwon committed
58
59

class LlamaMLP(nn.Module):
60

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

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


class LlamaAttention(nn.Module):

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

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

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

156
157
158
159
        is_neox_style = True
        if quant_config is not None and quant_config.get_name() == "gguf":
            is_neox_style = False

160
161
162
163
164
165
        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,
166
            is_neox_style=is_neox_style,
167
        )
168
169
170
171
172
173
174
175
        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,
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
176
177
178

    def forward(
        self,
179
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
180
        hidden_states: torch.Tensor,
181
182
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
183
    ) -> torch.Tensor:
184
        qkv, _ = self.qkv_proj(hidden_states)
Zhuohan Li's avatar
Zhuohan Li committed
185
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
186
        q, k = self.rotary_emb(positions, q, k)
187
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
188
189
190
191
192
193
        output, _ = self.o_proj(attn_output)
        return output


class LlamaDecoderLayer(nn.Module):

194
195
196
    def __init__(
        self,
        config: LlamaConfig,
197
        cache_config: Optional[CacheConfig] = None,
198
        quant_config: Optional[QuantizationConfig] = None,
199
        prefix: str = "",
200
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
201
202
        super().__init__()
        self.hidden_size = config.hidden_size
203
204
        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
205
206
207
208
        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)
209
210
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
211
212
213
214
        # 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
215
        self.self_attn = LlamaAttention(
216
            config=config,
217
218
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
219
220
            num_kv_heads=getattr(config, "num_key_value_heads",
                                 config.num_attention_heads),
221
222
223
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
224
            quant_config=quant_config,
225
            bias=attention_bias,
226
            cache_config=cache_config,
227
            prefix=f"{prefix}.self_attn",
Woosuk Kwon's avatar
Woosuk Kwon committed
228
229
        )
        self.mlp = LlamaMLP(
230
231
232
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
233
            quant_config=quant_config,
234
            bias=getattr(config, "mlp_bias", False),
235
            prefix=f"{prefix}.mlp",
Woosuk Kwon's avatar
Woosuk Kwon committed
236
        )
237
238
239
240
        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
241
242
243

    def forward(
        self,
244
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
245
        hidden_states: torch.Tensor,
246
247
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
248
249
        residual: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
250
        # Self Attention
251
252
253
254
255
256
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
257
258
259
260
        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
261
262

        # Fully Connected
263
264
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
265
        hidden_states = self.mlp(hidden_states)
266
        return hidden_states, residual
Woosuk Kwon's avatar
Woosuk Kwon committed
267
268


269
@support_compile_llama_style
Woosuk Kwon's avatar
Woosuk Kwon committed
270
271
class LlamaModel(nn.Module):

272
273
274
    def __init__(
        self,
        config: LlamaConfig,
275
        cache_config: Optional[CacheConfig] = None,
276
        quant_config: Optional[QuantizationConfig] = None,
277
        lora_config: Optional[LoRAConfig] = None,
278
        prefix: str = "",
279
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
280
281
282
        super().__init__()
        self.config = config
        self.padding_idx = config.pad_token_id
283
284
285
286
        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
287
288
289
290
291
292
        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,
293
                quant_config=quant_config,
294
295
296
            )
        else:
            self.embed_tokens = PPMissingLayer()
297
        self.start_layer, self.end_layer, self.layers = make_layers(
298
            config.num_hidden_layers,
299
300
301
302
            lambda prefix: LlamaDecoderLayer(config=config,
                                             cache_config=cache_config,
                                             quant_config=quant_config,
                                             prefix=prefix),
303
304
            prefix=f"{prefix}.layers",
        )
305
306
307
308
        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
309

310
311
312
313
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))

314
315
316
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Woosuk Kwon's avatar
Woosuk Kwon committed
317
318
    def forward(
        self,
319
        input_ids: Optional[torch.Tensor],
320
        positions: torch.Tensor,
321
322
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
323
        intermediate_tensors: Optional[IntermediateTensors],
324
        inputs_embeds: Optional[torch.Tensor] = None,
325
326
327
328
329
330
331
    ) -> 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
332
        else:
333
334
335
336
337
            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
338
            layer = self.layers[i]
339
340
341
            hidden_states, residual = layer(positions, hidden_states,
                                            kv_caches[i - self.start_layer],
                                            attn_metadata, residual)
342
343
344
345
346
347
348

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

349
        hidden_states, _ = self.norm(hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
350
351
        return hidden_states

352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        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())
        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)
                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)

    # 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

            if is_hip():
                # 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
            if hasattr(layer_self_attn, "kv_scale"):
                layer_self_attn.attn._kv_scale = scaling_factor
            else:
                raise RuntimeError("Self attention has no KV cache scaling "
                                   "factor attribute!")

Woosuk Kwon's avatar
Woosuk Kwon committed
436

437
class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
Terry's avatar
Terry committed
438
    packed_modules_mapping = {
439
440
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"]
Terry's avatar
Terry committed
441
442
443
444
    }

    # LoRA specific attributes
    supported_lora_modules = [
445
446
        "qkv_proj", "o_proj", "gate_up_proj", "down_proj", "embed_tokens",
        "lm_head"
Terry's avatar
Terry committed
447
448
449
    ]
    embedding_modules = {
        "embed_tokens": "input_embeddings",
450
        "lm_head": "output_embeddings"
Terry's avatar
Terry committed
451
452
    }
    embedding_padding_modules = ["lm_head"]
453
454
455
456
457
458
459
460
461
462
463
464
465

    # BitandBytes specific attributes
    default_bitsandbytes_target_modules = [
        ".gate_proj.",
        ".down_proj.",
        ".up_proj.",
        ".q_proj.",
        ".k_proj.",
        ".v_proj.",
        ".o_proj.",
    ]
    # in TP, these weights are partitioned along the column dimension (dim=-1)
    column_parallel_weights_modules = [".down_proj.", ".o_proj."]
466
467
468
469
470
471
472
473
    bitsandbytes_stacked_params_mapping = {
        # shard_name, weight_name, index
        "q_proj": ("qkv_proj", 0),
        "k_proj": ("qkv_proj", 1),
        "v_proj": ("qkv_proj", 2),
        "gate_proj": ("gate_up_proj", 0),
        "up_proj": ("gate_up_proj", 1),
    }
474

475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
    # 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"
    }
494

495
496
497
    def __init__(
        self,
        config: LlamaConfig,
498
        cache_config: Optional[CacheConfig] = None,
499
        quant_config: Optional[QuantizationConfig] = None,
500
        lora_config: Optional[LoRAConfig] = None,
501
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
502
        super().__init__()
503

Woosuk Kwon's avatar
Woosuk Kwon committed
504
        self.config = config
505
506
        self.lora_config = lora_config

507
508
509
        self.model = LlamaModel(config,
                                cache_config,
                                quant_config,
510
511
                                lora_config=lora_config,
                                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
527
528
                quant_config=quant_config,
            )
            if config.tie_word_embeddings:
529
530
                self.lm_head = self.lm_head.tie_weights(
                    self.model.embed_tokens)
531
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)
            self.sampler = Sampler()
        else:
            self.lm_head = PPMissingLayer()
539
540
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Woosuk Kwon's avatar
Woosuk Kwon committed
541
542
543

    def forward(
        self,
544
545
        input_ids: torch.Tensor,
        positions: torch.Tensor,
546
547
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
548
549
550
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        model_output = self.model(input_ids, positions, kv_caches,
Alphi's avatar
Alphi committed
551
                                  attn_metadata, intermediate_tensors)
552
        return model_output
553

554
555
556
557
558
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
559
        logits = self.logits_processor(self.lm_head, hidden_states,
560
561
562
                                       sampling_metadata)
        return logits

563
564
    def sample(self, logits: torch.Tensor,
               sampling_metadata: SamplingMetadata) -> Optional[SamplerOutput]:
565
        next_tokens = self.sampler(logits, sampling_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
566
567
        return next_tokens

568
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
569
570
571
572
573
574
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=(["lm_head."]
                           if self.config.tie_word_embeddings else None),
        )
        loader.load_weights(
575
            self.maybe_remap_mistral(name, loaded_weight)
576
            for name, loaded_weight in weights)
577
578

    def load_kv_cache_scales(self, quantization_param_path: str) -> None:
579
        self.model.load_kv_cache_scales(quantization_param_path)
580
581
582
583

    # This function is used to remap the mistral format as
    # used by Mistral and Llama <=2
    def maybe_remap_mistral(
584
585
586
587
        self,
        name: str,
        loaded_weight: torch.Tensor,
    ) -> Tuple[str, torch.Tensor]:
588

589
        def permute(w: torch.Tensor, n_heads: int):
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
            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