We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
KeyError: 'decoder.layers.23.fc2.weight'  when test model "facebook/opt-6.7b"
<img width="1133" alt="image" src="https://github.com/vllm-project/vllm/assets/73372584/23ac48c4-2b1f-42e3-a55f-a252fee7901b">
It looks like the model structure and parameters don't match, how should I fix this 

--- 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/opt.py`
Content:
```
1 # coding=utf-8
2 # Adapted from
3 # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/opt/modeling_opt.py
4 # Copyright 2023 The vLLM team.
5 # Copyright 2022 The Fairseq Authors and The HuggingFace Inc. team. All rights
6 # reserved.
7 #
8 # Licensed under the Apache License, Version 2.0 (the "License");
9 # you may not use this file except in compliance with the License.
10 # You may obtain a copy of the License at
11 #
12 #     http://www.apache.org/licenses/LICENSE-2.0
13 #
14 # Unless required by applicable law or agreed to in writing, software
15 # distributed under the License is distributed on an "AS IS" BASIS,
16 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 # See the License for the specific language governing permissions and
18 # limitations under the License.
19 """Inference-only OPT model compatible with HuggingFace weights."""
20 from typing import List, Optional, Tuple
21 
22 import torch
23 from torch import nn
24 from transformers import OPTConfig
25 
26 from vllm.model_executor.input_metadata import InputMetadata
27 from vllm.model_executor.layers.activation import get_act_fn
28 from vllm.model_executor.layers.attention import PagedAttention
29 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
30                                                LinearMethodBase,
31                                                QKVParallelLinear,
32                                                ReplicatedLinear,
33                                                RowParallelLinear)
34 from vllm.model_executor.layers.sampler import Sampler
35 from vllm.model_executor.layers.vocab_parallel_embedding import (
36     VocabParallelEmbedding)
37 from vllm.model_executor.parallel_utils.parallel_state import (
38     get_tensor_model_parallel_world_size)
39 from vllm.model_executor.weight_utils import (default_weight_loader,
40                                               hf_model_weights_iterator)
41 from vllm.sequence import SamplerOutput
42 
43 KVCache = Tuple[torch.Tensor, torch.Tensor]
44 
45 
46 class OPTLearnedPositionalEmbedding(nn.Embedding):
47 
48     def __init__(self, num_embeddings: int, embedding_dim: int):
49         # OPT is set up so that if padding_idx is specified then offset the
50         # embedding ids by 2 and adjust num_embeddings appropriately. Other
51         # models don't have this hack
52         self.offset = 2
53         super().__init__(num_embeddings + self.offset, embedding_dim)
54 
55     def forward(self, positions: torch.Tensor):
56         return super().forward(positions + self.offset)
57 
58 
59 class OPTAttention(nn.Module):
60 
61     def __init__(
62         self,
63         embed_dim: int,
64         num_heads: int,
65         bias: bool = True,
66         linear_method: Optional[LinearMethodBase] = None,
67     ) -> None:
68         super().__init__()
69         self.embed_dim = embed_dim
70         tensor_model_parallel_world_size = (
71             get_tensor_model_parallel_world_size())
72         total_num_heads = num_heads
73         assert num_heads % tensor_model_parallel_world_size == 0
74         self.num_heads = total_num_heads // tensor_model_parallel_world_size
75         self.head_dim = embed_dim // total_num_heads
76         self.scaling = self.head_dim**-0.5
77 
78         self.qkv_proj = QKVParallelLinear(
79             embed_dim,
80             self.head_dim,
81             total_num_heads,
82             bias=bias,
83             linear_method=linear_method,
84         )
85         self.out_proj = RowParallelLinear(
86             embed_dim,
87             embed_dim,
88             bias=bias,
89             linear_method=linear_method,
90         )
91         self.attn = PagedAttention(self.num_heads,
92                                    self.head_dim,
93                                    scale=self.scaling)
94 
95     def forward(
96         self,
97         hidden_states: torch.Tensor,
98         kv_cache: KVCache,
99         input_metadata: InputMetadata,
100         cache_event: Optional[torch.cuda.Event],
101     ) -> torch.Tensor:
102         qkv, _ = self.qkv_proj(hidden_states)
103         q, k, v = qkv.chunk(chunks=3, dim=-1)
104         key_cache, value_cache = kv_cache
105         attn_output = self.attn(q, k, v, key_cache, value_cache,
106                                 input_metadata, cache_event)
107         output, _ = self.out_proj(attn_output)
108         return output
109 
110 
111 class OPTDecoderLayer(nn.Module):
112 
113     def __init__(
114         self,
115         config: OPTConfig,
116         linear_method: Optional[LinearMethodBase] = None,
117     ):
118         super().__init__()
119         self.config = config
120         self.embed_dim = config.hidden_size
121         self.self_attn = OPTAttention(
122             embed_dim=self.embed_dim,
123             num_heads=config.num_attention_heads,
124             bias=config.enable_bias,
125             linear_method=linear_method,
126         )
127         self.do_layer_norm_before = config.do_layer_norm_before
128 
129         self.self_attn_layer_norm = nn.LayerNorm(
130             self.embed_dim,
131             elementwise_affine=config.layer_norm_elementwise_affine)
132         self.fc1 = ColumnParallelLinear(
133             self.embed_dim,
134             config.ffn_dim,
135             bias=config.enable_bias,
136             linear_method=linear_method,
137         )
138         quant_config = getattr(linear_method, "quant_config", None)
139         self.activation_fn = get_act_fn(config.activation_function,
140                                         quant_config, config.ffn_dim)
141         self.fc2 = RowParallelLinear(
142             config.ffn_dim,
143             self.embed_dim,
144             bias=config.enable_bias,
145             linear_method=linear_method,
146         )
147         self.final_layer_norm = nn.LayerNorm(
148             self.embed_dim,
149             elementwise_affine=config.layer_norm_elementwise_affine)
150 
151     def forward(
152         self,
153         hidden_states: torch.Tensor,
154         kv_cache: KVCache,
155         input_metadata: InputMetadata,
156         cache_event: Optional[torch.cuda.Event],
157     ) -> torch.Tensor:
158         # Self Attention
159         residual = hidden_states
160         # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention
161         if self.do_layer_norm_before:
162             hidden_states = self.self_attn_layer_norm(hidden_states)
163         hidden_states = self.self_attn(hidden_states=hidden_states,
164                                        kv_cache=kv_cache,
165                                        input_metadata=input_metadata,
166                                        cache_event=cache_event)
167         hidden_states = residual + hidden_states
168         # 350m applies layer norm AFTER attention
169         if not self.do_layer_norm_before:
170             hidden_states = self.self_attn_layer_norm(hidden_states)
171 
172         # Fully Connected
173         residual = hidden_states
174         # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention
175         if self.do_layer_norm_before:
176             hidden_states = self.final_layer_norm(hidden_states)
177         hidden_states, _ = self.fc1(hidden_states)
178         hidden_states = self.activation_fn(hidden_states)
179         hidden_states, _ = self.fc2(hidden_states)
180         hidden_states = residual + hidden_states
181         # 350m applies layer norm AFTER attention
182         if not self.do_layer_norm_before:
183             hidden_states = self.final_layer_norm(hidden_states)
184         return hidden_states
185 
186 
187 class OPTDecoder(nn.Module):
188 
189     def __init__(
190         self,
191         config: OPTConfig,
192         linear_method: Optional[LinearMethodBase] = None,
193     ):
194         super().__init__()
195         self.config = config
196         self.padding_idx = config.pad_token_id
197         self.max_target_positions = config.max_position_embeddings
198         self.vocab_size = config.vocab_size
199 
200         self.embed_tokens = VocabParallelEmbedding(
201             config.vocab_size,
202             config.word_embed_proj_dim,
203         )
204         # Positional embeddings are replicated (not sharded).
205         self.embed_positions = OPTLearnedPositionalEmbedding(
206             config.max_position_embeddings, config.hidden_size)
207 
208         # Project out & in will be replicated if they exist.
209         if config.word_embed_proj_dim != config.hidden_size:
210             self.project_out = ReplicatedLinear(config.hidden_size,
211                                                 config.word_embed_proj_dim,
212                                                 bias=False,
213                                                 linear_method=linear_method)
214         else:
215             self.project_out = None
216 
217         if config.word_embed_proj_dim != config.hidden_size:
218             self.project_in = ReplicatedLinear(config.word_embed_proj_dim,
219                                                config.hidden_size,
220                                                bias=False,
221                                                linear_method=linear_method)
222         else:
223             self.project_in = None
224 
225         # Note that the only purpose of `config._remove_final_layer_norm` is to
226         # keep backward compatibility with checkpoints that have been fine-tuned
227         # before transformers v4.20.1
228         # see https://github.com/facebookresearch/metaseq/pull/164
229         if config.do_layer_norm_before and not config._remove_final_layer_norm:
230             self.final_layer_norm = nn.LayerNorm(
231                 config.hidden_size,
232                 elementwise_affine=config.layer_norm_elementwise_affine)
233         else:
234             self.final_layer_norm = None
235 
236         self.layers = nn.ModuleList([
237             OPTDecoderLayer(config, linear_method)
238             for _ in range(config.num_hidden_layers)
239         ])
240 
241     def forward(
242         self,
243         input_ids: torch.Tensor,
244         positions: torch.Tensor,
245         kv_caches: List[KVCache],
246         input_metadata: InputMetadata,
247         cache_events: Optional[List[torch.cuda.Event]],
248     ) -> torch.Tensor:
249         inputs_embeds = self.embed_tokens(input_ids)
250         pos_embeds = self.embed_positions(positions)
251         if self.project_in is not None:
252             inputs_embeds, _ = self.project_in(inputs_embeds)
253         hidden_states = inputs_embeds + pos_embeds
254 
255         for i in range(len(self.layers)):
256             cache_event = None if cache_events is None else cache_events[i]
257             layer = self.layers[i]
258             hidden_states = layer(hidden_states, kv_caches[i], input_metadata,
259                                   cache_event)
260 
261         if self.final_layer_norm is not None:
262             hidden_states = self.final_layer_norm(hidden_states)
263         if self.project_out is not None:
264             hidden_states, _ = self.project_out(hidden_states)
265         return hidden_states
266 
267 
268 class OPTModel(nn.Module):
269 
270     def __init__(
271         self,
272         config: OPTConfig,
273         linear_method: Optional[LinearMethodBase] = None,
274     ):
275         super().__init__()
276         self.decoder = OPTDecoder(config, linear_method)
277 
278     def forward(
279         self,
280         input_ids: torch.Tensor,
281         positions: torch.Tensor,
282         kv_caches: List[KVCache],
283         input_metadata: InputMetadata,
284         cache_events: Optional[List[torch.cuda.Event]],
285     ) -> torch.Tensor:
286         return self.decoder(input_ids, positions, kv_caches, input_metadata,
287                             cache_events)
288 
289 
290 class OPTForCausalLM(nn.Module):
291 
292     def __init__(
293         self,
294         config,
295         linear_method: Optional[LinearMethodBase] = None,
296     ):
297         super().__init__()
298         self.config = config
299         self.linear_method = linear_method
300         self.model = OPTModel(config, linear_method)
301         self.lm_head_weight = self.model.decoder.embed_tokens.weight
302         self.sampler = Sampler(config.vocab_size)
303 
304     def forward(
305         self,
306         input_ids: torch.Tensor,
307         positions: torch.Tensor,
308         kv_caches: List[KVCache],
309         input_metadata: InputMetadata,
310         cache_events: Optional[List[torch.cuda.Event]],
311     ) -> SamplerOutput:
312         hidden_states = self.model(input_ids, positions, kv_caches,
313                                    input_metadata, cache_events)
314         next_tokens = self.sampler(self.lm_head_weight, hidden_states,
315                                    input_metadata)
316         return next_tokens
317 
318     def load_weights(self,
319                      model_name_or_path: str,
320                      cache_dir: Optional[str] = None,
321                      load_format: str = "auto",
322                      revision: Optional[str] = None):
323         stacked_params_mapping = [
324             # (param_name, shard_name, shard_id)
325             ("qkv_proj", "q_proj", "q"),
326             ("qkv_proj", "k_proj", "k"),
327             ("qkv_proj", "v_proj", "v"),
328         ]
329         params_dict = dict(self.named_parameters(remove_duplicate=False))
330         for name, loaded_weight in hf_model_weights_iterator(
331                 model_name_or_path, cache_dir, load_format, revision):
332             if "lm_head.weight" in name:
333                 continue
334             for (param_name, weight_name, shard_id) in stacked_params_mapping:
335                 if weight_name not in name:
336                     continue
337                 param = params_dict[name.replace(weight_name, param_name)]
338                 weight_loader = param.weight_loader
339                 weight_loader(param, loaded_weight, shard_id)
340                 break
341             else:
342                 param = params_dict[name]
343                 weight_loader = getattr(param, "weight_loader",
344                                         default_weight_loader)
345                 weight_loader(param, loaded_weight)
346
```
--- 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/opt.py b/vllm/model_executor/models/opt.py
--- a/vllm/model_executor/models/opt.py
+++ b/vllm/model_executor/models/opt.py
@@ -331,6 +331,9 @@
                 model_name_or_path, cache_dir, load_format, revision):
             if "lm_head.weight" in name:
                 continue
+            if name.startswith("decoder."):
+                name = "model." + name
+
             for (param_name, weight_name, shard_id) in stacked_params_mapping:
                 if weight_name not in name:
                     continue

{"golden_diff": "diff --git a/vllm/model_executor/models/opt.py b/vllm/model_executor/models/opt.py\n--- a/vllm/model_executor/models/opt.py\n+++ b/vllm/model_executor/models/opt.py\n@@ -331,6 +331,9 @@\n                 model_name_or_path, cache_dir, load_format, revision):\n             if \"lm_head.weight\" in name:\n                 continue\n+            if name.startswith(\"decoder.\"):\n+                name = \"model.\" + name\n+\n             for (param_name, weight_name, shard_id) in stacked_params_mapping:\n                 if weight_name not in name:\n                     continue\n", "issue": "KeyError: 'decoder.layers.23.fc2.weight'  when test model \"facebook/opt-6.7b\"\n<img width=\"1133\" alt=\"image\" src=\"https://github.com/vllm-project/vllm/assets/73372584/23ac48c4-2b1f-42e3-a55f-a252fee7901b\">\r\nIt looks like the model structure and parameters don't match, how should I fix this \n", "before_files": [{"content": "# coding=utf-8\n# Adapted from\n# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/opt/modeling_opt.py\n# Copyright 2023 The vLLM team.\n# Copyright 2022 The Fairseq Authors and The HuggingFace Inc. team. All rights\n# reserved.\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 OPT model compatible with HuggingFace weights.\"\"\"\nfrom typing import List, Optional, Tuple\n\nimport torch\nfrom torch import nn\nfrom transformers import OPTConfig\n\nfrom vllm.model_executor.input_metadata import InputMetadata\nfrom vllm.model_executor.layers.activation import get_act_fn\nfrom vllm.model_executor.layers.attention import PagedAttention\nfrom vllm.model_executor.layers.linear import (ColumnParallelLinear,\n                                               LinearMethodBase,\n                                               QKVParallelLinear,\n                                               ReplicatedLinear,\n                                               RowParallelLinear)\nfrom vllm.model_executor.layers.sampler import Sampler\nfrom vllm.model_executor.layers.vocab_parallel_embedding import (\n    VocabParallelEmbedding)\nfrom vllm.model_executor.parallel_utils.parallel_state import (\n    get_tensor_model_parallel_world_size)\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 OPTLearnedPositionalEmbedding(nn.Embedding):\n\n    def __init__(self, num_embeddings: int, embedding_dim: int):\n        # OPT is set up so that if padding_idx is specified then offset the\n        # embedding ids by 2 and adjust num_embeddings appropriately. Other\n        # models don't have this hack\n        self.offset = 2\n        super().__init__(num_embeddings + self.offset, embedding_dim)\n\n    def forward(self, positions: torch.Tensor):\n        return super().forward(positions + self.offset)\n\n\nclass OPTAttention(nn.Module):\n\n    def __init__(\n        self,\n        embed_dim: int,\n        num_heads: int,\n        bias: bool = True,\n        linear_method: Optional[LinearMethodBase] = None,\n    ) -> None:\n        super().__init__()\n        self.embed_dim = embed_dim\n        tensor_model_parallel_world_size = (\n            get_tensor_model_parallel_world_size())\n        total_num_heads = num_heads\n        assert num_heads % tensor_model_parallel_world_size == 0\n        self.num_heads = total_num_heads // tensor_model_parallel_world_size\n        self.head_dim = embed_dim // total_num_heads\n        self.scaling = self.head_dim**-0.5\n\n        self.qkv_proj = QKVParallelLinear(\n            embed_dim,\n            self.head_dim,\n            total_num_heads,\n            bias=bias,\n            linear_method=linear_method,\n        )\n        self.out_proj = RowParallelLinear(\n            embed_dim,\n            embed_dim,\n            bias=bias,\n            linear_method=linear_method,\n        )\n        self.attn = PagedAttention(self.num_heads,\n                                   self.head_dim,\n                                   scale=self.scaling)\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        kv_cache: KVCache,\n        input_metadata: InputMetadata,\n        cache_event: Optional[torch.cuda.Event],\n    ) -> torch.Tensor:\n        qkv, _ = self.qkv_proj(hidden_states)\n        q, k, v = qkv.chunk(chunks=3, dim=-1)\n        key_cache, value_cache = kv_cache\n        attn_output = self.attn(q, k, v, key_cache, value_cache,\n                                input_metadata, cache_event)\n        output, _ = self.out_proj(attn_output)\n        return output\n\n\nclass OPTDecoderLayer(nn.Module):\n\n    def __init__(\n        self,\n        config: OPTConfig,\n        linear_method: Optional[LinearMethodBase] = None,\n    ):\n        super().__init__()\n        self.config = config\n        self.embed_dim = config.hidden_size\n        self.self_attn = OPTAttention(\n            embed_dim=self.embed_dim,\n            num_heads=config.num_attention_heads,\n            bias=config.enable_bias,\n            linear_method=linear_method,\n        )\n        self.do_layer_norm_before = config.do_layer_norm_before\n\n        self.self_attn_layer_norm = nn.LayerNorm(\n            self.embed_dim,\n            elementwise_affine=config.layer_norm_elementwise_affine)\n        self.fc1 = ColumnParallelLinear(\n            self.embed_dim,\n            config.ffn_dim,\n            bias=config.enable_bias,\n            linear_method=linear_method,\n        )\n        quant_config = getattr(linear_method, \"quant_config\", None)\n        self.activation_fn = get_act_fn(config.activation_function,\n                                        quant_config, config.ffn_dim)\n        self.fc2 = RowParallelLinear(\n            config.ffn_dim,\n            self.embed_dim,\n            bias=config.enable_bias,\n            linear_method=linear_method,\n        )\n        self.final_layer_norm = nn.LayerNorm(\n            self.embed_dim,\n            elementwise_affine=config.layer_norm_elementwise_affine)\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        kv_cache: KVCache,\n        input_metadata: InputMetadata,\n        cache_event: Optional[torch.cuda.Event],\n    ) -> torch.Tensor:\n        # Self Attention\n        residual = hidden_states\n        # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention\n        if self.do_layer_norm_before:\n            hidden_states = self.self_attn_layer_norm(hidden_states)\n        hidden_states = self.self_attn(hidden_states=hidden_states,\n                                       kv_cache=kv_cache,\n                                       input_metadata=input_metadata,\n                                       cache_event=cache_event)\n        hidden_states = residual + hidden_states\n        # 350m applies layer norm AFTER attention\n        if not self.do_layer_norm_before:\n            hidden_states = self.self_attn_layer_norm(hidden_states)\n\n        # Fully Connected\n        residual = hidden_states\n        # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention\n        if self.do_layer_norm_before:\n            hidden_states = self.final_layer_norm(hidden_states)\n        hidden_states, _ = self.fc1(hidden_states)\n        hidden_states = self.activation_fn(hidden_states)\n        hidden_states, _ = self.fc2(hidden_states)\n        hidden_states = residual + hidden_states\n        # 350m applies layer norm AFTER attention\n        if not self.do_layer_norm_before:\n            hidden_states = self.final_layer_norm(hidden_states)\n        return hidden_states\n\n\nclass OPTDecoder(nn.Module):\n\n    def __init__(\n        self,\n        config: OPTConfig,\n        linear_method: Optional[LinearMethodBase] = None,\n    ):\n        super().__init__()\n        self.config = config\n        self.padding_idx = config.pad_token_id\n        self.max_target_positions = config.max_position_embeddings\n        self.vocab_size = config.vocab_size\n\n        self.embed_tokens = VocabParallelEmbedding(\n            config.vocab_size,\n            config.word_embed_proj_dim,\n        )\n        # Positional embeddings are replicated (not sharded).\n        self.embed_positions = OPTLearnedPositionalEmbedding(\n            config.max_position_embeddings, config.hidden_size)\n\n        # Project out & in will be replicated if they exist.\n        if config.word_embed_proj_dim != config.hidden_size:\n            self.project_out = ReplicatedLinear(config.hidden_size,\n                                                config.word_embed_proj_dim,\n                                                bias=False,\n                                                linear_method=linear_method)\n        else:\n            self.project_out = None\n\n        if config.word_embed_proj_dim != config.hidden_size:\n            self.project_in = ReplicatedLinear(config.word_embed_proj_dim,\n                                               config.hidden_size,\n                                               bias=False,\n                                               linear_method=linear_method)\n        else:\n            self.project_in = None\n\n        # Note that the only purpose of `config._remove_final_layer_norm` is to\n        # keep backward compatibility with checkpoints that have been fine-tuned\n        # before transformers v4.20.1\n        # see https://github.com/facebookresearch/metaseq/pull/164\n        if config.do_layer_norm_before and not config._remove_final_layer_norm:\n            self.final_layer_norm = nn.LayerNorm(\n                config.hidden_size,\n                elementwise_affine=config.layer_norm_elementwise_affine)\n        else:\n            self.final_layer_norm = None\n\n        self.layers = nn.ModuleList([\n            OPTDecoderLayer(config, linear_method)\n            for _ in range(config.num_hidden_layers)\n        ])\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        cache_events: Optional[List[torch.cuda.Event]],\n    ) -> torch.Tensor:\n        inputs_embeds = self.embed_tokens(input_ids)\n        pos_embeds = self.embed_positions(positions)\n        if self.project_in is not None:\n            inputs_embeds, _ = self.project_in(inputs_embeds)\n        hidden_states = inputs_embeds + pos_embeds\n\n        for i in range(len(self.layers)):\n            cache_event = None if cache_events is None else cache_events[i]\n            layer = self.layers[i]\n            hidden_states = layer(hidden_states, kv_caches[i], input_metadata,\n                                  cache_event)\n\n        if self.final_layer_norm is not None:\n            hidden_states = self.final_layer_norm(hidden_states)\n        if self.project_out is not None:\n            hidden_states, _ = self.project_out(hidden_states)\n        return hidden_states\n\n\nclass OPTModel(nn.Module):\n\n    def __init__(\n        self,\n        config: OPTConfig,\n        linear_method: Optional[LinearMethodBase] = None,\n    ):\n        super().__init__()\n        self.decoder = OPTDecoder(config, linear_method)\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        cache_events: Optional[List[torch.cuda.Event]],\n    ) -> torch.Tensor:\n        return self.decoder(input_ids, positions, kv_caches, input_metadata,\n                            cache_events)\n\n\nclass OPTForCausalLM(nn.Module):\n\n    def __init__(\n        self,\n        config,\n        linear_method: Optional[LinearMethodBase] = None,\n    ):\n        super().__init__()\n        self.config = config\n        self.linear_method = linear_method\n        self.model = OPTModel(config, linear_method)\n        self.lm_head_weight = self.model.decoder.embed_tokens.weight\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        cache_events: Optional[List[torch.cuda.Event]],\n    ) -> SamplerOutput:\n        hidden_states = self.model(input_ids, positions, kv_caches,\n                                   input_metadata, cache_events)\n        next_tokens = self.sampler(self.lm_head_weight, hidden_states,\n                                   input_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        ]\n        params_dict = dict(self.named_parameters(remove_duplicate=False))\n        for name, loaded_weight in hf_model_weights_iterator(\n                model_name_or_path, cache_dir, load_format, revision):\n            if \"lm_head.weight\" 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                param = params_dict[name.replace(weight_name, param_name)]\n                weight_loader = param.weight_loader\n                weight_loader(param, loaded_weight, shard_id)\n                break\n            else:\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/opt.py"}], "after_files": [{"content": "# coding=utf-8\n# Adapted from\n# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/opt/modeling_opt.py\n# Copyright 2023 The vLLM team.\n# Copyright 2022 The Fairseq Authors and The HuggingFace Inc. team. All rights\n# reserved.\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 OPT model compatible with HuggingFace weights.\"\"\"\nfrom typing import List, Optional, Tuple\n\nimport torch\nfrom torch import nn\nfrom transformers import OPTConfig\n\nfrom vllm.model_executor.input_metadata import InputMetadata\nfrom vllm.model_executor.layers.activation import get_act_fn\nfrom vllm.model_executor.layers.attention import PagedAttention\nfrom vllm.model_executor.layers.linear import (ColumnParallelLinear,\n                                               LinearMethodBase,\n                                               QKVParallelLinear,\n                                               ReplicatedLinear,\n                                               RowParallelLinear)\nfrom vllm.model_executor.layers.sampler import Sampler\nfrom vllm.model_executor.layers.vocab_parallel_embedding import (\n    VocabParallelEmbedding)\nfrom vllm.model_executor.parallel_utils.parallel_state import (\n    get_tensor_model_parallel_world_size)\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 OPTLearnedPositionalEmbedding(nn.Embedding):\n\n    def __init__(self, num_embeddings: int, embedding_dim: int):\n        # OPT is set up so that if padding_idx is specified then offset the\n        # embedding ids by 2 and adjust num_embeddings appropriately. Other\n        # models don't have this hack\n        self.offset = 2\n        super().__init__(num_embeddings + self.offset, embedding_dim)\n\n    def forward(self, positions: torch.Tensor):\n        return super().forward(positions + self.offset)\n\n\nclass OPTAttention(nn.Module):\n\n    def __init__(\n        self,\n        embed_dim: int,\n        num_heads: int,\n        bias: bool = True,\n        linear_method: Optional[LinearMethodBase] = None,\n    ) -> None:\n        super().__init__()\n        self.embed_dim = embed_dim\n        tensor_model_parallel_world_size = (\n            get_tensor_model_parallel_world_size())\n        total_num_heads = num_heads\n        assert num_heads % tensor_model_parallel_world_size == 0\n        self.num_heads = total_num_heads // tensor_model_parallel_world_size\n        self.head_dim = embed_dim // total_num_heads\n        self.scaling = self.head_dim**-0.5\n\n        self.qkv_proj = QKVParallelLinear(\n            embed_dim,\n            self.head_dim,\n            total_num_heads,\n            bias=bias,\n            linear_method=linear_method,\n        )\n        self.out_proj = RowParallelLinear(\n            embed_dim,\n            embed_dim,\n            bias=bias,\n            linear_method=linear_method,\n        )\n        self.attn = PagedAttention(self.num_heads,\n                                   self.head_dim,\n                                   scale=self.scaling)\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        kv_cache: KVCache,\n        input_metadata: InputMetadata,\n        cache_event: Optional[torch.cuda.Event],\n    ) -> torch.Tensor:\n        qkv, _ = self.qkv_proj(hidden_states)\n        q, k, v = qkv.chunk(chunks=3, dim=-1)\n        key_cache, value_cache = kv_cache\n        attn_output = self.attn(q, k, v, key_cache, value_cache,\n                                input_metadata, cache_event)\n        output, _ = self.out_proj(attn_output)\n        return output\n\n\nclass OPTDecoderLayer(nn.Module):\n\n    def __init__(\n        self,\n        config: OPTConfig,\n        linear_method: Optional[LinearMethodBase] = None,\n    ):\n        super().__init__()\n        self.config = config\n        self.embed_dim = config.hidden_size\n        self.self_attn = OPTAttention(\n            embed_dim=self.embed_dim,\n            num_heads=config.num_attention_heads,\n            bias=config.enable_bias,\n            linear_method=linear_method,\n        )\n        self.do_layer_norm_before = config.do_layer_norm_before\n\n        self.self_attn_layer_norm = nn.LayerNorm(\n            self.embed_dim,\n            elementwise_affine=config.layer_norm_elementwise_affine)\n        self.fc1 = ColumnParallelLinear(\n            self.embed_dim,\n            config.ffn_dim,\n            bias=config.enable_bias,\n            linear_method=linear_method,\n        )\n        quant_config = getattr(linear_method, \"quant_config\", None)\n        self.activation_fn = get_act_fn(config.activation_function,\n                                        quant_config, config.ffn_dim)\n        self.fc2 = RowParallelLinear(\n            config.ffn_dim,\n            self.embed_dim,\n            bias=config.enable_bias,\n            linear_method=linear_method,\n        )\n        self.final_layer_norm = nn.LayerNorm(\n            self.embed_dim,\n            elementwise_affine=config.layer_norm_elementwise_affine)\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        kv_cache: KVCache,\n        input_metadata: InputMetadata,\n        cache_event: Optional[torch.cuda.Event],\n    ) -> torch.Tensor:\n        # Self Attention\n        residual = hidden_states\n        # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention\n        if self.do_layer_norm_before:\n            hidden_states = self.self_attn_layer_norm(hidden_states)\n        hidden_states = self.self_attn(hidden_states=hidden_states,\n                                       kv_cache=kv_cache,\n                                       input_metadata=input_metadata,\n                                       cache_event=cache_event)\n        hidden_states = residual + hidden_states\n        # 350m applies layer norm AFTER attention\n        if not self.do_layer_norm_before:\n            hidden_states = self.self_attn_layer_norm(hidden_states)\n\n        # Fully Connected\n        residual = hidden_states\n        # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention\n        if self.do_layer_norm_before:\n            hidden_states = self.final_layer_norm(hidden_states)\n        hidden_states, _ = self.fc1(hidden_states)\n        hidden_states = self.activation_fn(hidden_states)\n        hidden_states, _ = self.fc2(hidden_states)\n        hidden_states = residual + hidden_states\n        # 350m applies layer norm AFTER attention\n        if not self.do_layer_norm_before:\n            hidden_states = self.final_layer_norm(hidden_states)\n        return hidden_states\n\n\nclass OPTDecoder(nn.Module):\n\n    def __init__(\n        self,\n        config: OPTConfig,\n        linear_method: Optional[LinearMethodBase] = None,\n    ):\n        super().__init__()\n        self.config = config\n        self.padding_idx = config.pad_token_id\n        self.max_target_positions = config.max_position_embeddings\n        self.vocab_size = config.vocab_size\n\n        self.embed_tokens = VocabParallelEmbedding(\n            config.vocab_size,\n            config.word_embed_proj_dim,\n        )\n        # Positional embeddings are replicated (not sharded).\n        self.embed_positions = OPTLearnedPositionalEmbedding(\n            config.max_position_embeddings, config.hidden_size)\n\n        # Project out & in will be replicated if they exist.\n        if config.word_embed_proj_dim != config.hidden_size:\n            self.project_out = ReplicatedLinear(config.hidden_size,\n                                                config.word_embed_proj_dim,\n                                                bias=False,\n                                                linear_method=linear_method)\n        else:\n            self.project_out = None\n\n        if config.word_embed_proj_dim != config.hidden_size:\n            self.project_in = ReplicatedLinear(config.word_embed_proj_dim,\n                                               config.hidden_size,\n                                               bias=False,\n                                               linear_method=linear_method)\n        else:\n            self.project_in = None\n\n        # Note that the only purpose of `config._remove_final_layer_norm` is to\n        # keep backward compatibility with checkpoints that have been fine-tuned\n        # before transformers v4.20.1\n        # see https://github.com/facebookresearch/metaseq/pull/164\n        if config.do_layer_norm_before and not config._remove_final_layer_norm:\n            self.final_layer_norm = nn.LayerNorm(\n                config.hidden_size,\n                elementwise_affine=config.layer_norm_elementwise_affine)\n        else:\n            self.final_layer_norm = None\n\n        self.layers = nn.ModuleList([\n            OPTDecoderLayer(config, linear_method)\n            for _ in range(config.num_hidden_layers)\n        ])\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        cache_events: Optional[List[torch.cuda.Event]],\n    ) -> torch.Tensor:\n        inputs_embeds = self.embed_tokens(input_ids)\n        pos_embeds = self.embed_positions(positions)\n        if self.project_in is not None:\n            inputs_embeds, _ = self.project_in(inputs_embeds)\n        hidden_states = inputs_embeds + pos_embeds\n\n        for i in range(len(self.layers)):\n            cache_event = None if cache_events is None else cache_events[i]\n            layer = self.layers[i]\n            hidden_states = layer(hidden_states, kv_caches[i], input_metadata,\n                                  cache_event)\n\n        if self.final_layer_norm is not None:\n            hidden_states = self.final_layer_norm(hidden_states)\n        if self.project_out is not None:\n            hidden_states, _ = self.project_out(hidden_states)\n        return hidden_states\n\n\nclass OPTModel(nn.Module):\n\n    def __init__(\n        self,\n        config: OPTConfig,\n        linear_method: Optional[LinearMethodBase] = None,\n    ):\n        super().__init__()\n        self.decoder = OPTDecoder(config, linear_method)\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        cache_events: Optional[List[torch.cuda.Event]],\n    ) -> torch.Tensor:\n        return self.decoder(input_ids, positions, kv_caches, input_metadata,\n                            cache_events)\n\n\nclass OPTForCausalLM(nn.Module):\n\n    def __init__(\n        self,\n        config,\n        linear_method: Optional[LinearMethodBase] = None,\n    ):\n        super().__init__()\n        self.config = config\n        self.linear_method = linear_method\n        self.model = OPTModel(config, linear_method)\n        self.lm_head_weight = self.model.decoder.embed_tokens.weight\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        cache_events: Optional[List[torch.cuda.Event]],\n    ) -> SamplerOutput:\n        hidden_states = self.model(input_ids, positions, kv_caches,\n                                   input_metadata, cache_events)\n        next_tokens = self.sampler(self.lm_head_weight, hidden_states,\n                                   input_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        ]\n        params_dict = dict(self.named_parameters(remove_duplicate=False))\n        for name, loaded_weight in hf_model_weights_iterator(\n                model_name_or_path, cache_dir, load_format, revision):\n            if \"lm_head.weight\" in name:\n                continue\n            if name.startswith(\"decoder.\"):\n                name = \"model.\" + name\n\n            for (param_name, weight_name, shard_id) in stacked_params_mapping:\n                if weight_name not in name:\n                    continue\n                param = params_dict[name.replace(weight_name, param_name)]\n                weight_loader = param.weight_loader\n                weight_loader(param, loaded_weight, shard_id)\n                break\n            else:\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/opt.py"}]}