llama.py 30.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

import torch
from torch import nn
from transformers import LlamaConfig
zhuwenwen's avatar
zhuwenwen committed
29
import os
gaoqiong's avatar
gaoqiong committed
30
import re
Woosuk Kwon's avatar
Woosuk Kwon committed
31

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

58
from .interfaces import SupportsLoRA, SupportsPP
59
from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
60
                    make_empty_intermediate_tensors_factory, make_layers)
Woosuk Kwon's avatar
Woosuk Kwon committed
61

gaoqiong's avatar
gaoqiong committed
62
from vllm import _custom_ops as ops
63
64
from vllm.model_executor.utils import pad_weight, gemm_bank_conf

Woosuk Kwon's avatar
Woosuk Kwon committed
65
66

class LlamaMLP(nn.Module):
67

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

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


class LlamaAttention(nn.Module):

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

145
        self.qkv_proj = QKVParallelLinear(
146
147
148
149
            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,
150
            bias=bias,
151
            quant_config=quant_config,
152
            prefix=f"{prefix}.qkv_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
153
        )
154

155
        self.o_proj = RowParallelLinear(
156
157
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
158
            bias=bias,
159
            quant_config=quant_config,
160
            prefix=f"{prefix}.o_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
161
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
162

163
164
165
166
        is_neox_style = True
        if quant_config is not None and quant_config.get_name() == "gguf":
            is_neox_style = False

167
168
169
170
171
172
        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,
173
            is_neox_style=is_neox_style,
174
        )
175

176
177
178
179
180
181
182
        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,
183
        )
184
185
186
187
188
        
        self.quant_method = None
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
            self.quant_config=quant_config
Woosuk Kwon's avatar
Woosuk Kwon committed
189
190
191

    def forward(
        self,
192
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
193
        hidden_states: torch.Tensor,
194
195
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
196
    ) -> torch.Tensor:
197
        qkv, _ = self.qkv_proj(hidden_states)
198
        if os.environ.get('FA_PAD') == '1' and self.quant_method is None:
199
            qkv = qkv[...,:-32]
Zhuohan Li's avatar
Zhuohan Li committed
200
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
201
        q, k = self.rotary_emb(positions, q, k)
202
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
203
204
205
206
207
208
        output, _ = self.o_proj(attn_output)
        return output


class LlamaDecoderLayer(nn.Module):

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

    def forward(
        self,
259
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
260
        hidden_states: torch.Tensor,
261
262
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
263
264
        residual: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
265
        # Self Attention
266
267
268
269
270
271
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
272
273
274
275
        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
276
277

        # Fully Connected
278
279
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
280
        hidden_states = self.mlp(hidden_states)
281
        return hidden_states, residual
Woosuk Kwon's avatar
Woosuk Kwon committed
282
283


284
285
286
287
288
289
290
@support_torch_compile(
    dynamic_arg_dims={
        "input_ids": 0,
        "positions": 0,
        "inputs_embeds": 0,
        "intermediate_tensors": 0,
    })
Woosuk Kwon's avatar
Woosuk Kwon committed
291
292
class LlamaModel(nn.Module):

293
294
295
    def __init__(
        self,
        config: LlamaConfig,
296
        cache_config: Optional[CacheConfig] = None,
297
        quant_config: Optional[QuantizationConfig] = None,
298
        lora_config: Optional[LoRAConfig] = None,
299
        prefix: str = "",
300
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
301
302
303
        super().__init__()
        self.config = config
        self.padding_idx = config.pad_token_id
304
305
306
307
        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
308
309
310
311
312
313
        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,
314
                quant_config=quant_config,
315
316
317
            )
        else:
            self.embed_tokens = PPMissingLayer()
318
        self.start_layer, self.end_layer, self.layers = make_layers(
319
            config.num_hidden_layers,
320
321
322
323
            lambda prefix: LlamaDecoderLayer(config=config,
                                             cache_config=cache_config,
                                             quant_config=quant_config,
                                             prefix=prefix),
324
325
            prefix=f"{prefix}.layers",
        )
326
327
328
329
        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
330

331
332
333
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
334
335
336
337
338
339
340
341
342
343
        
        self.quant_method = None
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
            self.quant_config=quant_config
              
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
        self.use_gemm_pad = os.environ.get('GEMM_PAD') == '1'
        self.use_fa_pad = os.environ.get('FA_PAD') == '1'
        self.use_awq_pad = os.environ.get('AWQ_PAD') == '1'
344

345
346
347
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Woosuk Kwon's avatar
Woosuk Kwon committed
348
349
    def forward(
        self,
350
        input_ids: Optional[torch.Tensor],
351
        positions: torch.Tensor,
352
353
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
354
        intermediate_tensors: Optional[IntermediateTensors],
355
        inputs_embeds: Optional[torch.Tensor] = None,
356
357
358
359
360
361
362
    ) -> 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
363
        else:
364
365
366
367
368
            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
369
            layer = self.layers[i]
370
371
372
            hidden_states, residual = layer(positions, hidden_states,
                                            kv_caches[i - self.start_layer],
                                            attn_metadata, residual)
373
374
375
376
377
378
379

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

380
        hidden_states, _ = self.norm(hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
381
382
        return hidden_states

383
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
384
385
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
386
387
388
389
390
            (".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),
Zhuohan Li's avatar
Zhuohan Li committed
391
        ]
392
        params_dict = dict(self.named_parameters())
393
        for name, loaded_weight in weights:
394
395
            if "rotary_emb.inv_freq" in name:
                continue
396
397
398
399
            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.
400
                continue
401
402
403
404
405
406
407
408
            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
409
            for param_name, weight_name, shard_id in stacked_params_mapping:
Zhuohan Li's avatar
Zhuohan Li committed
410
                if weight_name not in name:
411
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
412
413
414
415
                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
416
417
418
419

                if is_pp_missing_parameter(name, self):
                    continue

CHU Tianxiang's avatar
CHU Tianxiang committed
420
                param = params_dict[name]
421
422
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
423

424
                break
425
            else:
CHU Tianxiang's avatar
CHU Tianxiang committed
426
427
428
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
429
                # Remapping the name of FP8 kv-scale.
430
431
432
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue
433
434
435
436

                if is_pp_missing_parameter(name, self):
                    continue

437
438
439
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
440
                weight_loader(param, loaded_weight) 
gaoqiong's avatar
gaoqiong committed
441
            
442
        if self.use_llama_nn and self.quant_method is None :
gaoqiong's avatar
gaoqiong committed
443
444
445
446
            lay_key_words = [
                "self_attn.qkv_proj.weight",
                "self_attn.o_proj.weight",
                "mlp.gate_up_proj.weight",
447
                "mlp.down_proj.weight",
448
                # "lm_head.weight"
gaoqiong's avatar
gaoqiong committed
449
450
451
            ]
            combined_words = "|".join(lay_key_words)
            
zhuwenwen's avatar
zhuwenwen committed
452
453
454
            lay_qkv_words = ["self_attn.qkv_proj.weight"]   
            qkv_words = "|".join(lay_qkv_words)          
            
gaoqiong's avatar
gaoqiong committed
455
456
            for layername, weight in params_dict.items():
                matches = re.findall(combined_words, layername)
457
458
459
460
                if matches:         
                    if self.use_gemm_pad and gemm_bank_conf(weight.data.shape[0]):
                        weight.data = pad_weight(weight.data, 32)  
                        
zhuwenwen's avatar
zhuwenwen committed
461
462
463
                    if self.use_fa_pad and (re.findall(qkv_words, layername)):
                        if not gemm_bank_conf(weight.data.shape[0]):
                            weight.data = pad_weight(weight.data, 32)
464
                                 
gaoqiong's avatar
gaoqiong committed
465
466
467
                    _weight = torch.zeros_like(weight.data)
                    ori_shape =_weight.shape
                    
zhuwenwen's avatar
zhuwenwen committed
468
                    ops.trans_w16_gemm(_weight, weight.data, _weight.shape[0], _weight.shape[1])
gaoqiong's avatar
gaoqiong committed
469
470
                    weight.data.copy_(_weight)
                    
zhuwenwen's avatar
zhuwenwen committed
471
                    weight.data=weight.data.reshape(ori_shape[1], -1)
gaoqiong's avatar
gaoqiong committed
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
     
        if self.quant_method == "awq":
            lay_key_words = [
                "self_attn.qkv_proj.qweight",
                "self_attn.o_proj.qweight",
                "mlp.gate_up_proj.qweight",
                "mlp.down_proj.qweight"
            ]
            combined_words = "|".join(lay_key_words)
            
            for layername, weight in params_dict.items():
                
                matches = re.findall(combined_words, layername)
                if matches:
                    qweight =params_dict[layername]
                    qzeros=params_dict[layername.replace("qweight", "qzeros")]
                    scales=params_dict[layername.replace("qweight", "scales")]
                    zeros_and_scalse =params_dict[layername.replace("qweight", "zeros_and_scales")]
                    
                    group_size= self.quant_config.group_size 
                   
                    dim_n = scales.data.shape[1]
                    dim_k = qweight.data.shape[0]
                    pad_group=2              
                    
gaoqiong's avatar
gaoqiong committed
497
                    _qw, _sz=ops.convert_s4(qweight,qzeros,scales,int(group_size)) 
gaoqiong's avatar
gaoqiong committed
498
                    
gaoqiong's avatar
gaoqiong committed
499
                    sz = ops.sz_permute(_sz).reshape(-1,dim_n)       
gaoqiong's avatar
gaoqiong committed
500
501
502
                    
                    zeros_and_scalse.data.copy_(sz)
                    qweight.data.copy_(_qw)
gaoqiong's avatar
gaoqiong committed
503
                    
gaoqiong's avatar
gaoqiong committed
504
505
506
507
                    #reshape
                    zeros_and_scalse.data=zeros_and_scalse.reshape(dim_n,-1)    #[k/greop_size,n]------>[n,k/group_size]
                    qweight.data=qweight.data.reshape(dim_n,-1)                      #[k,n/8]---->[n,k/8]  
                
zhuwenwen's avatar
zhuwenwen committed
508
                    if dim_k % 4096==0 and self.use_awq_pad:
gaoqiong's avatar
gaoqiong committed
509
510
511
512
                        zeros_and_scalse_pad= torch.zeros(dim_n,pad_group,dtype=torch.int32).cuda()
                        zeros_and_scalse.data=torch.cat((zeros_and_scalse.data,zeros_and_scalse_pad),dim=1).contiguous()
                        qweight_pad= torch.zeros(dim_n,int(group_size//4),dtype=torch.int32).cuda()
                        qweight.data=torch.cat((qweight.data,qweight_pad),dim=1).contiguous()
zhuwenwen's avatar
zhuwenwen committed
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
            
        if self.quant_method == "compressed_tensors":
            lay_key_words = [
                "self_attn.qkv_proj.weight",
                "self_attn.o_proj.weight",
                "mlp.gate_up_proj.weight",
                "mlp.down_proj.weight",
            ]
            combined_words = "|".join(lay_key_words)
            
            for layername, weight in params_dict.items():  
                matches = re.findall(combined_words, layername)
                if matches:
                    weight_data =params_dict[layername]
                    k=weight_data.shape[0]
                    _weight=weight_data.T.contiguous().reshape(k,-1)
                    weight_data.data.copy_(_weight)   
530
                    
531
532
533
534
535
536
537
538
539
540
    # 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):
541
542
            if not isinstance(self.layers[layer_idx], nn.Identity):
                layer_self_attn = self.layers[layer_idx].self_attn
543
544
545
546
547
548
549
550

            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"):
551
                layer_self_attn.attn._kv_scale = scaling_factor
552
553
554
            else:
                raise RuntimeError("Self attention has no KV cache scaling "
                                   "factor attribute!")
555

Woosuk Kwon's avatar
Woosuk Kwon committed
556

557
class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
Terry's avatar
Terry committed
558
    packed_modules_mapping = {
559
560
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"]
Terry's avatar
Terry committed
561
562
563
564
    }

    # LoRA specific attributes
    supported_lora_modules = [
565
566
        "qkv_proj", "o_proj", "gate_up_proj", "down_proj", "embed_tokens",
        "lm_head"
Terry's avatar
Terry committed
567
568
569
    ]
    embedding_modules = {
        "embed_tokens": "input_embeddings",
570
        "lm_head": "output_embeddings"
Terry's avatar
Terry committed
571
572
    }
    embedding_padding_modules = ["lm_head"]
573
574
575
576
577
578
579
580
581
582
583
584
585

    # 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."]
586
587
588
589
590
591
592
593
    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),
    }
594

595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
    # 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"
    }
614

615
616
617
    def __init__(
        self,
        config: LlamaConfig,
618
        cache_config: Optional[CacheConfig] = None,
619
        quant_config: Optional[QuantizationConfig] = None,
620
        lora_config: Optional[LoRAConfig] = None,
621
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
622
        super().__init__()
623

Woosuk Kwon's avatar
Woosuk Kwon committed
624
        self.config = config
625
626
        self.lora_config = lora_config

627
628
629
        self.model = LlamaModel(config,
                                cache_config,
                                quant_config,
630
631
                                lora_config=lora_config,
                                prefix="model")
632
633
634
635
636
637
638
639
        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,
640
641
642
643
644
645
                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),
646
647
648
                quant_config=quant_config,
            )
            if config.tie_word_embeddings:
649
650
                self.lm_head = self.lm_head.tie_weights(
                    self.model.embed_tokens)
651
652
653
654
655
656
657
658

            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()
659
            
660
661
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Woosuk Kwon's avatar
Woosuk Kwon committed
662
663
664

    def forward(
        self,
665
666
        input_ids: torch.Tensor,
        positions: torch.Tensor,
667
668
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
669
670
671
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        model_output = self.model(input_ids, positions, kv_caches,
Alphi's avatar
Alphi committed
672
                                  attn_metadata, intermediate_tensors)
673
        return model_output
674

675
676
677
678
679
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
680
        logits = self.logits_processor(self.lm_head, hidden_states,
681
682
683
                                       sampling_metadata)
        return logits

684
685
    def sample(self, logits: torch.Tensor,
               sampling_metadata: SamplingMetadata) -> Optional[SamplerOutput]:
686
        next_tokens = self.sampler(logits, sampling_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
687
688
        return next_tokens

689
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
690
691
692
693
694
695
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=(["lm_head."]
                           if self.config.tie_word_embeddings else None),
        )
        loader.load_weights(
696
            self.maybe_remap_mistral(name, loaded_weight)
697
            for name, loaded_weight in weights)
698
699

    def load_kv_cache_scales(self, quantization_param_path: str) -> None:
700
        self.model.load_kv_cache_scales(quantization_param_path)
701
702
703
704

    # This function is used to remap the mistral format as
    # used by Mistral and Llama <=2
    def maybe_remap_mistral(
705
706
707
708
        self,
        name: str,
        loaded_weight: torch.Tensor,
    ) -> Tuple[str, torch.Tensor]:
709

710
        def permute(w: torch.Tensor, n_heads: int):
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
            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
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780


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:
zhuwenwen's avatar
zhuwenwen committed
781
        self.model.load_kv_cache_scales(quantization_param_path)