qwen2.py 20.7 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

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

63
from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
64
65
66
67
68
69
70
71
72
from .utils import (
    AutoWeightsLoader,
    PPMissingLayer,
    extract_layer_index,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
73

Junyang Lin's avatar
Junyang Lin committed
74
75
76
77
78
79
80

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

        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=True,
154
            quant_config=quant_config,
155
            prefix=f"{prefix}.qkv_proj",
Junyang Lin's avatar
Junyang Lin committed
156
157
158
159
160
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
161
            quant_config=quant_config,
162
            prefix=f"{prefix}.o_proj",
Junyang Lin's avatar
Junyang Lin committed
163
164
165
166
167
168
        )

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
169
            rope_parameters=rope_parameters,
170
            dual_chunk_attention_config=dual_chunk_attention_config,
Junyang Lin's avatar
Junyang Lin committed
171
        )
172
173
174
175
176
        attn_cls = (
            EncoderOnlyAttention
            if attn_type == AttentionType.ENCODER_ONLY
            else Attention
        )
177
        self.attn = attn_cls(
178
179
180
181
182
183
184
185
186
187
188
            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,
189
190
191
192
            }
            if dual_chunk_attention_config
            else {},
        )
Junyang Lin's avatar
Junyang Lin committed
193
194
195
196
197
198
199
200
201

    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)
202
        attn_output = self.attn(q, k, v)
Junyang Lin's avatar
Junyang Lin committed
203
204
205
206
207
208
209
210
        output, _ = self.o_proj(attn_output)
        return output


class Qwen2DecoderLayer(nn.Module):
    def __init__(
        self,
        config: Qwen2Config,
211
212
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
213
        prefix: str = "",
Junyang Lin's avatar
Junyang Lin committed
214
215
216
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
217
        set_default_rope_theta(config, default_theta=1000000)
218
219
220
        dual_chunk_attention_config = getattr(
            config, "dual_chunk_attention_config", None
        )
221
222
223
224
225
226
227
228
229
230

        # 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
231
232
233
234
235
        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,
236
            cache_config=cache_config,
237
            quant_config=quant_config,
238
            rope_parameters=config.rope_parameters,
239
            prefix=f"{prefix}.self_attn",
240
            attn_type=attn_type,
241
            dual_chunk_attention_config=dual_chunk_attention_config,
242
        )
Junyang Lin's avatar
Junyang Lin committed
243
244
245
246
        self.mlp = Qwen2MLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
247
            quant_config=quant_config,
248
            prefix=f"{prefix}.mlp",
Junyang Lin's avatar
Junyang Lin committed
249
        )
250
251
252
253
        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
254
255
256
257
258

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
259
        residual: torch.Tensor | None,
260
    ) -> tuple[torch.Tensor, torch.Tensor]:
Junyang Lin's avatar
Junyang Lin committed
261
262
263
264
265
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
266
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
Junyang Lin's avatar
Junyang Lin committed
267
268
269
270
271
272
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
        )

        # Fully Connected
273
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
Junyang Lin's avatar
Junyang Lin committed
274
275
276
277
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
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]
        )


310
311
312
313
314
315
316
317
@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,
318
319
    },
    shape_invariants=qwen_2_model_invariants,
320
)
Junyang Lin's avatar
Junyang Lin committed
321
class Qwen2Model(nn.Module):
322
323
324
325
326
327
328
    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer,
    ):
Junyang Lin's avatar
Junyang Lin committed
329
        super().__init__()
330

331
        config = vllm_config.model_config.hf_config.get_text_config()
332
333
334
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

335
        # TODO (@robertgshaw2): see if this can be moved out
336
        if is_interleaved(vllm_config.model_config.hf_text_config):
337
338
339
340
341
342
343
            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,
344
345
                )
            )
346

Junyang Lin's avatar
Junyang Lin committed
347
        self.config = config
348
        self.quant_config = quant_config
Junyang Lin's avatar
Junyang Lin committed
349
350
        self.vocab_size = config.vocab_size

351
352
353
        if get_pp_group().is_first_rank or (
            config.tie_word_embeddings and get_pp_group().is_last_rank
        ):
354
355
356
357
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
358
                prefix=f"{prefix}.embed_tokens",
359
360
361
362
            )
        else:
            self.embed_tokens = PPMissingLayer()

363
364
        # Use the provided decoder layer type or default to Qwen2DecoderLayer
        decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer
365
366
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
367
368
369
370
371
372
            lambda prefix: decoder_layer_type(
                config=config,
                cache_config=cache_config,
                quant_config=quant_config,
                prefix=prefix,
            ),
373
374
375
            prefix=f"{prefix}.layers",
        )

376
377
378
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
379
380
381
382
        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
383

384
        self.aux_hidden_state_layers = tuple[int, ...]()
385

386
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
387
388
        return self.embed_tokens(input_ids)

Junyang Lin's avatar
Junyang Lin committed
389
390
391
392
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
393
394
395
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
396
397
398
399
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
400
                hidden_states = self.embed_input_ids(input_ids)
401
            residual = None
402
        else:
403
404
405
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]
406
407
408

        aux_hidden_states = []
        for idx, layer in enumerate(
409
410
            islice(self.layers, self.start_layer, self.end_layer)
        ):
411
412
413
414
            if idx in self.aux_hidden_state_layers:
                aux_hidden_states.append(hidden_states + residual)
            hidden_states, residual = layer(positions, hidden_states, residual)

415
        if not get_pp_group().is_last_rank:
416
417
418
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
419

Junyang Lin's avatar
Junyang Lin committed
420
        hidden_states, _ = self.norm(hidden_states, residual)
421
422
423
424

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

Junyang Lin's avatar
Junyang Lin committed
425
426
        return hidden_states

427
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
428
429
430
431
432
433
434
435
436
        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))
437
        loaded_params: set[str] = set()
438
439
440
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
441
442
443
            if self.quant_config is not None and (
                scale_name := self.quant_config.get_cache_scale(name)
            ):
444
                # Loading kv cache quantization scales
445
                param = params_dict[scale_name]
446
447
448
449
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                loaded_weight = (
                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
                )
450
451
452
                weight_loader(param, loaded_weight)
                loaded_params.add(scale_name)
                continue
453
            for param_name, weight_name, shard_id in stacked_params_mapping:
454
455
456
457
458
459
460
461
                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
462
463
464
465
466
                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
467
                param = params_dict[name]
468
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
469
470
471
472
                if weight_loader == default_weight_loader:
                    weight_loader(param, loaded_weight)
                else:
                    weight_loader(param, loaded_weight, shard_id)
473
474
475
476
477
478
479
480
481
482
483
484
                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]
485
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
486
                weight_loader(param, loaded_weight)
487
488
            loaded_params.add(name)
        return loaded_params
489

Junyang Lin's avatar
Junyang Lin committed
490

491
class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
492
493
494
495
496
497
498
499
500
501
502
503
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

504
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
505
506
507
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
508

Junyang Lin's avatar
Junyang Lin committed
509
        self.config = config
510

511
        self.quant_config = quant_config
512
513
514
        self.model = Qwen2Model(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
515

516
517
518
519
        if get_pp_group().is_last_rank:
            if config.tie_word_embeddings:
                self.lm_head = self.model.embed_tokens
            else:
520
521
522
523
524
525
                self.lm_head = ParallelLMHead(
                    config.vocab_size,
                    config.hidden_size,
                    quant_config=quant_config,
                    prefix=maybe_prefix(prefix, "lm_head"),
                )
526
        else:
527
            self.lm_head = PPMissingLayer()
528

529
        self.logits_processor = LogitsProcessor(config.vocab_size)
530

531
        self.make_empty_intermediate_tensors = (
532
533
            self.model.make_empty_intermediate_tensors
        )
Junyang Lin's avatar
Junyang Lin committed
534

535
536
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
537

538
    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
539
540
        self.model.aux_hidden_state_layers = layers

541
    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
542
543
544
        num_layers = len(self.model.layers)
        return (2, num_layers // 2, num_layers - 3)

Junyang Lin's avatar
Junyang Lin committed
545
546
547
548
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
549
550
551
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
552
553
554
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
Junyang Lin's avatar
Junyang Lin committed
555
556
        return hidden_states

557
558
559
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
560
    ) -> torch.Tensor | None:
561
        logits = self.logits_processor(self.lm_head, hidden_states)
562
563
        return logits

564
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
565
566
        loader = AutoWeightsLoader(
            self,
567
            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
568
        )
569
        return loader.load_weights(weights)