llama.py 25.8 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_torch_compile
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.pooler import Pooler, PoolingType
42
from vllm.model_executor.layers.quantization import QuantizationConfig
43
44
from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
    get_compressed_tensors_cache_scale)
45
from vllm.model_executor.layers.rotary_embedding import get_rope
46
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
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, kv_cache_scales_loader, maybe_remap_kv_scale_name)
51
from vllm.model_executor.pooling_metadata import PoolingMetadata
52
from vllm.model_executor.sampling_metadata import SamplingMetadata
53
from vllm.sequence import IntermediateTensors, PoolerOutput
54
from vllm.utils import is_hip
Woosuk Kwon's avatar
Woosuk Kwon committed
55

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

Woosuk Kwon's avatar
Woosuk Kwon committed
60
61

class LlamaMLP(nn.Module):
62

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

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


class LlamaAttention(nn.Module):

    def __init__(
        self,
103
        config: LlamaConfig,
104
105
106
107
108
109
        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,
110
        quant_config: Optional[QuantizationConfig] = None,
111
        bias: bool = False,
112
        cache_config: Optional[CacheConfig] = None,
113
        prefix: str = "",
114
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
115
        super().__init__()
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
159
160
161
        is_neox_style = True
        if quant_config is not None and quant_config.get_name() == "gguf":
            is_neox_style = False

162
163
164
165
166
167
        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,
168
            is_neox_style=is_neox_style,
169
        )
170
171
172
173
174
175
176
177
        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
178
179
180

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


class LlamaDecoderLayer(nn.Module):

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

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

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


271
272
273
274
275
276
277
@support_torch_compile(
    dynamic_arg_dims={
        "input_ids": 0,
        "positions": 0,
        "inputs_embeds": 0,
        "intermediate_tensors": 0,
    })
Woosuk Kwon's avatar
Woosuk Kwon committed
278
279
class LlamaModel(nn.Module):

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

318
319
320
321
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))

322
323
324
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

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

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

357
        hidden_states, _ = self.norm(hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
358
359
        return hidden_states

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
436
437
438
439
440
441
442
443
    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
444

445
class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
Terry's avatar
Terry committed
446
    packed_modules_mapping = {
447
448
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"]
Terry's avatar
Terry committed
449
450
451
452
    }

    # LoRA specific attributes
    supported_lora_modules = [
453
454
        "qkv_proj", "o_proj", "gate_up_proj", "down_proj", "embed_tokens",
        "lm_head"
Terry's avatar
Terry committed
455
456
457
    ]
    embedding_modules = {
        "embed_tokens": "input_embeddings",
458
        "lm_head": "output_embeddings"
Terry's avatar
Terry committed
459
460
    }
    embedding_padding_modules = ["lm_head"]
461
462
463
464
465
466
467
468
469
470
471
472
473

    # 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."]
474
475
476
477
478
479
480
481
    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),
    }
482

483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
    # 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"
    }
502

503
504
505
    def __init__(
        self,
        config: LlamaConfig,
506
        cache_config: Optional[CacheConfig] = None,
507
        quant_config: Optional[QuantizationConfig] = None,
508
        lora_config: Optional[LoRAConfig] = None,
509
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
510
        super().__init__()
511

Woosuk Kwon's avatar
Woosuk Kwon committed
512
        self.config = config
513
514
        self.lora_config = lora_config

515
516
517
        self.model = LlamaModel(config,
                                cache_config,
                                quant_config,
518
519
                                lora_config=lora_config,
                                prefix="model")
520
521
522
523
524
525
526
527
        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,
528
529
530
531
532
533
                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),
534
535
536
                quant_config=quant_config,
            )
            if config.tie_word_embeddings:
537
538
                self.lm_head = self.lm_head.tie_weights(
                    self.model.embed_tokens)
539
540
541
542
543
544
545
546

            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()
547
548
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Woosuk Kwon's avatar
Woosuk Kwon committed
549
550
551

    def forward(
        self,
552
553
        input_ids: torch.Tensor,
        positions: torch.Tensor,
554
555
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
556
557
558
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        model_output = self.model(input_ids, positions, kv_caches,
Alphi's avatar
Alphi committed
559
                                  attn_metadata, intermediate_tensors)
560
        return model_output
561

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

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

576
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
577
578
579
580
581
582
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=(["lm_head."]
                           if self.config.tie_word_embeddings else None),
        )
        loader.load_weights(
583
            self.maybe_remap_mistral(name, loaded_weight)
584
            for name, loaded_weight in weights)
585
586

    def load_kv_cache_scales(self, quantization_param_path: str) -> None:
587
        self.model.load_kv_cache_scales(quantization_param_path)
588
589
590
591

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

597
        def permute(w: torch.Tensor, n_heads: int):
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
            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
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668


class LlamaEmbeddingModel(nn.Module, SupportsPP):
    """
    A model that uses Llama with additional embedding functionalities.

    This class encapsulates the LlamaModel and provides an interface for
    embedding operations and customized pooling functions.

    Attributes:
        model: An instance of LlamaModel used for forward operations.
        _pooler: An instance of Pooler used for pooling operations.
    """

    def __init__(
        self,
        **kwargs,
    ) -> None:
        super().__init__()

        self.model = LlamaModel(**kwargs)
        self._pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)

    def forward(
        self,
        input_ids: Optional[torch.Tensor],
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        return self.model(input_ids, positions, kv_caches, attn_metadata,
                          intermediate_tensors, inputs_embeds)

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

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        self.model.load_weights(weights)

    def load_kv_cache_scales(self, quantization_param_path: str) -> None:
        self.model.load_kv_cache_scales(quantization_param_path)