qwen2.py 21.3 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

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

import torch
from torch import nn
from transformers import Qwen2Config

32
from vllm.attention import Attention, 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
    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
from .interfaces import SupportsLoRA, SupportsPP
56
57
from .utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,
                    is_pp_missing_parameter,
58
59
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
60

61
62
logger = init_logger(__name__)

Junyang Lin's avatar
Junyang Lin committed
63
64
65
66
67
68
69
70

class Qwen2MLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
71
        quant_config: Optional[QuantizationConfig] = None,
72
        prefix: str = "",
Junyang Lin's avatar
Junyang Lin committed
73
74
75
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
76
77
78
79
80
81
82
83
84
            hidden_size,
            [intermediate_size] * 2,
            bias=False,
            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
85
            bias=False,
86
87
88
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
Junyang Lin's avatar
Junyang Lin committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
        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,
109
                 cache_config: Optional[CacheConfig] = None,
110
                 quant_config: Optional[QuantizationConfig] = None,
111
                 rope_scaling: Optional[Tuple] = None,
112
113
                 prefix: str = "",
                 attn_type: str = AttentionType.DECODER) -> None:
Junyang Lin's avatar
Junyang Lin committed
114
115
116
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
        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,
142
            quant_config=quant_config,
143
            prefix=f"{prefix}.qkv_proj",
Junyang Lin's avatar
Junyang Lin committed
144
145
146
147
148
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
149
            quant_config=quant_config,
150
            prefix=f"{prefix}.o_proj",
Junyang Lin's avatar
Junyang Lin committed
151
152
153
154
155
156
157
        )

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

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> 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)
177
        attn_output = self.attn(q, k, v)
Junyang Lin's avatar
Junyang Lin committed
178
179
180
181
182
183
184
185
186
        output, _ = self.o_proj(attn_output)
        return output


class Qwen2DecoderLayer(nn.Module):

    def __init__(
        self,
        config: Qwen2Config,
187
        cache_config: Optional[CacheConfig] = None,
188
        quant_config: Optional[QuantizationConfig] = None,
189
        prefix: str = "",
Junyang Lin's avatar
Junyang Lin committed
190
191
192
193
194
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        # Requires transformers > 4.32.0
        rope_theta = getattr(config, "rope_theta", 1000000)
195
        rope_scaling = getattr(config, "rope_scaling", None)
196
197
198
199
200
201
202
203
204
205

        # 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

Junyang Lin's avatar
Junyang Lin committed
206
207
208
209
210
211
        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,
212
            cache_config=cache_config,
213
            quant_config=quant_config,
214
215
            rope_scaling=rope_scaling,
            prefix=f"{prefix}.self_attn",
216
            attn_type=attn_type,
217
        )
Junyang Lin's avatar
Junyang Lin committed
218
219
220
221
        self.mlp = Qwen2MLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
222
            quant_config=quant_config,
223
            prefix=f"{prefix}.mlp",
Junyang Lin's avatar
Junyang Lin committed
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
        )
        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,
        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,
        )

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


255
256
257
258
259
260
261
262
263
@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,
    })
Junyang Lin's avatar
Junyang Lin committed
264
265
class Qwen2Model(nn.Module):

266
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Junyang Lin's avatar
Junyang Lin committed
267
        super().__init__()
268
269
270
271
272

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

273
274
275
276
277
278
279
280
281
282
283
284
        # 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
285
        self.config = config
286
        self.quant_config = quant_config
Junyang Lin's avatar
Junyang Lin committed
287
288
        self.vocab_size = config.vocab_size

289
290
291
292
293
294
        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,
295
                prefix=f"{prefix}.embed_tokens",
296
297
298
299
            )
        else:
            self.embed_tokens = PPMissingLayer()

300
301
302
303
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: Qwen2DecoderLayer(config=config,
                                             cache_config=cache_config,
304
                                             quant_config=quant_config,
305
                                             prefix=prefix),
306
307
308
            prefix=f"{prefix}.layers",
        )

309
310
311
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
312
313
314
315
        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()
Junyang Lin's avatar
Junyang Lin committed
316

317
318
319
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Junyang Lin's avatar
Junyang Lin committed
320
321
322
323
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
324
325
        intermediate_tensors: Optional[IntermediateTensors] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
326
    ) -> Union[torch.Tensor, IntermediateTensors]:
327
328
329
330
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
331
                hidden_states = self.get_input_embeddings(input_ids)
332
            residual = None
333
        else:
334
335
336
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]
337
        for layer in self.layers[self.start_layer:self.end_layer]:
Junyang Lin's avatar
Junyang Lin committed
338
339
340
341
342
            hidden_states, residual = layer(
                positions,
                hidden_states,
                residual,
            )
343
344
345
346
347
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
                "residual": residual
            })
Junyang Lin's avatar
Junyang Lin committed
348
349
350
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states

351
352
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
353
354
355
356
357
358
359
360
361
        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),
        ]
        params_dict = dict(self.named_parameters(remove_duplicate=False))
362
        loaded_params: Set[str] = set()
363
364
365
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
366
367
            if (self.quant_config is not None and
                (scale_name := self.quant_config.get_cache_scale(name))):
368
                # Loading kv cache quantization scales
369
370
371
372
373
374
375
376
                param = params_dict[scale_name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
                                 loaded_weight[0])
                weight_loader(param, loaded_weight)
                loaded_params.add(scale_name)
                continue
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
            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
                if is_pp_missing_parameter(name, self):
                    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
                # Remapping the name of FP8 kv-scale.
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue
                if is_pp_missing_parameter(name, self):
                    continue
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
404
405
            loaded_params.add(name)
        return loaded_params
406

Junyang Lin's avatar
Junyang Lin committed
407

408
class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
409
410
411
412
413
414
415
416
417
418
419
420
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

421
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
422
423
424
425
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config
426

Junyang Lin's avatar
Junyang Lin committed
427
        self.config = config
428
429
        self.lora_config = lora_config

430
        self.quant_config = quant_config
431
        self.model = Qwen2Model(vllm_config=vllm_config,
432
                                prefix=maybe_prefix(prefix, "model"))
433

434
435
436
437
438
439
440
441
442
        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"))
443
        else:
444
            self.lm_head = PPMissingLayer()
445

446
        self.logits_processor = LogitsProcessor(config.vocab_size)
Joe Runde's avatar
Joe Runde committed
447
        self.sampler = get_sampler()
448

449
450
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Junyang Lin's avatar
Junyang Lin committed
451

452
453
454
    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
455
456
457
458
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
459
        intermediate_tensors: Optional[IntermediateTensors] = None,
460
        inputs_embeds: Optional[torch.Tensor] = None,
461
    ) -> Union[torch.Tensor, IntermediateTensors]:
462
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
463
                                   inputs_embeds)
Junyang Lin's avatar
Junyang Lin committed
464
465
        return hidden_states

466
467
468
469
470
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
471
        logits = self.logits_processor(self.lm_head, hidden_states,
472
473
474
                                       sampling_metadata)
        return logits

Junyang Lin's avatar
Junyang Lin committed
475
476
    def sample(
        self,
477
        logits: torch.Tensor,
Junyang Lin's avatar
Junyang Lin committed
478
479
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
480
        next_tokens = self.sampler(logits, sampling_metadata)
Junyang Lin's avatar
Junyang Lin committed
481
482
        return next_tokens

483
484
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
485
486
487
488
489
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=(["lm_head."]
                           if self.config.tie_word_embeddings else None),
        )
490
        return loader.load_weights(weights)
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505


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

506
507
    hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})

508
509
510
511
512
513
514
515
516
517
518
519
520
521
    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"))

522
523
        # TODO: Replace this model class with as_embedding_model(
        # Qwen2ForCausalLM) after changing the default pooling method
524
525
526
527
528
529
530
        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.")

531
532
533
534
535
536
537
538
539
540
541
542
        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,
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> torch.Tensor:
543
        return self.model(input_ids, positions, intermediate_tensors)
544
545
546
547
548
549
550
551

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

552
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
553
        weights = self.hf_to_vllm_mapper.apply(weights)
554
555
        weights = ((name, data) for name, data in weights
                   if not name.startswith("lm_head."))
556
        self.model.load_weights(weights)