We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Support Lora for qwen2


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `vllm/model_executor/models/qwen2.py`
Content:
```
1 # coding=utf-8
2 # Adapted from
3 # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py
4 # Copyright 2024 The Qwen team.
5 # Copyright 2023 The vLLM team.
6 # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
7 #
8 # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
9 # and OPT implementations in this library. It has been modified from its
10 # original forms to accommodate minor architectural differences compared
11 # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
12 #
13 # Licensed under the Apache License, Version 2.0 (the "License");
14 # you may not use this file except in compliance with the License.
15 # You may obtain a copy of the License at
16 #
17 #     http://www.apache.org/licenses/LICENSE-2.0
18 #
19 # Unless required by applicable law or agreed to in writing, software
20 # distributed under the License is distributed on an "AS IS" BASIS,
21 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22 # See the License for the specific language governing permissions and
23 # limitations under the License.
24 """Inference-only Qwen2 model compatible with HuggingFace weights."""
25 from typing import List, Optional, Tuple
26 
27 import torch
28 from torch import nn
29 from transformers import Qwen2Config
30 
31 from vllm.model_executor.input_metadata import InputMetadata
32 from vllm.model_executor.layers.activation import SiluAndMul
33 from vllm.model_executor.layers.attention import PagedAttention
34 from vllm.model_executor.layers.layernorm import RMSNorm
35 from vllm.model_executor.layers.linear import (LinearMethodBase,
36                                                MergedColumnParallelLinear,
37                                                QKVParallelLinear,
38                                                RowParallelLinear)
39 from vllm.model_executor.layers.rotary_embedding import get_rope
40 from vllm.model_executor.layers.sampler import Sampler
41 from vllm.model_executor.layers.vocab_parallel_embedding import (
42     VocabParallelEmbedding, ParallelLMHead)
43 from vllm.model_executor.parallel_utils.parallel_state import (
44     get_tensor_model_parallel_world_size)
45 from vllm.model_executor.sampling_metadata import SamplingMetadata
46 from vllm.model_executor.weight_utils import (default_weight_loader,
47                                               hf_model_weights_iterator)
48 from vllm.sequence import SamplerOutput
49 
50 KVCache = Tuple[torch.Tensor, torch.Tensor]
51 
52 
53 class Qwen2MLP(nn.Module):
54 
55     def __init__(
56         self,
57         hidden_size: int,
58         intermediate_size: int,
59         hidden_act: str,
60         linear_method: Optional[LinearMethodBase] = None,
61     ) -> None:
62         super().__init__()
63         self.gate_up_proj = MergedColumnParallelLinear(
64             hidden_size, [intermediate_size] * 2,
65             bias=False,
66             linear_method=linear_method)
67         self.down_proj = RowParallelLinear(intermediate_size,
68                                            hidden_size,
69                                            bias=False,
70                                            linear_method=linear_method)
71         if hidden_act != "silu":
72             raise ValueError(f"Unsupported activation: {hidden_act}. "
73                              "Only silu is supported for now.")
74         self.act_fn = SiluAndMul()
75 
76     def forward(self, x):
77         gate_up, _ = self.gate_up_proj(x)
78         x = self.act_fn(gate_up)
79         x, _ = self.down_proj(x)
80         return x
81 
82 
83 class Qwen2Attention(nn.Module):
84 
85     def __init__(self,
86                  hidden_size: int,
87                  num_heads: int,
88                  num_kv_heads: int,
89                  max_position: int = 4096 * 32,
90                  rope_theta: float = 10000,
91                  use_sliding_window: bool = False,
92                  linear_method: Optional[LinearMethodBase] = None,
93                  sliding_window: Optional[int] = None) -> None:
94         super().__init__()
95         self.hidden_size = hidden_size
96         tp_size = get_tensor_model_parallel_world_size()
97         self.total_num_heads = num_heads
98         assert self.total_num_heads % tp_size == 0
99         self.num_heads = self.total_num_heads // tp_size
100         self.total_num_kv_heads = num_kv_heads
101         if self.total_num_kv_heads >= tp_size:
102             # Number of KV heads is greater than TP size, so we partition
103             # the KV heads across multiple tensor parallel GPUs.
104             assert self.total_num_kv_heads % tp_size == 0
105         else:
106             # Number of KV heads is less than TP size, so we replicate
107             # the KV heads across multiple tensor parallel GPUs.
108             assert tp_size % self.total_num_kv_heads == 0
109         self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
110         self.head_dim = hidden_size // self.total_num_heads
111         self.q_size = self.num_heads * self.head_dim
112         self.kv_size = self.num_kv_heads * self.head_dim
113         self.scaling = self.head_dim**-0.5
114         self.rope_theta = rope_theta
115         self.sliding_window = sliding_window if use_sliding_window else None
116 
117         self.qkv_proj = QKVParallelLinear(
118             hidden_size,
119             self.head_dim,
120             self.total_num_heads,
121             self.total_num_kv_heads,
122             bias=True,
123             linear_method=linear_method,
124         )
125         self.o_proj = RowParallelLinear(
126             self.total_num_heads * self.head_dim,
127             hidden_size,
128             bias=False,
129             linear_method=linear_method,
130         )
131 
132         self.rotary_emb = get_rope(
133             self.head_dim,
134             rotary_dim=self.head_dim,
135             max_position=max_position,
136             base=self.rope_theta,
137         )
138         self.attn = PagedAttention(self.num_heads,
139                                    self.head_dim,
140                                    self.scaling,
141                                    num_kv_heads=self.num_kv_heads,
142                                    sliding_window=self.sliding_window)
143 
144     def forward(
145         self,
146         positions: torch.Tensor,
147         hidden_states: torch.Tensor,
148         kv_cache: KVCache,
149         input_metadata: InputMetadata,
150     ) -> torch.Tensor:
151         qkv, _ = self.qkv_proj(hidden_states)
152         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
153         q, k = self.rotary_emb(positions, q, k)
154         k_cache, v_cache = kv_cache
155         attn_output = self.attn(q, k, v, k_cache, v_cache, input_metadata)
156         output, _ = self.o_proj(attn_output)
157         return output
158 
159 
160 class Qwen2DecoderLayer(nn.Module):
161 
162     def __init__(
163         self,
164         config: Qwen2Config,
165         layer_idx: int,
166         linear_method: Optional[LinearMethodBase] = None,
167     ) -> None:
168         super().__init__()
169         self.hidden_size = config.hidden_size
170         # Requires transformers > 4.32.0
171         rope_theta = getattr(config, "rope_theta", 1000000)
172         use_sliding_window = config.use_sliding_window and layer_idx < config.max_window_layers
173         self.self_attn = Qwen2Attention(
174             hidden_size=self.hidden_size,
175             num_heads=config.num_attention_heads,
176             max_position=config.max_position_embeddings,
177             num_kv_heads=config.num_key_value_heads,
178             rope_theta=rope_theta,
179             use_sliding_window=use_sliding_window,
180             linear_method=linear_method,
181             sliding_window=config.sliding_window)
182         self.mlp = Qwen2MLP(
183             hidden_size=self.hidden_size,
184             intermediate_size=config.intermediate_size,
185             hidden_act=config.hidden_act,
186             linear_method=linear_method,
187         )
188         self.input_layernorm = RMSNorm(config.hidden_size,
189                                        eps=config.rms_norm_eps)
190         self.post_attention_layernorm = RMSNorm(config.hidden_size,
191                                                 eps=config.rms_norm_eps)
192 
193     def forward(
194         self,
195         positions: torch.Tensor,
196         hidden_states: torch.Tensor,
197         kv_cache: KVCache,
198         input_metadata: InputMetadata,
199         residual: Optional[torch.Tensor],
200     ) -> Tuple[torch.Tensor, torch.Tensor]:
201         # Self Attention
202         if residual is None:
203             residual = hidden_states
204             hidden_states = self.input_layernorm(hidden_states)
205         else:
206             hidden_states, residual = self.input_layernorm(
207                 hidden_states, residual)
208         hidden_states = self.self_attn(
209             positions=positions,
210             hidden_states=hidden_states,
211             kv_cache=kv_cache,
212             input_metadata=input_metadata,
213         )
214 
215         # Fully Connected
216         hidden_states, residual = self.post_attention_layernorm(
217             hidden_states, residual)
218         hidden_states = self.mlp(hidden_states)
219         return hidden_states, residual
220 
221 
222 class Qwen2Model(nn.Module):
223 
224     def __init__(
225         self,
226         config: Qwen2Config,
227         linear_method: Optional[LinearMethodBase] = None,
228     ) -> None:
229         super().__init__()
230         self.config = config
231         self.padding_idx = config.pad_token_id
232         self.vocab_size = config.vocab_size
233 
234         self.embed_tokens = VocabParallelEmbedding(
235             config.vocab_size,
236             config.hidden_size,
237         )
238         self.layers = nn.ModuleList([
239             Qwen2DecoderLayer(config, layer_idx, linear_method)
240             for layer_idx in range(config.num_hidden_layers)
241         ])
242         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
243 
244     def forward(
245         self,
246         input_ids: torch.Tensor,
247         positions: torch.Tensor,
248         kv_caches: List[KVCache],
249         input_metadata: InputMetadata,
250     ) -> torch.Tensor:
251         hidden_states = self.embed_tokens(input_ids)
252         residual = None
253         for i in range(len(self.layers)):
254             layer = self.layers[i]
255             hidden_states, residual = layer(
256                 positions,
257                 hidden_states,
258                 kv_caches[i],
259                 input_metadata,
260                 residual,
261             )
262         hidden_states, _ = self.norm(hidden_states, residual)
263         return hidden_states
264 
265 
266 class Qwen2ForCausalLM(nn.Module):
267 
268     def __init__(
269         self,
270         config: Qwen2Config,
271         linear_method: Optional[LinearMethodBase] = None,
272     ) -> None:
273         super().__init__()
274         self.config = config
275         self.linear_method = linear_method
276         self.model = Qwen2Model(config, linear_method)
277         self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
278         self.sampler = Sampler(config.vocab_size)
279 
280     def forward(
281         self,
282         input_ids: torch.Tensor,
283         positions: torch.Tensor,
284         kv_caches: List[KVCache],
285         input_metadata: InputMetadata,
286     ) -> torch.Tensor:
287         hidden_states = self.model(input_ids, positions, kv_caches,
288                                    input_metadata)
289         return hidden_states
290 
291     def sample(
292         self,
293         hidden_states: torch.Tensor,
294         sampling_metadata: SamplingMetadata,
295     ) -> Optional[SamplerOutput]:
296         next_tokens = self.sampler(self.lm_head.weight, hidden_states,
297                                    sampling_metadata)
298         return next_tokens
299 
300     def load_weights(self,
301                      model_name_or_path: str,
302                      cache_dir: Optional[str] = None,
303                      load_format: str = "auto",
304                      revision: Optional[str] = None):
305         stacked_params_mapping = [
306             # (param_name, shard_name, shard_id)
307             ("qkv_proj", "q_proj", "q"),
308             ("qkv_proj", "k_proj", "k"),
309             ("qkv_proj", "v_proj", "v"),
310             ("gate_up_proj", "gate_proj", 0),
311             ("gate_up_proj", "up_proj", 1),
312         ]
313         params_dict = dict(self.named_parameters())
314         for name, loaded_weight in hf_model_weights_iterator(
315                 model_name_or_path, cache_dir, load_format, revision):
316             if "rotary_emb.inv_freq" in name:
317                 continue
318             for (param_name, weight_name, shard_id) in stacked_params_mapping:
319                 if weight_name not in name:
320                     continue
321                 name = name.replace(weight_name, param_name)
322                 # Skip loading extra bias for GPTQ models.
323                 if name.endswith(".bias") and name not in params_dict:
324                     continue
325                 param = params_dict[name]
326                 weight_loader = param.weight_loader
327                 weight_loader(param, loaded_weight, shard_id)
328                 break
329             else:
330                 # Skip loading extra bias for GPTQ models.
331                 if name.endswith(".bias") and name not in params_dict:
332                     continue
333                 param = params_dict[name]
334                 weight_loader = getattr(param, "weight_loader",
335                                         default_weight_loader)
336                 weight_loader(param, loaded_weight)
337
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py
--- a/vllm/model_executor/models/qwen2.py
+++ b/vllm/model_executor/models/qwen2.py
@@ -46,6 +46,7 @@
 from vllm.model_executor.weight_utils import (default_weight_loader,
                                               hf_model_weights_iterator)
 from vllm.sequence import SamplerOutput
+from vllm.config import LoRAConfig
 
 KVCache = Tuple[torch.Tensor, torch.Tensor]
 
@@ -264,12 +265,35 @@
 
 
 class Qwen2ForCausalLM(nn.Module):
+    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 = []
 
     def __init__(
         self,
         config: Qwen2Config,
         linear_method: Optional[LinearMethodBase] = None,
+        lora_config: Optional[LoRAConfig] = None,
     ) -> None:
+        del lora_config
         super().__init__()
         self.config = config
         self.linear_method = linear_method

{"golden_diff": "diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py\n--- a/vllm/model_executor/models/qwen2.py\n+++ b/vllm/model_executor/models/qwen2.py\n@@ -46,6 +46,7 @@\n from vllm.model_executor.weight_utils import (default_weight_loader,\n                                               hf_model_weights_iterator)\n from vllm.sequence import SamplerOutput\n+from vllm.config import LoRAConfig\n \n KVCache = Tuple[torch.Tensor, torch.Tensor]\n \n@@ -264,12 +265,35 @@\n \n \n class Qwen2ForCausalLM(nn.Module):\n+    packed_modules_mapping = {\n+        \"qkv_proj\": [\n+            \"q_proj\",\n+            \"k_proj\",\n+            \"v_proj\",\n+        ],\n+        \"gate_up_proj\": [\n+            \"gate_proj\",\n+            \"up_proj\",\n+        ],\n+    }\n+\n+    # LoRA specific attributes\n+    supported_lora_modules = [\n+        \"qkv_proj\",\n+        \"o_proj\",\n+        \"gate_up_proj\",\n+        \"down_proj\",\n+    ]\n+    embedding_modules = {}\n+    embedding_padding_modules = []\n \n     def __init__(\n         self,\n         config: Qwen2Config,\n         linear_method: Optional[LinearMethodBase] = None,\n+        lora_config: Optional[LoRAConfig] = None,\n     ) -> None:\n+        del lora_config\n         super().__init__()\n         self.config = config\n         self.linear_method = linear_method\n", "issue": "Support Lora for qwen2\n\n", "before_files": [{"content": "# coding=utf-8\n# Adapted from\n# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py\n# Copyright 2024 The Qwen team.\n# Copyright 2023 The vLLM team.\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX\n# and OPT implementations in this library. It has been modified from its\n# original forms to accommodate minor architectural differences compared\n# to GPT-NeoX and OPT used by the Meta AI team that trained the model.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Inference-only Qwen2 model compatible with HuggingFace weights.\"\"\"\nfrom typing import List, Optional, Tuple\n\nimport torch\nfrom torch import nn\nfrom transformers import Qwen2Config\n\nfrom vllm.model_executor.input_metadata import InputMetadata\nfrom vllm.model_executor.layers.activation import SiluAndMul\nfrom vllm.model_executor.layers.attention import PagedAttention\nfrom vllm.model_executor.layers.layernorm import RMSNorm\nfrom vllm.model_executor.layers.linear import (LinearMethodBase,\n                                               MergedColumnParallelLinear,\n                                               QKVParallelLinear,\n                                               RowParallelLinear)\nfrom vllm.model_executor.layers.rotary_embedding import get_rope\nfrom vllm.model_executor.layers.sampler import Sampler\nfrom vllm.model_executor.layers.vocab_parallel_embedding import (\n    VocabParallelEmbedding, ParallelLMHead)\nfrom vllm.model_executor.parallel_utils.parallel_state import (\n    get_tensor_model_parallel_world_size)\nfrom vllm.model_executor.sampling_metadata import SamplingMetadata\nfrom vllm.model_executor.weight_utils import (default_weight_loader,\n                                              hf_model_weights_iterator)\nfrom vllm.sequence import SamplerOutput\n\nKVCache = Tuple[torch.Tensor, torch.Tensor]\n\n\nclass Qwen2MLP(nn.Module):\n\n    def __init__(\n        self,\n        hidden_size: int,\n        intermediate_size: int,\n        hidden_act: str,\n        linear_method: Optional[LinearMethodBase] = None,\n    ) -> None:\n        super().__init__()\n        self.gate_up_proj = MergedColumnParallelLinear(\n            hidden_size, [intermediate_size] * 2,\n            bias=False,\n            linear_method=linear_method)\n        self.down_proj = RowParallelLinear(intermediate_size,\n                                           hidden_size,\n                                           bias=False,\n                                           linear_method=linear_method)\n        if hidden_act != \"silu\":\n            raise ValueError(f\"Unsupported activation: {hidden_act}. \"\n                             \"Only silu is supported for now.\")\n        self.act_fn = SiluAndMul()\n\n    def forward(self, x):\n        gate_up, _ = self.gate_up_proj(x)\n        x = self.act_fn(gate_up)\n        x, _ = self.down_proj(x)\n        return x\n\n\nclass Qwen2Attention(nn.Module):\n\n    def __init__(self,\n                 hidden_size: int,\n                 num_heads: int,\n                 num_kv_heads: int,\n                 max_position: int = 4096 * 32,\n                 rope_theta: float = 10000,\n                 use_sliding_window: bool = False,\n                 linear_method: Optional[LinearMethodBase] = None,\n                 sliding_window: Optional[int] = None) -> None:\n        super().__init__()\n        self.hidden_size = hidden_size\n        tp_size = get_tensor_model_parallel_world_size()\n        self.total_num_heads = num_heads\n        assert self.total_num_heads % tp_size == 0\n        self.num_heads = self.total_num_heads // tp_size\n        self.total_num_kv_heads = num_kv_heads\n        if self.total_num_kv_heads >= tp_size:\n            # Number of KV heads is greater than TP size, so we partition\n            # the KV heads across multiple tensor parallel GPUs.\n            assert self.total_num_kv_heads % tp_size == 0\n        else:\n            # Number of KV heads is less than TP size, so we replicate\n            # the KV heads across multiple tensor parallel GPUs.\n            assert tp_size % self.total_num_kv_heads == 0\n        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)\n        self.head_dim = hidden_size // self.total_num_heads\n        self.q_size = self.num_heads * self.head_dim\n        self.kv_size = self.num_kv_heads * self.head_dim\n        self.scaling = self.head_dim**-0.5\n        self.rope_theta = rope_theta\n        self.sliding_window = sliding_window if use_sliding_window else None\n\n        self.qkv_proj = QKVParallelLinear(\n            hidden_size,\n            self.head_dim,\n            self.total_num_heads,\n            self.total_num_kv_heads,\n            bias=True,\n            linear_method=linear_method,\n        )\n        self.o_proj = RowParallelLinear(\n            self.total_num_heads * self.head_dim,\n            hidden_size,\n            bias=False,\n            linear_method=linear_method,\n        )\n\n        self.rotary_emb = get_rope(\n            self.head_dim,\n            rotary_dim=self.head_dim,\n            max_position=max_position,\n            base=self.rope_theta,\n        )\n        self.attn = PagedAttention(self.num_heads,\n                                   self.head_dim,\n                                   self.scaling,\n                                   num_kv_heads=self.num_kv_heads,\n                                   sliding_window=self.sliding_window)\n\n    def forward(\n        self,\n        positions: torch.Tensor,\n        hidden_states: torch.Tensor,\n        kv_cache: KVCache,\n        input_metadata: InputMetadata,\n    ) -> torch.Tensor:\n        qkv, _ = self.qkv_proj(hidden_states)\n        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)\n        q, k = self.rotary_emb(positions, q, k)\n        k_cache, v_cache = kv_cache\n        attn_output = self.attn(q, k, v, k_cache, v_cache, input_metadata)\n        output, _ = self.o_proj(attn_output)\n        return output\n\n\nclass Qwen2DecoderLayer(nn.Module):\n\n    def __init__(\n        self,\n        config: Qwen2Config,\n        layer_idx: int,\n        linear_method: Optional[LinearMethodBase] = None,\n    ) -> None:\n        super().__init__()\n        self.hidden_size = config.hidden_size\n        # Requires transformers > 4.32.0\n        rope_theta = getattr(config, \"rope_theta\", 1000000)\n        use_sliding_window = config.use_sliding_window and layer_idx < config.max_window_layers\n        self.self_attn = Qwen2Attention(\n            hidden_size=self.hidden_size,\n            num_heads=config.num_attention_heads,\n            max_position=config.max_position_embeddings,\n            num_kv_heads=config.num_key_value_heads,\n            rope_theta=rope_theta,\n            use_sliding_window=use_sliding_window,\n            linear_method=linear_method,\n            sliding_window=config.sliding_window)\n        self.mlp = Qwen2MLP(\n            hidden_size=self.hidden_size,\n            intermediate_size=config.intermediate_size,\n            hidden_act=config.hidden_act,\n            linear_method=linear_method,\n        )\n        self.input_layernorm = RMSNorm(config.hidden_size,\n                                       eps=config.rms_norm_eps)\n        self.post_attention_layernorm = RMSNorm(config.hidden_size,\n                                                eps=config.rms_norm_eps)\n\n    def forward(\n        self,\n        positions: torch.Tensor,\n        hidden_states: torch.Tensor,\n        kv_cache: KVCache,\n        input_metadata: InputMetadata,\n        residual: Optional[torch.Tensor],\n    ) -> Tuple[torch.Tensor, torch.Tensor]:\n        # Self Attention\n        if residual is None:\n            residual = hidden_states\n            hidden_states = self.input_layernorm(hidden_states)\n        else:\n            hidden_states, residual = self.input_layernorm(\n                hidden_states, residual)\n        hidden_states = self.self_attn(\n            positions=positions,\n            hidden_states=hidden_states,\n            kv_cache=kv_cache,\n            input_metadata=input_metadata,\n        )\n\n        # Fully Connected\n        hidden_states, residual = self.post_attention_layernorm(\n            hidden_states, residual)\n        hidden_states = self.mlp(hidden_states)\n        return hidden_states, residual\n\n\nclass Qwen2Model(nn.Module):\n\n    def __init__(\n        self,\n        config: Qwen2Config,\n        linear_method: Optional[LinearMethodBase] = None,\n    ) -> None:\n        super().__init__()\n        self.config = config\n        self.padding_idx = config.pad_token_id\n        self.vocab_size = config.vocab_size\n\n        self.embed_tokens = VocabParallelEmbedding(\n            config.vocab_size,\n            config.hidden_size,\n        )\n        self.layers = nn.ModuleList([\n            Qwen2DecoderLayer(config, layer_idx, linear_method)\n            for layer_idx in range(config.num_hidden_layers)\n        ])\n        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        positions: torch.Tensor,\n        kv_caches: List[KVCache],\n        input_metadata: InputMetadata,\n    ) -> torch.Tensor:\n        hidden_states = self.embed_tokens(input_ids)\n        residual = None\n        for i in range(len(self.layers)):\n            layer = self.layers[i]\n            hidden_states, residual = layer(\n                positions,\n                hidden_states,\n                kv_caches[i],\n                input_metadata,\n                residual,\n            )\n        hidden_states, _ = self.norm(hidden_states, residual)\n        return hidden_states\n\n\nclass Qwen2ForCausalLM(nn.Module):\n\n    def __init__(\n        self,\n        config: Qwen2Config,\n        linear_method: Optional[LinearMethodBase] = None,\n    ) -> None:\n        super().__init__()\n        self.config = config\n        self.linear_method = linear_method\n        self.model = Qwen2Model(config, linear_method)\n        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)\n        self.sampler = Sampler(config.vocab_size)\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        positions: torch.Tensor,\n        kv_caches: List[KVCache],\n        input_metadata: InputMetadata,\n    ) -> torch.Tensor:\n        hidden_states = self.model(input_ids, positions, kv_caches,\n                                   input_metadata)\n        return hidden_states\n\n    def sample(\n        self,\n        hidden_states: torch.Tensor,\n        sampling_metadata: SamplingMetadata,\n    ) -> Optional[SamplerOutput]:\n        next_tokens = self.sampler(self.lm_head.weight, hidden_states,\n                                   sampling_metadata)\n        return next_tokens\n\n    def load_weights(self,\n                     model_name_or_path: str,\n                     cache_dir: Optional[str] = None,\n                     load_format: str = \"auto\",\n                     revision: Optional[str] = None):\n        stacked_params_mapping = [\n            # (param_name, shard_name, shard_id)\n            (\"qkv_proj\", \"q_proj\", \"q\"),\n            (\"qkv_proj\", \"k_proj\", \"k\"),\n            (\"qkv_proj\", \"v_proj\", \"v\"),\n            (\"gate_up_proj\", \"gate_proj\", 0),\n            (\"gate_up_proj\", \"up_proj\", 1),\n        ]\n        params_dict = dict(self.named_parameters())\n        for name, loaded_weight in hf_model_weights_iterator(\n                model_name_or_path, cache_dir, load_format, revision):\n            if \"rotary_emb.inv_freq\" in name:\n                continue\n            for (param_name, weight_name, shard_id) in stacked_params_mapping:\n                if weight_name not in name:\n                    continue\n                name = name.replace(weight_name, param_name)\n                # Skip loading extra bias for GPTQ models.\n                if name.endswith(\".bias\") and name not in params_dict:\n                    continue\n                param = params_dict[name]\n                weight_loader = param.weight_loader\n                weight_loader(param, loaded_weight, shard_id)\n                break\n            else:\n                # Skip loading extra bias for GPTQ models.\n                if name.endswith(\".bias\") and name not in params_dict:\n                    continue\n                param = params_dict[name]\n                weight_loader = getattr(param, \"weight_loader\",\n                                        default_weight_loader)\n                weight_loader(param, loaded_weight)\n", "path": "vllm/model_executor/models/qwen2.py"}], "after_files": [{"content": "# coding=utf-8\n# Adapted from\n# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py\n# Copyright 2024 The Qwen team.\n# Copyright 2023 The vLLM team.\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX\n# and OPT implementations in this library. It has been modified from its\n# original forms to accommodate minor architectural differences compared\n# to GPT-NeoX and OPT used by the Meta AI team that trained the model.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Inference-only Qwen2 model compatible with HuggingFace weights.\"\"\"\nfrom typing import List, Optional, Tuple\n\nimport torch\nfrom torch import nn\nfrom transformers import Qwen2Config\n\nfrom vllm.model_executor.input_metadata import InputMetadata\nfrom vllm.model_executor.layers.activation import SiluAndMul\nfrom vllm.model_executor.layers.attention import PagedAttention\nfrom vllm.model_executor.layers.layernorm import RMSNorm\nfrom vllm.model_executor.layers.linear import (LinearMethodBase,\n                                               MergedColumnParallelLinear,\n                                               QKVParallelLinear,\n                                               RowParallelLinear)\nfrom vllm.model_executor.layers.rotary_embedding import get_rope\nfrom vllm.model_executor.layers.sampler import Sampler\nfrom vllm.model_executor.layers.vocab_parallel_embedding import (\n    VocabParallelEmbedding, ParallelLMHead)\nfrom vllm.model_executor.parallel_utils.parallel_state import (\n    get_tensor_model_parallel_world_size)\nfrom vllm.model_executor.sampling_metadata import SamplingMetadata\nfrom vllm.model_executor.weight_utils import (default_weight_loader,\n                                              hf_model_weights_iterator)\nfrom vllm.sequence import SamplerOutput\nfrom vllm.config import LoRAConfig\n\nKVCache = Tuple[torch.Tensor, torch.Tensor]\n\n\nclass Qwen2MLP(nn.Module):\n\n    def __init__(\n        self,\n        hidden_size: int,\n        intermediate_size: int,\n        hidden_act: str,\n        linear_method: Optional[LinearMethodBase] = None,\n    ) -> None:\n        super().__init__()\n        self.gate_up_proj = MergedColumnParallelLinear(\n            hidden_size, [intermediate_size] * 2,\n            bias=False,\n            linear_method=linear_method)\n        self.down_proj = RowParallelLinear(intermediate_size,\n                                           hidden_size,\n                                           bias=False,\n                                           linear_method=linear_method)\n        if hidden_act != \"silu\":\n            raise ValueError(f\"Unsupported activation: {hidden_act}. \"\n                             \"Only silu is supported for now.\")\n        self.act_fn = SiluAndMul()\n\n    def forward(self, x):\n        gate_up, _ = self.gate_up_proj(x)\n        x = self.act_fn(gate_up)\n        x, _ = self.down_proj(x)\n        return x\n\n\nclass Qwen2Attention(nn.Module):\n\n    def __init__(self,\n                 hidden_size: int,\n                 num_heads: int,\n                 num_kv_heads: int,\n                 max_position: int = 4096 * 32,\n                 rope_theta: float = 10000,\n                 use_sliding_window: bool = False,\n                 linear_method: Optional[LinearMethodBase] = None,\n                 sliding_window: Optional[int] = None) -> None:\n        super().__init__()\n        self.hidden_size = hidden_size\n        tp_size = get_tensor_model_parallel_world_size()\n        self.total_num_heads = num_heads\n        assert self.total_num_heads % tp_size == 0\n        self.num_heads = self.total_num_heads // tp_size\n        self.total_num_kv_heads = num_kv_heads\n        if self.total_num_kv_heads >= tp_size:\n            # Number of KV heads is greater than TP size, so we partition\n            # the KV heads across multiple tensor parallel GPUs.\n            assert self.total_num_kv_heads % tp_size == 0\n        else:\n            # Number of KV heads is less than TP size, so we replicate\n            # the KV heads across multiple tensor parallel GPUs.\n            assert tp_size % self.total_num_kv_heads == 0\n        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)\n        self.head_dim = hidden_size // self.total_num_heads\n        self.q_size = self.num_heads * self.head_dim\n        self.kv_size = self.num_kv_heads * self.head_dim\n        self.scaling = self.head_dim**-0.5\n        self.rope_theta = rope_theta\n        self.sliding_window = sliding_window if use_sliding_window else None\n\n        self.qkv_proj = QKVParallelLinear(\n            hidden_size,\n            self.head_dim,\n            self.total_num_heads,\n            self.total_num_kv_heads,\n            bias=True,\n            linear_method=linear_method,\n        )\n        self.o_proj = RowParallelLinear(\n            self.total_num_heads * self.head_dim,\n            hidden_size,\n            bias=False,\n            linear_method=linear_method,\n        )\n\n        self.rotary_emb = get_rope(\n            self.head_dim,\n            rotary_dim=self.head_dim,\n            max_position=max_position,\n            base=self.rope_theta,\n        )\n        self.attn = PagedAttention(self.num_heads,\n                                   self.head_dim,\n                                   self.scaling,\n                                   num_kv_heads=self.num_kv_heads,\n                                   sliding_window=self.sliding_window)\n\n    def forward(\n        self,\n        positions: torch.Tensor,\n        hidden_states: torch.Tensor,\n        kv_cache: KVCache,\n        input_metadata: InputMetadata,\n    ) -> torch.Tensor:\n        qkv, _ = self.qkv_proj(hidden_states)\n        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)\n        q, k = self.rotary_emb(positions, q, k)\n        k_cache, v_cache = kv_cache\n        attn_output = self.attn(q, k, v, k_cache, v_cache, input_metadata)\n        output, _ = self.o_proj(attn_output)\n        return output\n\n\nclass Qwen2DecoderLayer(nn.Module):\n\n    def __init__(\n        self,\n        config: Qwen2Config,\n        layer_idx: int,\n        linear_method: Optional[LinearMethodBase] = None,\n    ) -> None:\n        super().__init__()\n        self.hidden_size = config.hidden_size\n        # Requires transformers > 4.32.0\n        rope_theta = getattr(config, \"rope_theta\", 1000000)\n        use_sliding_window = config.use_sliding_window and layer_idx < config.max_window_layers\n        self.self_attn = Qwen2Attention(\n            hidden_size=self.hidden_size,\n            num_heads=config.num_attention_heads,\n            max_position=config.max_position_embeddings,\n            num_kv_heads=config.num_key_value_heads,\n            rope_theta=rope_theta,\n            use_sliding_window=use_sliding_window,\n            linear_method=linear_method,\n            sliding_window=config.sliding_window)\n        self.mlp = Qwen2MLP(\n            hidden_size=self.hidden_size,\n            intermediate_size=config.intermediate_size,\n            hidden_act=config.hidden_act,\n            linear_method=linear_method,\n        )\n        self.input_layernorm = RMSNorm(config.hidden_size,\n                                       eps=config.rms_norm_eps)\n        self.post_attention_layernorm = RMSNorm(config.hidden_size,\n                                                eps=config.rms_norm_eps)\n\n    def forward(\n        self,\n        positions: torch.Tensor,\n        hidden_states: torch.Tensor,\n        kv_cache: KVCache,\n        input_metadata: InputMetadata,\n        residual: Optional[torch.Tensor],\n    ) -> Tuple[torch.Tensor, torch.Tensor]:\n        # Self Attention\n        if residual is None:\n            residual = hidden_states\n            hidden_states = self.input_layernorm(hidden_states)\n        else:\n            hidden_states, residual = self.input_layernorm(\n                hidden_states, residual)\n        hidden_states = self.self_attn(\n            positions=positions,\n            hidden_states=hidden_states,\n            kv_cache=kv_cache,\n            input_metadata=input_metadata,\n        )\n\n        # Fully Connected\n        hidden_states, residual = self.post_attention_layernorm(\n            hidden_states, residual)\n        hidden_states = self.mlp(hidden_states)\n        return hidden_states, residual\n\n\nclass Qwen2Model(nn.Module):\n\n    def __init__(\n        self,\n        config: Qwen2Config,\n        linear_method: Optional[LinearMethodBase] = None,\n    ) -> None:\n        super().__init__()\n        self.config = config\n        self.padding_idx = config.pad_token_id\n        self.vocab_size = config.vocab_size\n\n        self.embed_tokens = VocabParallelEmbedding(\n            config.vocab_size,\n            config.hidden_size,\n        )\n        self.layers = nn.ModuleList([\n            Qwen2DecoderLayer(config, layer_idx, linear_method)\n            for layer_idx in range(config.num_hidden_layers)\n        ])\n        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        positions: torch.Tensor,\n        kv_caches: List[KVCache],\n        input_metadata: InputMetadata,\n    ) -> torch.Tensor:\n        hidden_states = self.embed_tokens(input_ids)\n        residual = None\n        for i in range(len(self.layers)):\n            layer = self.layers[i]\n            hidden_states, residual = layer(\n                positions,\n                hidden_states,\n                kv_caches[i],\n                input_metadata,\n                residual,\n            )\n        hidden_states, _ = self.norm(hidden_states, residual)\n        return hidden_states\n\n\nclass Qwen2ForCausalLM(nn.Module):\n    packed_modules_mapping = {\n        \"qkv_proj\": [\n            \"q_proj\",\n            \"k_proj\",\n            \"v_proj\",\n        ],\n        \"gate_up_proj\": [\n            \"gate_proj\",\n            \"up_proj\",\n        ],\n    }\n\n    # LoRA specific attributes\n    supported_lora_modules = [\n        \"qkv_proj\",\n        \"o_proj\",\n        \"gate_up_proj\",\n        \"down_proj\",\n    ]\n    embedding_modules = {}\n    embedding_padding_modules = []\n\n    def __init__(\n        self,\n        config: Qwen2Config,\n        linear_method: Optional[LinearMethodBase] = None,\n        lora_config: Optional[LoRAConfig] = None,\n    ) -> None:\n        del lora_config\n        super().__init__()\n        self.config = config\n        self.linear_method = linear_method\n        self.model = Qwen2Model(config, linear_method)\n        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)\n        self.sampler = Sampler(config.vocab_size)\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        positions: torch.Tensor,\n        kv_caches: List[KVCache],\n        input_metadata: InputMetadata,\n    ) -> torch.Tensor:\n        hidden_states = self.model(input_ids, positions, kv_caches,\n                                   input_metadata)\n        return hidden_states\n\n    def sample(\n        self,\n        hidden_states: torch.Tensor,\n        sampling_metadata: SamplingMetadata,\n    ) -> Optional[SamplerOutput]:\n        next_tokens = self.sampler(self.lm_head.weight, hidden_states,\n                                   sampling_metadata)\n        return next_tokens\n\n    def load_weights(self,\n                     model_name_or_path: str,\n                     cache_dir: Optional[str] = None,\n                     load_format: str = \"auto\",\n                     revision: Optional[str] = None):\n        stacked_params_mapping = [\n            # (param_name, shard_name, shard_id)\n            (\"qkv_proj\", \"q_proj\", \"q\"),\n            (\"qkv_proj\", \"k_proj\", \"k\"),\n            (\"qkv_proj\", \"v_proj\", \"v\"),\n            (\"gate_up_proj\", \"gate_proj\", 0),\n            (\"gate_up_proj\", \"up_proj\", 1),\n        ]\n        params_dict = dict(self.named_parameters())\n        for name, loaded_weight in hf_model_weights_iterator(\n                model_name_or_path, cache_dir, load_format, revision):\n            if \"rotary_emb.inv_freq\" in name:\n                continue\n            for (param_name, weight_name, shard_id) in stacked_params_mapping:\n                if weight_name not in name:\n                    continue\n                name = name.replace(weight_name, param_name)\n                # Skip loading extra bias for GPTQ models.\n                if name.endswith(\".bias\") and name not in params_dict:\n                    continue\n                param = params_dict[name]\n                weight_loader = param.weight_loader\n                weight_loader(param, loaded_weight, shard_id)\n                break\n            else:\n                # Skip loading extra bias for GPTQ models.\n                if name.endswith(\".bias\") and name not in params_dict:\n                    continue\n                param = params_dict[name]\n                weight_loader = getattr(param, \"weight_loader\",\n                                        default_weight_loader)\n                weight_loader(param, loaded_weight)\n", "path": "vllm/model_executor/models/qwen2.py"}]}