qwen2.py 29.6 KB
Newer Older
Junyang Lin's avatar
Junyang Lin committed
1
2
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py
3
4
# Copyright 2024 The Qwen team.
# Copyright 2023 The vLLM team.
Junyang Lin's avatar
Junyang Lin committed
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.
23
"""Inference-only Qwen2 model compatible with HuggingFace weights."""
24
from typing import Iterable, List, Optional, Set, Tuple, Union
Junyang Lin's avatar
Junyang Lin committed
25
26
27
28

import torch
from torch import nn
from transformers import Qwen2Config
zhuwenwen's avatar
zhuwenwen committed
29
import os
gaoqiong's avatar
gaoqiong committed
30
import re
Junyang Lin's avatar
Junyang Lin committed
31

32
from vllm.attention import Attention, AttentionMetadata, AttentionType
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
from vllm.logger import init_logger
Junyang Lin's avatar
Junyang Lin committed
37
38
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
39
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
Junyang Lin's avatar
Junyang Lin committed
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
from vllm.model_executor.layers.rotary_embedding import get_rope
Joe Runde's avatar
Joe Runde committed
46
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
Junyang Lin's avatar
Junyang Lin committed
47
from vllm.model_executor.layers.vocab_parallel_embedding import (
48
    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
49
50
from vllm.model_executor.model_loader.weight_utils import (
    default_weight_loader, maybe_remap_kv_scale_name)
51
from vllm.model_executor.pooling_metadata import PoolingMetadata
Junyang Lin's avatar
Junyang Lin committed
52
from vllm.model_executor.sampling_metadata import SamplingMetadata
53
from vllm.sequence import IntermediateTensors, PoolerOutput
Junyang Lin's avatar
Junyang Lin committed
54
55


56
from .interfaces import SupportsLoRA, SupportsPP
57
58
from .utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,
                    is_pp_missing_parameter,
59
60
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
gaoqiong's avatar
gaoqiong committed
61
from vllm import _custom_ops as ops
62
from vllm.model_executor.utils import pad_weight, gemm_bank_conf
zhuwenwen's avatar
zhuwenwen committed
63
from vllm.utils import W8a8GetCacheJSON
64

65
66
logger = init_logger(__name__)

67

Junyang Lin's avatar
Junyang Lin committed
68
69
70
71
72
73
74
class Qwen2MLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
75
        quant_config: Optional[QuantizationConfig] = None,
76
        prefix: str = "",
Junyang Lin's avatar
Junyang Lin committed
77
78
79
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
80
81
            hidden_size,
            [intermediate_size] * 2,
Junyang Lin's avatar
Junyang Lin committed
82
            bias=False,
83
84
85
86
87
88
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
Junyang Lin's avatar
Junyang Lin committed
89
            bias=False,
90
91
92
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
Junyang Lin's avatar
Junyang Lin committed
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
        self.act_fn = SiluAndMul()

    def forward(self, x):
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x


class Qwen2Attention(nn.Module):

    def __init__(self,
                 hidden_size: int,
                 num_heads: int,
                 num_kv_heads: int,
                 max_position: int = 4096 * 32,
                 rope_theta: float = 10000,
113
                 cache_config: Optional[CacheConfig] = None,
114
                 quant_config: Optional[QuantizationConfig] = None,
115
116
                 rope_scaling: Optional[Tuple] = None,
                 prefix: str = "") -> None:
Junyang Lin's avatar
Junyang Lin committed
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
        super().__init__()
        self.hidden_size = hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        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)
        self.head_dim = hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5
        self.rope_theta = rope_theta

        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=True,
145
            quant_config=quant_config,
146
            prefix=f"{prefix}.qkv_proj",
Junyang Lin's avatar
Junyang Lin committed
147
148
149
150
151
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
152
            quant_config=quant_config,
153
            prefix=f"{prefix}.o_proj",
Junyang Lin's avatar
Junyang Lin committed
154
155
156
157
158
159
160
        )

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
            base=self.rope_theta,
161
            rope_scaling=rope_scaling,
Junyang Lin's avatar
Junyang Lin committed
162
        )
163
164
165
166
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
167
                              cache_config=cache_config,
168
169
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
Junyang Lin's avatar
Junyang Lin committed
170

171
172
173
174
        self.quant_method = None
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
            self.quant_config=quant_config
Junyang Lin's avatar
Junyang Lin committed
175
176
177
178
179

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
180
181
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
182
        attn_type: str = AttentionType.DECODER,
Junyang Lin's avatar
Junyang Lin committed
183
184
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
zhuwenwen's avatar
zhuwenwen committed
185
186
        # if os.environ.get('FA_PAD') == '1' and self.quant_method is None:
        #     qkv = qkv[...,:-32]
Junyang Lin's avatar
Junyang Lin committed
187
188
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
189
190
191
192
193
194
        attn_output = self.attn(q,
                                k,
                                v,
                                kv_cache,
                                attn_metadata,
                                attn_type=attn_type)
Junyang Lin's avatar
Junyang Lin committed
195
196
197
198
199
200
201
202
203
        output, _ = self.o_proj(attn_output)
        return output


class Qwen2DecoderLayer(nn.Module):

    def __init__(
        self,
        config: Qwen2Config,
204
        cache_config: Optional[CacheConfig] = None,
205
        quant_config: Optional[QuantizationConfig] = None,
206
        prefix: str = "",
Junyang Lin's avatar
Junyang Lin committed
207
208
209
210
211
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        # Requires transformers > 4.32.0
        rope_theta = getattr(config, "rope_theta", 1000000)
212
        rope_scaling = getattr(config, "rope_scaling", None)
Junyang Lin's avatar
Junyang Lin committed
213
214
215
216
217
218
        self.self_attn = Qwen2Attention(
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            max_position=config.max_position_embeddings,
            num_kv_heads=config.num_key_value_heads,
            rope_theta=rope_theta,
219
            cache_config=cache_config,
220
            quant_config=quant_config,
221
222
223
            rope_scaling=rope_scaling,
            prefix=f"{prefix}.self_attn",
        )
Junyang Lin's avatar
Junyang Lin committed
224
225
226
227
        self.mlp = Qwen2MLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
228
            quant_config=quant_config,
229
            prefix=f"{prefix}.mlp",
Junyang Lin's avatar
Junyang Lin committed
230
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)

236
237
238
239
240
241
242
243
244
        # By default, Qwen2 uses causal attention as it is a decoder-only model.
        # You can override the HF config with `is_causal=False` to enable
        # bidirectional attention, which is used in some embedding models
        # (e.g. Alibaba-NLP/gte-Qwen2-7B-instruct)
        if getattr(config, "is_causal", True):
            self._attn_type = AttentionType.DECODER
        else:
            self._attn_type = AttentionType.ENCODER_ONLY

Junyang Lin's avatar
Junyang Lin committed
245
246
247
248
    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
249
250
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Junyang Lin's avatar
Junyang Lin committed
251
252
253
254
255
256
257
258
259
260
261
262
263
        residual: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
            kv_cache=kv_cache,
264
            attn_metadata=attn_metadata,
265
            attn_type=self._attn_type,
Junyang Lin's avatar
Junyang Lin committed
266
267
268
269
270
271
272
273
274
        )

        # Fully Connected
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


275
@support_torch_compile
Junyang Lin's avatar
Junyang Lin committed
276
277
class Qwen2Model(nn.Module):

278
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Junyang Lin's avatar
Junyang Lin committed
279
        super().__init__()
280
281
282
283
284

        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

285
286
287
288
289
290
291
292
293
294
295
296
        # TODO (@robertgshaw2): see if this can be moved out
        if (cache_config.sliding_window is not None
                and hasattr(config, "max_window_layers")):
            raise ValueError("Sliding window for some but all layers is not "
                             "supported. This model uses sliding window "
                             "but `max_window_layers` = {} is less than "
                             "`num_hidden_layers` = {}. Please open an issue "
                             "to discuss this feature.".format(
                                 config.max_window_layers,
                                 config.num_hidden_layers,
                             ))

Junyang Lin's avatar
Junyang Lin committed
297
298
299
300
        self.config = config
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

301
302
303
304
305
306
        if get_pp_group().is_first_rank or (config.tie_word_embeddings
                                            and get_pp_group().is_last_rank):
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
307
                prefix=f"{prefix}.embed_tokens",
308
309
310
311
            )
        else:
            self.embed_tokens = PPMissingLayer()

312
313
314
315
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: Qwen2DecoderLayer(config=config,
                                             cache_config=cache_config,
316
317
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.layers"),
318
319
320
            prefix=f"{prefix}.layers",
        )

321
322
323
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
324
325
326
327
        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()
zhuwenwen's avatar
zhuwenwen committed
328
329
330
331
332
            
        self.quant_method = None
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
            self.quant_config=quant_config
zhuwenwen's avatar
zhuwenwen committed
333
        self.tritonsingleton= W8a8GetCacheJSON()
zhuwenwen's avatar
zhuwenwen committed
334
335
336
337
338
            
        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'
zhuwenwen's avatar
zhuwenwen committed
339
        self.w8a8_strategy=int(os.getenv('W8A8_SUPPORT_METHODS', '1'))
Junyang Lin's avatar
Junyang Lin committed
340

341
342
343
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Junyang Lin's avatar
Junyang Lin committed
344
345
346
347
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
348
349
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
350
351
        intermediate_tensors: Optional[IntermediateTensors] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
352
    ) -> Union[torch.Tensor, IntermediateTensors]:
353
354
355
356
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
357
                hidden_states = self.get_input_embeddings(input_ids)
358
            residual = None
359
        else:
360
361
362
363
            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):
Junyang Lin's avatar
Junyang Lin committed
364
365
366
367
            layer = self.layers[i]
            hidden_states, residual = layer(
                positions,
                hidden_states,
368
                kv_caches[i - self.start_layer],
369
                attn_metadata,
Junyang Lin's avatar
Junyang Lin committed
370
371
                residual,
            )
372
373
374
375
376
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
                "residual": residual
            })
Junyang Lin's avatar
Junyang Lin committed
377
378
379
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states

380
381
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
Junyang Lin's avatar
Junyang Lin committed
382
383
384
385
386
387
388
389
        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),
        ]
Roy's avatar
Roy committed
390
        params_dict = dict(self.named_parameters(remove_duplicate=False))
391
        loaded_params: Set[str] = set()
392
        for name, loaded_weight in weights:
Junyang Lin's avatar
Junyang Lin committed
393
394
395
396
397
398
399
400
401
            if "rotary_emb.inv_freq" in name:
                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
402
403
                if is_pp_missing_parameter(name, self):
                    continue
Junyang Lin's avatar
Junyang Lin committed
404
405
406
407
408
409
410
411
                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
412
                # Remapping the name of FP8 kv-scale.
413
414
415
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue
416
417
                if is_pp_missing_parameter(name, self):
                    continue
Junyang Lin's avatar
Junyang Lin committed
418
419
420
421
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
422
            loaded_params.add(name)
zhuwenwen's avatar
zhuwenwen committed
423
            
424
        if self.use_llama_nn and self.quant_method is None:
gaoqiong's avatar
gaoqiong committed
425
426
427
428
            lay_key_words = [
                "self_attn.qkv_proj.weight",
                "self_attn.o_proj.weight",
                "mlp.gate_up_proj.weight",
429
                "mlp.down_proj.weight",
gaoqiong's avatar
gaoqiong committed
430
431
432
            ]
            combined_words = "|".join(lay_key_words)
            
zhuwenwen's avatar
zhuwenwen committed
433
434
435
436
437
438
            lay_qkv_words = ["self_attn.qkv_proj.weight"]   
            qkv_words = "|".join(lay_qkv_words)  
            
            lay_qkv_bias_words = ["self_attn.qkv_proj.bias"]   
            qkv_bias_words = "|".join(lay_qkv_bias_words) 
            
gaoqiong's avatar
gaoqiong committed
439
            for layername, weight in params_dict.items():
440
441
442
443
444
445
                if "lm_head.weight" in layername and weight.shape[1] >= 3584:
                    lay_key_words.append("lm_head.weight")
                    combined_words = "|".join(lay_key_words)
                    os.environ['LM_NN'] = '1'  
                else:
                    os.environ['LM_NN'] = '0' 
zhuwenwen's avatar
zhuwenwen committed
446
447
                # if self.use_fa_pad and (re.findall(qkv_bias_words, layername)):
                #     weight.data = pad_weight(weight.data, 32)
zhuwenwen's avatar
zhuwenwen committed
448
                    
gaoqiong's avatar
gaoqiong committed
449
                matches = re.findall(combined_words, layername)
zhuwenwen's avatar
zhuwenwen committed
450
                if matches:   
zhuwenwen's avatar
zhuwenwen committed
451
452
                    # 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
453
                    
zhuwenwen's avatar
zhuwenwen committed
454
455
456
                    # 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)
457
                        
gaoqiong's avatar
gaoqiong committed
458
459
460
                    _weight = torch.zeros_like(weight.data)
                    ori_shape =_weight.shape
                    
zhuwenwen's avatar
zhuwenwen committed
461
                    ops.trans_w16_gemm(_weight, weight.data, _weight.shape[0], _weight.shape[1])
gaoqiong's avatar
gaoqiong committed
462
463
464
465
                    weight.data.copy_(_weight)
                    
                    weight.data=weight.data.reshape(ori_shape[1],-1)
                    
gaoqiong's avatar
gaoqiong committed
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
        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
490
                    _qw, _sz=ops.convert_s4(qweight,qzeros,scales,int(group_size)) 
gaoqiong's avatar
gaoqiong committed
491
                    
gaoqiong's avatar
gaoqiong committed
492
                    sz = ops.sz_permute(_sz).reshape(-1,dim_n)       
gaoqiong's avatar
gaoqiong committed
493
494
495
496
497
498
499
500
                    
                    zeros_and_scalse.data.copy_(sz)
                    qweight.data.copy_(_qw)
                    
                    #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
501
                    if dim_k % 4096==0 and self.use_awq_pad:
gaoqiong's avatar
gaoqiong committed
502
503
504
505
                        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
506
507
508
509
510
511
512
513
514
                   
        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)
zhuwenwen's avatar
zhuwenwen committed
515
516
            weight_shapes=[]
            all_json={}
zhuwenwen's avatar
zhuwenwen committed
517
518
519
            
            for layername, weight in params_dict.items():  
                matches = re.findall(combined_words, layername)
zhuwenwen's avatar
zhuwenwen committed
520
                if matches and "scale" not in layername:
zhuwenwen's avatar
zhuwenwen committed
521
                    weight_data =params_dict[layername]
zhuwenwen's avatar
zhuwenwen committed
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
                    n=weight_data.shape[0]
                    
                    #rocblas和cutlass目前都需要weight做处理,但是triton不用
                    if self.w8a8_strategy!=1:
                        _weight=weight_data.T.contiguous().reshape(n,-1)
                        weight_data.data.copy_(_weight)  
                    
                    #下面是针对模型记录模型出现k和n值 
                    elif len(weight_shapes)<4: 
                        k=weight_data.shape[1]
                        weight_shapes.append({n,k})
                
                        json_file=self.tritonsingleton.get_w8a8json_name(n,k)
                        configs_dict=self.tritonsingleton.get_triton_cache(json_file,n,k)
                        if configs_dict:
                            all_json.update(configs_dict)
                                              
            if self.w8a8_strategy==1:
                self.tritonsingleton.triton_json_dict.append(all_json)
                #找到的所有config都进行一次warmup
                for key, value in all_json.items():
                    m=int(key.split('_')[0])
                    n=int(key.split('_')[1])
                    k=int(key.split('_')[2])
                    ops.triton_int8_gemm_helper(m=m,n=n,k=k,per_token_act_quant=True,per_out_channel_weight_quant=True,use_bias=False,best_config=value)
 
zhuwenwen's avatar
zhuwenwen committed
548
                    
549
        return loaded_params
550

Junyang Lin's avatar
Junyang Lin committed
551

552
class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # LoRA specific attributes
    supported_lora_modules = [
        "qkv_proj",
        "o_proj",
        "gate_up_proj",
        "down_proj",
    ]
    embedding_modules = {}
    embedding_padding_modules = []
chenqianfzh's avatar
chenqianfzh committed
574
575

    # BitandBytes specific attributes
576
577
578
579
580
581
582
583
    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),
    }
Junyang Lin's avatar
Junyang Lin committed
584

585
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
586
587
588
589
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config
590

Junyang Lin's avatar
Junyang Lin committed
591
        self.config = config
592
593
        self.lora_config = lora_config

594
        self.quant_config = quant_config
595
        self.model = Qwen2Model(vllm_config=vllm_config,
596
                                prefix=maybe_prefix(prefix, "model"))
597

598
599
600
601
602
603
604
605
606
        if get_pp_group().is_last_rank:
            if config.tie_word_embeddings:
                self.lm_head = self.model.embed_tokens
            else:
                self.lm_head = ParallelLMHead(config.vocab_size,
                                              config.hidden_size,
                                              quant_config=quant_config,
                                              prefix=maybe_prefix(
                                                  prefix, "lm_head"))
607
        else:
608
            self.lm_head = PPMissingLayer()
609

610
        self.logits_processor = LogitsProcessor(config.vocab_size)
Joe Runde's avatar
Joe Runde committed
611
        self.sampler = get_sampler()
zhuwenwen's avatar
zhuwenwen committed
612
613
614
615
616
617
618
619
620
621
        
        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'
622

623
624
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Junyang Lin's avatar
Junyang Lin committed
625

626
627
628
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

Junyang Lin's avatar
Junyang Lin committed
629
630
631
632
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
633
634
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
635
        intermediate_tensors: Optional[IntermediateTensors] = None,
636
        inputs_embeds: Optional[torch.Tensor] = None,
637
    ) -> Union[torch.Tensor, IntermediateTensors]:
Junyang Lin's avatar
Junyang Lin committed
638
        hidden_states = self.model(input_ids, positions, kv_caches,
639
640
                                   attn_metadata, intermediate_tensors,
                                   inputs_embeds)
Junyang Lin's avatar
Junyang Lin committed
641
642
        return hidden_states

643
644
645
646
647
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
648
        logits = self.logits_processor(self.lm_head, hidden_states,
649
650
651
                                       sampling_metadata)
        return logits

Junyang Lin's avatar
Junyang Lin committed
652
653
    def sample(
        self,
654
        logits: torch.Tensor,
Junyang Lin's avatar
Junyang Lin committed
655
656
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
657
        next_tokens = self.sampler(logits, sampling_metadata)
Junyang Lin's avatar
Junyang Lin committed
658
659
        return next_tokens

660
661
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
662
663
664
665
666
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=(["lm_head."]
                           if self.config.tie_word_embeddings else None),
        )
667
        return loader.load_weights(weights)
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692


class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # LoRA specific attributes
    supported_lora_modules = [
        "qkv_proj",
        "o_proj",
        "gate_up_proj",
        "down_proj",
    ]
    embedding_modules = {}
    embedding_padding_modules = []

693
694
    hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})

695
696
697
698
699
700
701
702
703
704
705
706
707
708
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config
        pooler_config = vllm_config.model_config.pooler_config

        self.config = config
        self.lora_config = lora_config

        self.quant_config = quant_config
        self.model = Qwen2Model(vllm_config=vllm_config,
                                prefix=maybe_prefix(prefix, "model"))

709
710
        # TODO: Replace this model class with as_embedding_model(
        # Qwen2ForCausalLM) after changing the default pooling method
711
712
713
714
715
716
717
        if pooler_config.pooling_type is None:
            logger.warning(
                "This embedding model will default to last-token pooling in "
                "an upcoming version. To avoid breaking changes, you should "
                "pass `--override-pooler-config '{\"pooling_type\": \"MEAN\"}'`"
                " explicitly.")

718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
        self._pooler = Pooler.from_config_with_defaults(
            pooler_config,
            pooling_type=PoolingType.MEAN,
            normalize=True,
            softmax=False)

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

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

742
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
743
        weights = self.hf_to_vllm_mapper.apply(weights)
744
745
        weights = ((name, data) for name, data in weights
                   if not name.startswith("lm_head."))
746
        self.model.load_weights(weights)