"tests/kernels/attention/test_attention.py" did not exist on "bf33700ecd6db472c4aeb489c5d42aa47a735198"
qwen2.py 14 KB
Newer Older
Junyang Lin's avatar
Junyang Lin committed
1
2
3
# coding=utf-8
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py
4
5
# Copyright 2024 The Qwen team.
# Copyright 2023 The vLLM team.
Junyang Lin's avatar
Junyang Lin committed
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# 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.
24
"""Inference-only Qwen2 model compatible with HuggingFace weights."""
25
from typing import Iterable, List, Optional, Tuple
Junyang Lin's avatar
Junyang Lin committed
26
27
28
29
30

import torch
from torch import nn
from transformers import Qwen2Config

31
from vllm.attention import Attention, AttentionMetadata
32
from vllm.config import CacheConfig, LoRAConfig
33
from vllm.distributed import get_tensor_model_parallel_world_size
Junyang Lin's avatar
Junyang Lin committed
34
35
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
36
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
Junyang Lin's avatar
Junyang Lin committed
37
38
                                               QKVParallelLinear,
                                               RowParallelLinear)
39
from vllm.model_executor.layers.logits_processor import LogitsProcessor
40
41
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
42
from vllm.model_executor.layers.rotary_embedding import get_rope
Junyang Lin's avatar
Junyang Lin committed
43
44
from vllm.model_executor.layers.sampler import Sampler
from vllm.model_executor.layers.vocab_parallel_embedding import (
45
    ParallelLMHead, VocabParallelEmbedding)
46
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
Junyang Lin's avatar
Junyang Lin committed
47
48
49
50
51
52
53
54
55
56
57
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.sequence import SamplerOutput


class Qwen2MLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
58
        quant_config: Optional[QuantizationConfig] = None,
Junyang Lin's avatar
Junyang Lin committed
59
60
61
62
63
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size, [intermediate_size] * 2,
            bias=False,
64
            quant_config=quant_config)
Junyang Lin's avatar
Junyang Lin committed
65
66
67
        self.down_proj = RowParallelLinear(intermediate_size,
                                           hidden_size,
                                           bias=False,
68
                                           quant_config=quant_config)
Junyang Lin's avatar
Junyang Lin committed
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
        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,
                 use_sliding_window: bool = False,
90
                 cache_config: Optional[CacheConfig] = None,
91
                 quant_config: Optional[QuantizationConfig] = None,
92
93
                 sliding_window: Optional[int] = None,
                 rope_scaling: Optional[Tuple] = None) -> None:
Junyang Lin's avatar
Junyang Lin committed
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
        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.sliding_window = sliding_window if use_sliding_window else None

        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=True,
123
            quant_config=quant_config,
Junyang Lin's avatar
Junyang Lin committed
124
125
126
127
128
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
129
            quant_config=quant_config,
Junyang Lin's avatar
Junyang Lin committed
130
131
132
133
134
135
136
        )

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
            base=self.rope_theta,
137
            rope_scaling=rope_scaling,
Junyang Lin's avatar
Junyang Lin committed
138
        )
139
140
141
142
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
143
144
                              sliding_window=self.sliding_window,
                              cache_config=cache_config)
Junyang Lin's avatar
Junyang Lin committed
145
146
147
148
149

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
150
151
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Junyang Lin's avatar
Junyang Lin committed
152
153
154
155
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
156
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Junyang Lin's avatar
Junyang Lin committed
157
158
159
160
161
162
163
164
165
166
        output, _ = self.o_proj(attn_output)
        return output


class Qwen2DecoderLayer(nn.Module):

    def __init__(
        self,
        config: Qwen2Config,
        layer_idx: int,
167
        cache_config: Optional[CacheConfig] = None,
168
        quant_config: Optional[QuantizationConfig] = None,
Junyang Lin's avatar
Junyang Lin committed
169
170
171
172
173
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        # Requires transformers > 4.32.0
        rope_theta = getattr(config, "rope_theta", 1000000)
174
        rope_scaling = getattr(config, "rope_scaling", None)
175
176
        use_sliding_window = (config.use_sliding_window
                              and layer_idx < config.max_window_layers)
Junyang Lin's avatar
Junyang Lin committed
177
178
179
180
181
182
183
        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,
            use_sliding_window=use_sliding_window,
184
            cache_config=cache_config,
185
            quant_config=quant_config,
186
187
            sliding_window=config.sliding_window,
            rope_scaling=rope_scaling)
Junyang Lin's avatar
Junyang Lin committed
188
189
190
191
        self.mlp = Qwen2MLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
192
            quant_config=quant_config,
Junyang Lin's avatar
Junyang Lin committed
193
194
195
196
197
198
199
200
201
202
        )
        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)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
203
204
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Junyang Lin's avatar
Junyang Lin committed
205
206
207
208
209
210
211
212
213
214
215
216
217
        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,
218
            attn_metadata=attn_metadata,
Junyang Lin's avatar
Junyang Lin committed
219
220
221
222
223
224
225
226
227
228
229
230
231
232
        )

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


class Qwen2Model(nn.Module):

    def __init__(
        self,
        config: Qwen2Config,
233
        cache_config: Optional[CacheConfig] = None,
234
        quant_config: Optional[QuantizationConfig] = None,
Junyang Lin's avatar
Junyang Lin committed
235
236
237
238
239
240
241
242
243
244
245
    ) -> None:
        super().__init__()
        self.config = config
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
        )
        self.layers = nn.ModuleList([
246
            Qwen2DecoderLayer(config, layer_idx, cache_config, quant_config)
Junyang Lin's avatar
Junyang Lin committed
247
248
249
250
251
252
253
254
            for layer_idx in range(config.num_hidden_layers)
        ])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
255
256
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
Junyang Lin's avatar
Junyang Lin committed
257
258
259
260
261
262
263
264
265
    ) -> torch.Tensor:
        hidden_states = self.embed_tokens(input_ids)
        residual = None
        for i in range(len(self.layers)):
            layer = self.layers[i]
            hidden_states, residual = layer(
                positions,
                hidden_states,
                kv_caches[i],
266
                attn_metadata,
Junyang Lin's avatar
Junyang Lin committed
267
268
269
270
271
272
273
                residual,
            )
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


class Qwen2ForCausalLM(nn.Module):
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
    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 = []
Junyang Lin's avatar
Junyang Lin committed
295
296
297
298

    def __init__(
        self,
        config: Qwen2Config,
299
        cache_config: Optional[CacheConfig] = None,
300
        quant_config: Optional[QuantizationConfig] = None,
301
        lora_config: Optional[LoRAConfig] = None,
Junyang Lin's avatar
Junyang Lin committed
302
    ) -> None:
303
        del lora_config
Junyang Lin's avatar
Junyang Lin committed
304
305
        super().__init__()
        self.config = config
306
        self.quant_config = quant_config
307
        self.model = Qwen2Model(config, cache_config, quant_config)
308

309
310
311
        if config.tie_word_embeddings:
            self.lm_head_weight = self.model.embed_tokens.weight
        else:
312
313
            self.lm_head = ParallelLMHead(config.vocab_size,
                                          config.hidden_size)
314
            self.lm_head_weight = self.lm_head.weight
315

316
317
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()
Junyang Lin's avatar
Junyang Lin committed
318
319
320
321
322

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
323
324
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
Junyang Lin's avatar
Junyang Lin committed
325
326
    ) -> torch.Tensor:
        hidden_states = self.model(input_ids, positions, kv_caches,
327
                                   attn_metadata)
Junyang Lin's avatar
Junyang Lin committed
328
329
        return hidden_states

330
331
332
333
334
335
    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
        logits = self.logits_processor(self.lm_head_weight, hidden_states,
                                       sampling_metadata)
        return logits

Junyang Lin's avatar
Junyang Lin committed
336
337
    def sample(
        self,
338
        logits: torch.Tensor,
Junyang Lin's avatar
Junyang Lin committed
339
340
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
341
        next_tokens = self.sampler(logits, sampling_metadata)
Junyang Lin's avatar
Junyang Lin committed
342
343
        return next_tokens

344
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
Junyang Lin's avatar
Junyang Lin committed
345
346
347
348
349
350
351
352
        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
353
        params_dict = dict(self.named_parameters(remove_duplicate=False))
354
        for name, loaded_weight in weights:
Junyang Lin's avatar
Junyang Lin committed
355
356
            if "rotary_emb.inv_freq" in name:
                continue
357
358
            if self.config.tie_word_embeddings and "lm_head.weight" in name:
                continue
Junyang Lin's avatar
Junyang Lin committed
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
            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
                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
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)