qwen2.py 25.4 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

Junyang Lin's avatar
Junyang Lin committed
4
5
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py
6
7
# Copyright 2024 The Qwen team.
# Copyright 2023 The vLLM team.
Junyang Lin's avatar
Junyang Lin committed
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# 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.
26
"""Inference-only Qwen2 model compatible with HuggingFace weights."""
27

28
from collections.abc import Iterable
29
from itertools import islice
30
from typing import Any
Junyang Lin's avatar
Junyang Lin committed
31
32
33
34
35

import torch
from torch import nn
from transformers import Qwen2Config

zhuwenwen's avatar
zhuwenwen committed
36
import os
gaoqiong's avatar
gaoqiong committed
37
import re
zhuwenwen's avatar
zhuwenwen committed
38

39
from vllm.attention.layer import Attention
40
from vllm.compilation.decorators import support_torch_compile
41
from vllm.config import CacheConfig, VllmConfig
42
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
Junyang Lin's avatar
Junyang Lin committed
43
from vllm.model_executor.layers.activation import SiluAndMul
44
45
46
from vllm.model_executor.layers.attention.encoder_only_attention import (
    EncoderOnlyAttention,
)
Junyang Lin's avatar
Junyang Lin committed
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
49
50
51
52
from vllm.model_executor.layers.linear import (
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
53
from vllm.model_executor.layers.logits_processor import LogitsProcessor
54
from vllm.model_executor.layers.quantization import QuantizationConfig
55
from vllm.model_executor.layers.rotary_embedding import get_rope
Junyang Lin's avatar
Junyang Lin committed
56
from vllm.model_executor.layers.vocab_parallel_embedding import (
57
58
59
    ParallelLMHead,
    VocabParallelEmbedding,
)
60
from vllm.model_executor.model_loader.weight_utils import (
61
62
63
    default_weight_loader,
    maybe_remap_kv_scale_name,
)
64
from vllm.sequence import IntermediateTensors
65
from vllm.transformers_utils.config import is_interleaved, set_default_rope_theta
66
from vllm.v1.attention.backend import AttentionType
Junyang Lin's avatar
Junyang Lin committed
67

68
from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
69
70
71
72
73
74
75
76
77
from .utils import (
    AutoWeightsLoader,
    PPMissingLayer,
    extract_layer_index,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
gaoqiong's avatar
gaoqiong committed
78
from vllm import _custom_ops as ops
79
80
81
from vllm.model_executor.utils import pad_weight, gemm_bank_conf


Junyang Lin's avatar
Junyang Lin committed
82
83
84
85
86
87
class Qwen2MLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
88
        quant_config: QuantizationConfig | None = None,
89
        prefix: str = "",
Junyang Lin's avatar
Junyang Lin committed
90
91
92
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
93
94
            hidden_size,
            [intermediate_size] * 2,
Junyang Lin's avatar
Junyang Lin committed
95
            bias=False,
96
97
98
99
100
101
            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
102
            bias=False,
103
104
105
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
Junyang Lin's avatar
Junyang Lin committed
106
        if hidden_act != "silu":
107
108
109
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
Junyang Lin's avatar
Junyang Lin committed
110
111
112
113
114
115
116
117
118
119
        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):
120
    def __init__(
121
122
123
124
        self,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
125
        rope_parameters: dict[str, Any],
126
        max_position: int = 4096 * 32,
127
128
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
129
130
        prefix: str = "",
        attn_type: str = AttentionType.DECODER,
131
        dual_chunk_attention_config: dict[str, Any] | None = None,
132
133
        qk_norm: bool = False,
        rms_norm_eps: float = 1e-6,
134
    ) -> None:
Junyang Lin's avatar
Junyang Lin committed
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
        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
155
        self.dual_chunk_attention_config = dual_chunk_attention_config
156
        self.qk_norm = qk_norm
Junyang Lin's avatar
Junyang Lin committed
157
158
159
160
161
162
163

        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=True,
164
            quant_config=quant_config,
165
            prefix=f"{prefix}.qkv_proj",
Junyang Lin's avatar
Junyang Lin committed
166
167
168
169
170
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
171
            quant_config=quant_config,
172
            prefix=f"{prefix}.o_proj",
Junyang Lin's avatar
Junyang Lin committed
173
174
        )

175
176
177
178
179
        # QK Normalization support (used in BAGEL and some other models)
        if self.qk_norm:
            self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
            self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

Junyang Lin's avatar
Junyang Lin committed
180
181
182
        self.rotary_emb = get_rope(
            self.head_dim,
            max_position=max_position,
183
            rope_parameters=rope_parameters,
184
            dual_chunk_attention_config=dual_chunk_attention_config,
Junyang Lin's avatar
Junyang Lin committed
185
        )
186
187
188
189
190
        attn_cls = (
            EncoderOnlyAttention
            if attn_type == AttentionType.ENCODER_ONLY
            else Attention
        )
191
        self.attn = attn_cls(
192
193
194
195
196
197
198
199
200
201
202
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            quant_config=quant_config,
            attn_type=attn_type,
            prefix=f"{prefix}.attn",
            **{
                "layer_idx": extract_layer_index(prefix),
                "dual_chunk_attention_config": dual_chunk_attention_config,
203
204
205
206
            }
            if dual_chunk_attention_config
            else {},
        )
Junyang Lin's avatar
Junyang Lin committed
207

208
209
210
211
        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
212
213
214
215
216
217
218

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
zhuwenwen's avatar
zhuwenwen committed
219
220
        # if os.environ.get('FA_PAD') == '1' and self.quant_method is None:
        #     qkv = qkv[...,:-32]
Junyang Lin's avatar
Junyang Lin committed
221
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238

        # Apply QK normalization if enabled (before RoPE)
        if self.qk_norm:
            # Reshape to apply per-head normalization
            # q shape: (total_tokens, q_size) -> (total_tokens, num_heads, head_dim)
            total_tokens = q.shape[0]
            q = q.view(total_tokens, self.num_heads, self.head_dim)
            k = k.view(total_tokens, self.num_kv_heads, self.head_dim)

            # Apply normalization
            q = self.q_norm(q)
            k = self.k_norm(k)

            # Reshape back
            q = q.view(total_tokens, self.q_size)
            k = k.view(total_tokens, self.kv_size)

Junyang Lin's avatar
Junyang Lin committed
239
        q, k = self.rotary_emb(positions, q, k)
240
        attn_output = self.attn(q, k, v)
Junyang Lin's avatar
Junyang Lin committed
241
242
243
244
245
246
247
248
        output, _ = self.o_proj(attn_output)
        return output


class Qwen2DecoderLayer(nn.Module):
    def __init__(
        self,
        config: Qwen2Config,
249
250
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
251
        prefix: str = "",
Junyang Lin's avatar
Junyang Lin committed
252
253
254
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
255
        set_default_rope_theta(config, default_theta=1000000)
256
257
258
        dual_chunk_attention_config = getattr(
            config, "dual_chunk_attention_config", None
        )
259
260
261
262
263
264
265
266
267
268

        # 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):
            attn_type = AttentionType.DECODER
        else:
            attn_type = AttentionType.ENCODER_ONLY

269
270
271
        # Check if QK normalization is enabled (used in BAGEL and some other models)
        qk_norm = getattr(config, "qk_norm", False)

Junyang Lin's avatar
Junyang Lin committed
272
273
274
275
276
        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,
277
            cache_config=cache_config,
278
            quant_config=quant_config,
279
            rope_parameters=config.rope_parameters,
280
            prefix=f"{prefix}.self_attn",
281
            attn_type=attn_type,
282
            dual_chunk_attention_config=dual_chunk_attention_config,
283
284
            qk_norm=qk_norm,
            rms_norm_eps=config.rms_norm_eps,
285
        )
Junyang Lin's avatar
Junyang Lin committed
286
287
288
289
        self.mlp = Qwen2MLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
290
            quant_config=quant_config,
291
            prefix=f"{prefix}.mlp",
Junyang Lin's avatar
Junyang Lin committed
292
        )
293
294
295
296
        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
        )
Junyang Lin's avatar
Junyang Lin committed
297
298
299
300
301

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
302
        residual: torch.Tensor | None,
303
    ) -> tuple[torch.Tensor, torch.Tensor]:
Junyang Lin's avatar
Junyang Lin committed
304
305
306
307
308
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
309
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
Junyang Lin's avatar
Junyang Lin committed
310
311
312
313
314
315
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
        )

        # Fully Connected
316
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
Junyang Lin's avatar
Junyang Lin committed
317
318
319
320
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
def qwen_2_model_invariants(
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    intermediate_tensors: IntermediateTensors | None = None,
    inputs_embeds: torch.Tensor | None = None,
):
    """Shape invariants for Qwen2Model Model, those are translated to
    runtime assertions for unbacked dynamic shapes and are compiled away for
    backed"""
    # All these should be equal.
    # input_ids.size()[0]
    # positions.size()[-1]
    # intermediate_tensors["hidden_states"].size()[0]
    # inputs_embeds.size()[0]
    torch._check(input_ids.size()[0] == positions.size()[-1])
    if intermediate_tensors is not None:
        torch._check(
            input_ids.size()[0] == intermediate_tensors["hidden_states"].size()[0]
        )

    if inputs_embeds is not None:
        torch._check(input_ids.size()[0] == inputs_embeds.size()[0])

    # Hidden dimensions should match (hidden_size)
    # intermediate_tensors["hidden_states"].size()[1]
    # inputs_embeds.size()[1]
    if inputs_embeds is not None and intermediate_tensors is not None:
        torch._check(
            inputs_embeds.size()[1] == intermediate_tensors["hidden_states"].size()[1]
        )


353
354
355
356
357
358
359
360
@support_torch_compile(
    dynamic_arg_dims={
        "input_ids": 0,
        # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
        # otherwise (seq_len, ).
        "positions": -1,
        "intermediate_tensors": 0,
        "inputs_embeds": 0,
361
362
    },
    shape_invariants=qwen_2_model_invariants,
363
)
Junyang Lin's avatar
Junyang Lin committed
364
class Qwen2Model(nn.Module):
365
366
367
368
369
370
371
    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer,
    ):
Junyang Lin's avatar
Junyang Lin committed
372
        super().__init__()
373

374
        config = vllm_config.model_config.hf_config.get_text_config()
375
376
377
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

378
        # TODO (@robertgshaw2): see if this can be moved out
379
        if is_interleaved(vllm_config.model_config.hf_text_config):
380
381
382
383
384
385
386
            assert config.max_window_layers == config.num_hidden_layers, (
                "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,
387
388
                )
            )
389

Junyang Lin's avatar
Junyang Lin committed
390
        self.config = config
391
        self.quant_config = quant_config
Junyang Lin's avatar
Junyang Lin committed
392
393
        self.vocab_size = config.vocab_size

394
395
396
        if get_pp_group().is_first_rank or (
            config.tie_word_embeddings and get_pp_group().is_last_rank
        ):
397
398
399
400
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
401
                prefix=f"{prefix}.embed_tokens",
402
403
404
405
            )
        else:
            self.embed_tokens = PPMissingLayer()

406
407
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
408
409
410
411
412
413
            lambda prefix: decoder_layer_type(
                config=config,
                cache_config=cache_config,
                quant_config=quant_config,
                prefix=prefix,
            ),
414
415
416
            prefix=f"{prefix}.layers",
        )

417
418
419
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
420
421
422
423
        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
424
425
426
427
428
429
430
431
432
433
            
        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'
Junyang Lin's avatar
Junyang Lin committed
434

435
        self.aux_hidden_state_layers = tuple[int, ...]()
436

437
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
438
439
        return self.embed_tokens(input_ids)

Junyang Lin's avatar
Junyang Lin committed
440
441
    def forward(
        self,
442
        input_ids: torch.Tensor | None,
Junyang Lin's avatar
Junyang Lin committed
443
        positions: torch.Tensor,
444
445
446
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
447
448
449
450
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
451
                hidden_states = self.embed_input_ids(input_ids)
452
            residual = None
453
        else:
454
455
456
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]
457
458
459

        aux_hidden_states = []
        for idx, layer in enumerate(
460
461
            islice(self.layers, self.start_layer, self.end_layer)
        ):
462
463
464
465
            if idx in self.aux_hidden_state_layers:
                aux_hidden_states.append(hidden_states + residual)
            hidden_states, residual = layer(positions, hidden_states, residual)

466
        if not get_pp_group().is_last_rank:
467
468
469
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
470

Junyang Lin's avatar
Junyang Lin committed
471
        hidden_states, _ = self.norm(hidden_states, residual)
472
473
474
475

        if len(aux_hidden_states) > 0:
            return hidden_states, aux_hidden_states

Junyang Lin's avatar
Junyang Lin committed
476
477
        return hidden_states

478
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Junyang Lin's avatar
Junyang Lin committed
479
480
481
482
483
484
485
486
        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
487
        params_dict = dict(self.named_parameters(remove_duplicate=False))
488
        loaded_params: set[str] = set()
489
        for name, loaded_weight in weights:
zhuwenwen's avatar
zhuwenwen committed
490
491
492
            if self.use_llama_nn:
                current_count = loaded_weight.current_count 
                total_count = loaded_weight.total_count
Junyang Lin's avatar
Junyang Lin committed
493
494
            if "rotary_emb.inv_freq" in name:
                continue
495
496
497
            if self.quant_config is not None and (
                scale_name := self.quant_config.get_cache_scale(name)
            ):
498
                # Loading kv cache quantization scales
499
                param = params_dict[scale_name]
500
501
502
503
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                loaded_weight = (
                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
                )
504
505
506
                weight_loader(param, loaded_weight)
                loaded_params.add(scale_name)
                continue
507
            for param_name, weight_name, shard_id in stacked_params_mapping:
Junyang Lin's avatar
Junyang Lin committed
508
509
510
511
512
513
                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
514
515
                if is_pp_missing_parameter(name, self):
                    continue
516
517
518
519
520
                if name.endswith("scale"):
                    # Remapping the name of FP8 kv-scale.
                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue
Junyang Lin's avatar
Junyang Lin committed
521
                param = params_dict[name]
522
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
523
524
525
526
                if weight_loader == default_weight_loader:
                    weight_loader(param, loaded_weight)
                else:
                    weight_loader(param, loaded_weight, shard_id)
Junyang Lin's avatar
Junyang Lin committed
527
528
529
530
531
                break
            else:
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
532
                # Remapping the name of FP8 kv-scale.
533
534
535
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue
536
537
                if is_pp_missing_parameter(name, self):
                    continue
538
539
                if name not in params_dict:
                    continue
Junyang Lin's avatar
Junyang Lin committed
540
                param = params_dict[name]
541
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
Junyang Lin's avatar
Junyang Lin committed
542
                weight_loader(param, loaded_weight)
543
            loaded_params.add(name)
zhuwenwen's avatar
zhuwenwen committed
544
            
zhuwenwen's avatar
zhuwenwen committed
545
        if self.use_llama_nn and self.quant_method is None and current_count==total_count:
gaoqiong's avatar
gaoqiong committed
546
547
548
549
            lay_key_words = [
                "self_attn.qkv_proj.weight",
                "self_attn.o_proj.weight",
                "mlp.gate_up_proj.weight",
550
                "mlp.down_proj.weight",
gaoqiong's avatar
gaoqiong committed
551
552
553
            ]
            combined_words = "|".join(lay_key_words)
            
zhuwenwen's avatar
zhuwenwen committed
554
555
            # lay_qkv_words = ["self_attn.qkv_proj.weight"]   
            # qkv_words = "|".join(lay_qkv_words)  
zhuwenwen's avatar
zhuwenwen committed
556
            
zhuwenwen's avatar
zhuwenwen committed
557
558
            # lay_qkv_bias_words = ["self_attn.qkv_proj.bias"]   
            # qkv_bias_words = "|".join(lay_qkv_bias_words) 
zhuwenwen's avatar
zhuwenwen committed
559
            
zhuwenwen's avatar
zhuwenwen committed
560
            # for layername, weight in params_dict.items():
zhuwenwen's avatar
zhuwenwen committed
561
562
            # for layername in loaded_params:
            for layername in params_dict.keys():
zhuwenwen's avatar
zhuwenwen committed
563
                weight = params_dict[layername]
564
565
566
567
568
569
                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
570
571
                # if self.use_fa_pad and (re.findall(qkv_bias_words, layername)):
                #     weight.data = pad_weight(weight.data, 32)
zhuwenwen's avatar
zhuwenwen committed
572
                    
gaoqiong's avatar
gaoqiong committed
573
                matches = re.findall(combined_words, layername)
zhuwenwen's avatar
zhuwenwen committed
574
                if matches:   
zhuwenwen's avatar
zhuwenwen committed
575
576
                    # 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
577
                    
zhuwenwen's avatar
zhuwenwen committed
578
579
580
                    # 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)
581
                        
gaoqiong's avatar
gaoqiong committed
582
583
584
                    _weight = torch.zeros_like(weight.data)
                    ori_shape =_weight.shape
                    
zhuwenwen's avatar
zhuwenwen committed
585
                    ops.trans_w16_gemm(_weight, weight.data, _weight.shape[0], _weight.shape[1])
gaoqiong's avatar
gaoqiong committed
586
587
588
                    weight.data.copy_(_weight)
                    
                    weight.data=weight.data.reshape(ori_shape[1],-1)
589
        else:
zhuwenwen's avatar
zhuwenwen committed
590
            os.environ['LM_NN'] = '0'
591
            os.environ['LLAMA_NN'] = '0'
zhuwenwen's avatar
zhuwenwen committed
592
                    
593
        return loaded_params
594

Junyang Lin's avatar
Junyang Lin committed
595

596
class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
597
598
599
600
601
602
603
604
605
606
607
608
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

609
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
610
        super().__init__()
611
        config = vllm_config.model_config.hf_config.get_text_config()
612
        quant_config = vllm_config.quant_config
613

Junyang Lin's avatar
Junyang Lin committed
614
        self.config = config
615

616
        self.quant_config = quant_config
617
618
619
        self.model = Qwen2Model(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
620

621
622
623
624
        if get_pp_group().is_last_rank:
            if config.tie_word_embeddings:
                self.lm_head = self.model.embed_tokens
            else:
625
626
627
628
629
630
                self.lm_head = ParallelLMHead(
                    config.vocab_size,
                    config.hidden_size,
                    quant_config=quant_config,
                    prefix=maybe_prefix(prefix, "lm_head"),
                )
631
        else:
632
            self.lm_head = PPMissingLayer()
633

634
        self.logits_processor = LogitsProcessor(config.vocab_size)
zhuwenwen's avatar
zhuwenwen committed
635
636
637
638
639
640
641
642
643
644
        
        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'
645

646
        self.make_empty_intermediate_tensors = (
647
648
            self.model.make_empty_intermediate_tensors
        )
Junyang Lin's avatar
Junyang Lin committed
649

650
651
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
652

653
    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
654
655
        self.model.aux_hidden_state_layers = layers

656
    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
657
658
659
        num_layers = len(self.model.layers)
        return (2, num_layers // 2, num_layers - 3)

Junyang Lin's avatar
Junyang Lin committed
660
661
    def forward(
        self,
662
        input_ids: torch.Tensor | None,
Junyang Lin's avatar
Junyang Lin committed
663
        positions: torch.Tensor,
664
665
666
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
667
668
669
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
Junyang Lin's avatar
Junyang Lin committed
670
671
        return hidden_states

672
673
674
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
675
    ) -> torch.Tensor | None:
676
        logits = self.logits_processor(self.lm_head, hidden_states)
677
678
        return logits

679
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
680
681
        loader = AutoWeightsLoader(
            self,
682
            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
683
        )
684
        return loader.load_weights(weights)