qwen.py 9.98 KB
Newer Older
Qing's avatar
Qing committed
1
2
3
4
5
# coding=utf-8
# Adapted from
# https://huggingface.co/Qwen/Qwen-7B/blob/main/modeling_qwen.py
# Copyright (c) Alibaba Cloud.
# LICENSE: https://huggingface.co/Qwen/Qwen-7B/blob/main/LICENSE
Woosuk Kwon's avatar
Woosuk Kwon committed
6
"""Inference-only QWen model compatible with HuggingFace weights."""
7
from typing import Any, Dict, Iterable, List, Optional, Tuple
Qing's avatar
Qing committed
8

9
10
import torch
from torch import nn
11
from transformers import PretrainedConfig
Qing's avatar
Qing committed
12

13
from vllm.attention import Attention, AttentionMetadata
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.model_executor.layers.activation import SiluAndMul
16
from vllm.model_executor.layers.layernorm import RMSNorm
17
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
18
19
                                               QKVParallelLinear,
                                               RowParallelLinear)
20
from vllm.model_executor.layers.logits_processor import LogitsProcessor
21
22
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
23
from vllm.model_executor.layers.rotary_embedding import get_rope
24
25
from vllm.model_executor.layers.sampler import Sampler
from vllm.model_executor.layers.vocab_parallel_embedding import (
26
    ParallelLMHead, VocabParallelEmbedding)
27
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
28
29
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.sequence import SamplerOutput
Qing's avatar
Qing committed
30

31
32
33
34
35
36
37
38

class QWenMLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str = "silu",
39
        quant_config: Optional[QuantizationConfig] = None,
40
41
42
43
44
    ):
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size, [intermediate_size] * 2,
            bias=False,
45
            quant_config=quant_config)
46
47
48
        self.c_proj = RowParallelLinear(intermediate_size,
                                        hidden_size,
                                        bias=False,
49
                                        quant_config=quant_config)
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
        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.c_proj(x)
        return x


class QWenAttention(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        max_position_embeddings: int,
        rope_theta: float = 10000,
        rope_scaling: Optional[Dict[str, Any]] = None,
71
        quant_config: Optional[QuantizationConfig] = None,
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
    ):
        super().__init__()
        self.hidden_size = hidden_size
        tensor_model_parallel_world_size = get_tensor_model_parallel_world_size(
        )
        self.total_num_heads = num_heads
        assert self.total_num_heads % tensor_model_parallel_world_size == 0
        self.num_heads = (self.total_num_heads //
                          tensor_model_parallel_world_size)
        self.head_dim = hidden_size // self.total_num_heads
        self.c_attn = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            bias=True,
87
            quant_config=quant_config,
88
89
90
91
92
        )
        self.c_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
93
            quant_config=quant_config,
94
95
96
97
98
99
100
101
102
103
        )
        self.scaling = self.head_dim**-0.5

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position_embeddings,
            base=rope_theta,
            rope_scaling=rope_scaling,
        )
104
        self.attn = Attention(self.num_heads, self.head_dim, self.scaling)
105
106
107
108
109

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
110
111
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
112
113
114
115
    ) -> torch.Tensor:
        qkv, _ = self.c_attn(hidden_states)
        q, k, v = qkv.chunk(chunks=3, dim=-1)
        q, k = self.rotary_emb(positions, q, k)
116
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
117
118
119
120
121
122
123
124
        output, _ = self.c_proj(attn_output)
        return output


class QWenBlock(nn.Module):

    def __init__(
        self,
125
        config: PretrainedConfig,
126
        quant_config: Optional[QuantizationConfig] = None,
127
128
129
130
131
132
133
134
135
136
137
    ):
        super().__init__()
        self.ln_1 = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
        self.attn = QWenAttention(config.hidden_size,
                                  config.num_attention_heads,
                                  config.max_position_embeddings,
                                  rope_theta=rope_theta,
                                  rope_scaling=rope_scaling,
138
                                  quant_config=quant_config)
139
140
141
142
143

        self.ln_2 = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

        self.mlp = QWenMLP(config.hidden_size,
                           config.intermediate_size // 2,
144
                           quant_config=quant_config)
145
146
147
148
149

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
150
151
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
152
153
154
155
156
157
158
159
160
161
162
163
        residual: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.ln_1(hidden_states)
        else:
            hidden_states, residual = self.ln_1(hidden_states, residual)
        hidden_states = self.attn(
            positions=positions,
            hidden_states=hidden_states,
            kv_cache=kv_cache,
164
            attn_metadata=attn_metadata,
165
166
167
168
169
170
171
172
173
        )

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


class QWenModel(nn.Module):
Qing's avatar
Qing committed
174

175
176
    def __init__(
        self,
177
        config: PretrainedConfig,
178
        quant_config: Optional[QuantizationConfig] = None,
179
180
181
182
183
184
185
186
187
188
    ):
        super().__init__()
        self.config = config
        self.vocab_size = config.vocab_size

        self.wte = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
        )
        self.h = nn.ModuleList([
189
            QWenBlock(config, quant_config)
190
191
192
193
194
195
196
197
            for _ in range(config.num_hidden_layers)
        ])
        self.ln_f = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
198
199
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
200
201
202
203
204
205
206
207
208
    ) -> torch.Tensor:
        hidden_states = self.wte(input_ids)
        residual = None
        for i in range(len(self.h)):
            layer = self.h[i]
            hidden_states, residual = layer(
                positions,
                hidden_states,
                kv_caches[i],
209
                attn_metadata,
210
211
212
213
214
215
216
217
218
219
                residual,
            )
        hidden_states, _ = self.ln_f(hidden_states, residual)
        return hidden_states


class QWenLMHeadModel(nn.Module):

    def __init__(
        self,
220
        config: PretrainedConfig,
221
        quant_config: Optional[QuantizationConfig] = None,
222
223
224
    ):
        super().__init__()
        self.config = config
225
226
        self.quant_config = quant_config
        self.transformer = QWenModel(config, quant_config)
227
        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
228
229
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()
230
231
232
233
234

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
235
236
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
237
238
    ) -> torch.Tensor:
        hidden_states = self.transformer(input_ids, positions, kv_caches,
239
                                         attn_metadata)
240
241
        return hidden_states

242
243
244
245
246
247
    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

248
249
    def sample(
        self,
250
        logits: torch.Tensor,
251
252
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
253
        next_tokens = self.sampler(logits, sampling_metadata)
254
        return next_tokens
Qing's avatar
Qing committed
255

256
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
257
258
259
260
261
262
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "w2", 0),
            ("gate_up_proj", "w1", 1),
        ]
        params_dict = dict(self.named_parameters())
263
        for name, loaded_weight in weights:
Qing's avatar
Qing committed
264
265
            if "rotary_emb.inv_freq" in name:
                continue
266
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
Qing's avatar
Qing committed
267
268
                if weight_name not in name:
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
269
270
271
272
273
                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]
274
275
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
Qing's avatar
Qing committed
276
                break
277
            else:
CHU Tianxiang's avatar
CHU Tianxiang committed
278
279
280
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
281
282
283
284
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)