llama.py 25.6 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.config import CacheConfig, LoRAConfig
34
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
35
                              get_tensor_model_parallel_world_size)
Woosuk Kwon's avatar
Woosuk Kwon committed
36
from vllm.model_executor.layers.activation import SiluAndMul
37
from vllm.model_executor.layers.layernorm import RMSNorm
38
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
39
40
                                               QKVParallelLinear,
                                               RowParallelLinear)
41
from vllm.model_executor.layers.logits_processor import LogitsProcessor
42
43
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
44
45
from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
    get_compressed_tensors_cache_scale)
46
from vllm.model_executor.layers.rotary_embedding import get_rope
Woosuk Kwon's avatar
Woosuk Kwon committed
47
from vllm.model_executor.layers.sampler import Sampler
48
from vllm.model_executor.layers.vocab_parallel_embedding import (
49
    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
50
from vllm.model_executor.model_loader.weight_utils import (
51
    default_weight_loader, kv_cache_scales_loader, maybe_remap_kv_scale_name)
52
from vllm.model_executor.sampling_metadata import SamplingMetadata
53
from vllm.sequence import IntermediateTensors, SamplerOutput
54
from vllm.utils import is_hip
Woosuk Kwon's avatar
Woosuk Kwon committed
55

56
from .interfaces import SupportsLoRA
57
from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers
58

gaoqiong's avatar
gaoqiong committed
59
from vllm import _custom_ops as ops
60
61
from vllm.model_executor.utils import pad_weight, gemm_bank_conf

Woosuk Kwon's avatar
Woosuk Kwon committed
62
63

class LlamaMLP(nn.Module):
64

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

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


class LlamaAttention(nn.Module):

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

139
        self.qkv_proj = QKVParallelLinear(
140
141
142
143
            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,
144
            bias=bias,
145
            quant_config=quant_config,
146
            prefix=f"{prefix}.qkv_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
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
160
161
162
        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,
        )
163
164
165
166
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
167
168
                              cache_config=cache_config,
                              quant_config=quant_config)
Woosuk Kwon's avatar
Woosuk Kwon committed
169
170
171

    def forward(
        self,
172
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
173
        hidden_states: torch.Tensor,
174
175
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
176
    ) -> torch.Tensor:
177
        qkv, _ = self.qkv_proj(hidden_states)
zhuwenwen's avatar
zhuwenwen committed
178
        if os.environ.get('FA_PAD') == '1':
179
            qkv = qkv[...,:-32]
Zhuohan Li's avatar
Zhuohan Li committed
180
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
181
        q, k = self.rotary_emb(positions, q, k)
182
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
183
184
185
186
187
188
        output, _ = self.o_proj(attn_output)
        return output


class LlamaDecoderLayer(nn.Module):

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

    def forward(
        self,
239
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
240
        hidden_states: torch.Tensor,
241
242
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
243
244
        residual: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
245
        # Self Attention
246
247
248
249
250
251
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
252
253
254
255
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
            kv_cache=kv_cache,
256
            attn_metadata=attn_metadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
257
258
259
        )

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


class LlamaModel(nn.Module):

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

304
305
306
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

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

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

343
        hidden_states, _ = self.norm(hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
344
345
346
        return hidden_states


347
class LlamaForCausalLM(nn.Module, SupportsLoRA):
Terry's avatar
Terry committed
348
349
350
351
352
353
354
355
356
357
358
359
360
361
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # LoRA specific attributes
    supported_lora_modules = [
362
363
        "qkv_proj", "o_proj", "gate_up_proj", "down_proj", "embed_tokens",
        "lm_head"
Terry's avatar
Terry committed
364
365
366
367
368
369
    ]
    embedding_modules = {
        "embed_tokens": "input_embeddings",
        "lm_head": "output_embeddings",
    }
    embedding_padding_modules = ["lm_head"]
370
371
372
373
374
375
376
377
    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),
    }
378

379
380
381
    def __init__(
        self,
        config: LlamaConfig,
382
        cache_config: Optional[CacheConfig] = None,
383
        quant_config: Optional[QuantizationConfig] = None,
384
        lora_config: Optional[LoRAConfig] = None,
385
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
386
        super().__init__()
387

Woosuk Kwon's avatar
Woosuk Kwon committed
388
        self.config = config
389
390
        self.lora_config = lora_config

391
392
393
        self.model = LlamaModel(config,
                                cache_config,
                                quant_config,
394
395
                                lora_config=lora_config,
                                prefix="model")
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
        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,
                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,
                quant_config=quant_config,
            )
            if config.tie_word_embeddings:
                self.lm_head.weight = self.model.embed_tokens.weight

            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()
gaoqiong's avatar
gaoqiong committed
420
421
422
423
424
425
        self.quant_method =  None
     
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
            self.quant_config=quant_config
              
gaoqiong's avatar
gaoqiong committed
426
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
427
428
        self.use_gemm_pad = os.environ.get('GEMM_PAD') == '1'
        self.use_fa_pad = os.environ.get('FA_PAD') == '1'
Woosuk Kwon's avatar
Woosuk Kwon committed
429
430
431

    def forward(
        self,
432
433
        input_ids: torch.Tensor,
        positions: torch.Tensor,
434
435
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
436
437
438
439
440
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        model_output = self.model(input_ids, positions, kv_caches,
                                  attn_metadata, intermediate_tensors)
        return model_output
441

442
443
    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
444
        logits = self.logits_processor(self.lm_head, hidden_states,
445
446
447
                                       sampling_metadata)
        return logits

448
449
    def sample(
        self,
450
        logits: torch.Tensor,
451
        sampling_metadata: SamplingMetadata,
452
    ) -> Optional[SamplerOutput]:
453
        next_tokens = self.sampler(logits, sampling_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
454
455
        return next_tokens

456
457
458
459
460
461
462
463
464
465
466
467
468
469
    def make_empty_intermediate_tensors(
            self, batch_size: int, dtype: torch.dtype,
            device: torch.device) -> IntermediateTensors:
        return IntermediateTensors({
            "hidden_states":
            torch.zeros((batch_size, self.config.hidden_size),
                        dtype=dtype,
                        device=device),
            "residual":
            torch.zeros((batch_size, self.config.hidden_size),
                        dtype=dtype,
                        device=device),
        })

470
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
471
472
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
473
474
475
476
477
            (".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
478
        ]
479
        params_dict = dict(self.named_parameters())
480
        for name, loaded_weight in weights:
481
482
            if "rotary_emb.inv_freq" in name:
                continue
483
484
485
486
            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.
487
                continue
488
489
490
491
492
            # With tie_word_embeddings, we can skip lm_head.weight
            # The weight might appear unnecessarily in the files if the model is
            # processed with quantization, LoRA, fine-tuning, etc.
            if self.config.tie_word_embeddings and "lm_head.weight" in name:
                continue
493
494
495
496
497
498
499
500
            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
501
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
Zhuohan Li's avatar
Zhuohan Li committed
502
                if weight_name not in name:
503
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
504
505
506
507
                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
508
509
510
511

                if is_pp_missing_parameter(name, self):
                    continue

CHU Tianxiang's avatar
CHU Tianxiang committed
512
                param = params_dict[name]
513
514
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
515

516
                break
517
            else:
CHU Tianxiang's avatar
CHU Tianxiang committed
518
519
520
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
521
                # Remapping the name of FP8 kv-scale.
522
523
524
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue
525
526
527
528

                if is_pp_missing_parameter(name, self):
                    continue

529
530
531
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
gaoqiong's avatar
gaoqiong committed
532
533
534
535
536
537
538
                weight_loader(param, loaded_weight)  
            
        if self.use_llama_nn:
            lay_key_words = [
                "self_attn.qkv_proj.weight",
                "self_attn.o_proj.weight",
                "mlp.gate_up_proj.weight",
539
540
                "mlp.down_proj.weight",
                "lm_head.weight"
gaoqiong's avatar
gaoqiong committed
541
542
543
            ]
            combined_words = "|".join(lay_key_words)
            
zhuwenwen's avatar
zhuwenwen committed
544
545
546
            lay_qkv_words = ["self_attn.qkv_proj.weight"]   
            qkv_words = "|".join(lay_qkv_words)          
            
gaoqiong's avatar
gaoqiong committed
547
548
            for layername, weight in params_dict.items():
                matches = re.findall(combined_words, layername)
549
550
551
552
                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
553
554
555
                    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)
556
                                 
gaoqiong's avatar
gaoqiong committed
557
558
559
                    _weight = torch.zeros_like(weight.data)
                    ori_shape =_weight.shape
                    
zhuwenwen's avatar
zhuwenwen committed
560
                    ops.trans_w16_gemm(_weight, weight.data, _weight.shape[0], _weight.shape[1])
gaoqiong's avatar
gaoqiong committed
561
562
                    weight.data.copy_(_weight)
                    
zhuwenwen's avatar
zhuwenwen committed
563
                    weight.data=weight.data.reshape(ori_shape[1], -1)
gaoqiong's avatar
gaoqiong committed
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
     
        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
589
                    _qw, _sz=ops.convert_s4(qweight,qzeros,scales,int(group_size)) 
gaoqiong's avatar
gaoqiong committed
590
                    
gaoqiong's avatar
gaoqiong committed
591
                    sz = ops.sz_permute(_sz).reshape(-1,dim_n)       
gaoqiong's avatar
gaoqiong committed
592
593
594
                    
                    zeros_and_scalse.data.copy_(sz)
                    qweight.data.copy_(_qw)
gaoqiong's avatar
gaoqiong committed
595
                    
gaoqiong's avatar
gaoqiong committed
596
597
598
599
600
601
602
603
604
605
                    #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]  
                
                    if dim_k % 4096==0:
                        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()
                         
606
607
608
609
610
611
612
613
614
615
    # 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):
616
617
            if not isinstance(self.model.layers[layer_idx], nn.Identity):
                layer_self_attn = self.model.layers[layer_idx].self_attn
618
619
620
621
622
623
624
625

            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"):
626
                layer_self_attn.attn._kv_scale = scaling_factor
627
628
629
            else:
                raise RuntimeError("Self attention has no KV cache scaling "
                                   "factor attribute!")