We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Error while creating inference server with EleutherAI/gpt-j-6b model.
 I am trying to run GPTj-6b on an instance with 8 L4 GPU getting the below error. Using version 0.2. but tried with latest as well.

 python -m vllm.entrypoints.api_server     --model EleutherAI/gpt-j-6b     --tensor-parallel-size 8  --dtype float16 --host 0.0.0.0 --port 8000 --gpu-memory-utilization 0.95
  File "/opt/conda/envs/myenv/lib/python3.8/site-packages/ray/_private/auto_init_hook.py", line 24, in auto_init_wrapper
    return fn(args, **kwargs)
  File "/opt/conda/envs/myenv/lib/python3.8/site-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
    return func(args, kwargs)
  File "/opt/conda/envs/myenv/lib/python3.8/site-packages/ray/_private/worker.py", line 2547, in get
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AssertionError): ray::RayWorker.execute_method() (pid=17501, ip=10.138.15.207, actor_id=6393f8b00b5b463275043b0b01000000, repr=<vllm.engine.ray_utils.RayWorker object at 0x7f1ac80b5550>)
  File "/opt/conda/envs/myenv/lib/python3.8/site-packages/vllm/engine/ray_utils.py", line 32, in execute_method
    return executor(*args, kwargs)
  File "/opt/conda/envs/myenv/lib/python3.8/site-packages/vllm/worker/worker.py", line 68, in init_model
    self.model = get_model(self.model_config)
  File "/opt/conda/envs/myenv/lib/python3.8/site-packages/vllm/model_executor/model_loader.py", line 101, in get_model
    model.load_weights(model_config.model, model_config.download_dir,
  File "/opt/conda/envs/myenv/lib/python3.8/site-packages/vllm/model_executor/models/gpt_j.py", line 251, in load_weights
    assert param_slice.shape == loaded_weight.shape
AssertionError

--- 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/gpt_j.py`
Content:
```
1 # coding=utf-8
2 # Adapted from
3 # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/gptj/modeling_gptj.py
4 # Copyright 2023 The vLLM team.
5 # Copyright 2021 The EleutherAI and HuggingFace Teams. All rights reserved.
6 #
7 # Licensed under the Apache License, Version 2.0 (the "License");
8 # you may not use this file except in compliance with the License.
9 # You may obtain a copy of the License at
10 #
11 #     http://www.apache.org/licenses/LICENSE-2.0
12 #
13 # Unless required by applicable law or agreed to in writing, software
14 # distributed under the License is distributed on an "AS IS" BASIS,
15 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 # See the License for the specific language governing permissions and
17 # limitations under the License.
18 """Inference-only GPT-J model compatible with HuggingFace weights.
19 
20 The input of the model is flattened to a 1D tensor of tokens. The model uses
21 InputMetadata to extract the original 2D shape of the input.
22 """
23 from typing import List, Optional, Tuple
24 
25 import torch
26 from torch import nn
27 from transformers import GPTJConfig
28 
29 from vllm.model_executor.input_metadata import InputMetadata
30 from vllm.model_executor.layers.activation import get_act_fn
31 from vllm.model_executor.layers.attention import PagedAttentionWithRoPE
32 from vllm.model_executor.layers.sampler import Sampler
33 from vllm.model_executor.weight_utils import (hf_model_weights_iterator,
34                                               load_tensor_parallel_weights)
35 from vllm.model_executor.parallel_utils.parallel_state import (
36     get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)
37 from vllm.model_executor.parallel_utils.layers import (VocabParallelEmbedding,
38                                                        ColumnParallelLinear,
39                                                        RowParallelLinear)
40 from vllm.sequence import SamplerOutput
41 
42 KVCache = Tuple[torch.Tensor, torch.Tensor]
43 
44 
45 class GPTJAttention(nn.Module):
46 
47     def __init__(self, config: GPTJConfig):
48         super().__init__()
49         self.total_num_heads = config.num_attention_heads
50         self.hidden_size = config.hidden_size
51         self.head_size = self.hidden_size // self.total_num_heads
52 
53         self.qkv_proj = ColumnParallelLinear(
54             config.hidden_size,
55             3 * config.hidden_size,
56             bias=False,
57             gather_output=False,
58         )
59         self.out_proj = RowParallelLinear(
60             config.hidden_size,
61             config.hidden_size,
62             bias=False,
63             input_is_parallel=True,
64         )
65 
66         tp_world_size = get_tensor_model_parallel_world_size()
67         assert self.total_num_heads % tp_world_size == 0
68         self.num_heads = self.total_num_heads // tp_world_size
69 
70         scaling = self.head_size**-0.5
71         assert getattr(config, "rotary", True)
72         assert config.rotary_dim % 2 == 0
73         rope_theta = getattr(config, "rope_theta", 10000)
74         max_position_embeddings = getattr(config, "max_position_embeddings",
75                                           8192)
76         self.attn = PagedAttentionWithRoPE(
77             self.num_heads,
78             self.head_size,
79             scaling,
80             config.rotary_dim,
81             base=rope_theta,
82             max_position=max_position_embeddings,
83             is_neox_style=False)
84         self.warmup = False
85 
86     def forward(
87         self,
88         position_ids: torch.Tensor,
89         hidden_states: torch.Tensor,
90         kv_cache: KVCache,
91         input_metadata: InputMetadata,
92         cache_event: Optional[torch.cuda.Event],
93     ) -> torch.Tensor:
94         qkv, _ = self.qkv_proj(hidden_states)
95         q, k, v = qkv.chunk(chunks=3, dim=-1)
96         k_cache, v_cache = kv_cache
97         attn_output = self.attn(position_ids, q, k, v, k_cache, v_cache,
98                                 input_metadata, cache_event)
99         attn_output, _ = self.out_proj(attn_output)
100         return attn_output
101 
102 
103 class GPTJMLP(nn.Module):
104 
105     def __init__(self, intermediate_size: int, config: GPTJConfig):
106         super().__init__()
107         hidden_size = config.n_embd
108         self.fc_in = ColumnParallelLinear(
109             hidden_size,
110             intermediate_size,
111             gather_output=False,
112         )
113         self.fc_out = RowParallelLinear(
114             intermediate_size,
115             hidden_size,
116             input_is_parallel=True,
117         )
118         self.act = get_act_fn(config.activation_function)
119 
120     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
121         hidden_states, _ = self.fc_in(hidden_states)
122         hidden_states = self.act(hidden_states)
123         hidden_states, _ = self.fc_out(hidden_states)
124         return hidden_states
125 
126 
127 class GPTJBlock(nn.Module):
128 
129     def __init__(self, config: GPTJConfig):
130         super().__init__()
131         if config.n_inner is None:
132             inner_dim = 4 * config.n_embd
133         else:
134             inner_dim = config.n_inner
135         self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
136         self.attn = GPTJAttention(config)
137         self.mlp = GPTJMLP(inner_dim, config)
138 
139     def forward(
140         self,
141         position_ids: torch.Tensor,
142         hidden_states: torch.Tensor,
143         kv_cache: KVCache,
144         input_metadata: InputMetadata,
145         cache_event: Optional[torch.cuda.Event],
146     ) -> torch.Tensor:
147         residual = hidden_states
148         hidden_states = self.ln_1(hidden_states)
149         attn_output = self.attn(
150             position_ids=position_ids,
151             hidden_states=hidden_states,
152             kv_cache=kv_cache,
153             input_metadata=input_metadata,
154             cache_event=cache_event,
155         )
156         mlp_output = self.mlp(hidden_states)
157         hidden_states = attn_output + mlp_output + residual
158         return hidden_states
159 
160 
161 class GPTJModel(nn.Module):
162 
163     def __init__(self, config: GPTJConfig):
164         super().__init__()
165         self.config = config
166         self.embed_dim = config.n_embd
167         self.wte = VocabParallelEmbedding(
168             config.vocab_size,
169             self.embed_dim,
170         )
171         self.h = nn.ModuleList(
172             [GPTJBlock(config) for _ in range(config.n_layer)])
173         self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
174 
175     def forward(
176         self,
177         input_ids: torch.Tensor,
178         position_ids: torch.Tensor,
179         kv_caches: List[KVCache],
180         input_metadata: InputMetadata,
181         cache_events: Optional[List[torch.cuda.Event]],
182     ) -> torch.Tensor:
183         hidden_states = self.wte(input_ids)
184         for i in range(len(self.h)):
185             if cache_events is None:
186                 cache_event = None
187             else:
188                 cache_event = cache_events[i]
189             layer = self.h[i]
190             hidden_states = layer(
191                 position_ids,
192                 hidden_states,
193                 kv_caches[i],
194                 input_metadata,
195                 cache_event,
196             )
197         hidden_states = self.ln_f(hidden_states)
198         return hidden_states
199 
200 
201 class GPTJForCausalLM(nn.Module):
202 
203     def __init__(self, config: GPTJConfig):
204         super().__init__()
205         self.config = config
206         assert not config.tie_word_embeddings
207         self.transformer = GPTJModel(config)
208         self.lm_head = ColumnParallelLinear(
209             config.n_embd,
210             config.vocab_size,
211             gather_output=False,
212         )
213         self.sampler = Sampler(config.vocab_size)
214 
215     def forward(
216         self,
217         input_ids: torch.Tensor,
218         positions: torch.Tensor,
219         kv_caches: List[KVCache],
220         input_metadata: InputMetadata,
221         cache_events: Optional[List[torch.cuda.Event]],
222     ) -> SamplerOutput:
223         hidden_states = self.transformer(input_ids, positions, kv_caches,
224                                          input_metadata, cache_events)
225         next_tokens = self.sampler(self.lm_head.weight, hidden_states,
226                                    input_metadata, self.lm_head.bias)
227         return next_tokens
228 
229     _column_parallel_weights = [
230         "wte.weight", "fc_in.weight", "fc_in.bias", "lm_head.weight",
231         "lm_head.bias"
232     ]
233     _row_parallel_weights = ["out_proj.weight", "fc_out.weight"]
234 
235     def load_weights(self,
236                      model_name_or_path: str,
237                      cache_dir: Optional[str] = None,
238                      load_format: str = "auto",
239                      revision: Optional[str] = None):
240         tp_rank = get_tensor_model_parallel_rank()
241         state_dict = self.state_dict()
242         for name, loaded_weight in hf_model_weights_iterator(
243                 model_name_or_path, cache_dir, load_format, revision):
244             if "attn.bias" in name or "attn.masked_bias" in name:
245                 continue
246 
247             is_attention_weight = False
248             for stride_id, att_weight_name in enumerate(
249                 ["q_proj", "k_proj", "v_proj"]):
250                 if att_weight_name not in name:
251                     continue
252                 param = state_dict[name.replace(att_weight_name, "qkv_proj")]
253                 shard_size = param.shape[1]
254                 loaded_weight = loaded_weight[shard_size * tp_rank:shard_size *
255                                               (tp_rank + 1)]
256                 param_slice = param.data[shard_size * stride_id:shard_size *
257                                          (stride_id + 1)]
258                 assert param_slice.shape == loaded_weight.shape
259                 param_slice.copy_(loaded_weight)
260                 is_attention_weight = True
261                 break
262             if is_attention_weight:
263                 continue
264 
265             param = state_dict[name]
266             load_tensor_parallel_weights(param, loaded_weight, name,
267                                          self._column_parallel_weights,
268                                          self._row_parallel_weights, tp_rank)
269
```
--- 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/gpt_j.py b/vllm/model_executor/models/gpt_j.py
--- a/vllm/model_executor/models/gpt_j.py
+++ b/vllm/model_executor/models/gpt_j.py
@@ -250,7 +250,7 @@
                 if att_weight_name not in name:
                     continue
                 param = state_dict[name.replace(att_weight_name, "qkv_proj")]
-                shard_size = param.shape[1]
+                shard_size = param.shape[0] // 3
                 loaded_weight = loaded_weight[shard_size * tp_rank:shard_size *
                                               (tp_rank + 1)]
                 param_slice = param.data[shard_size * stride_id:shard_size *

{"golden_diff": "diff --git a/vllm/model_executor/models/gpt_j.py b/vllm/model_executor/models/gpt_j.py\n--- a/vllm/model_executor/models/gpt_j.py\n+++ b/vllm/model_executor/models/gpt_j.py\n@@ -250,7 +250,7 @@\n                 if att_weight_name not in name:\n                     continue\n                 param = state_dict[name.replace(att_weight_name, \"qkv_proj\")]\n-                shard_size = param.shape[1]\n+                shard_size = param.shape[0] // 3\n                 loaded_weight = loaded_weight[shard_size * tp_rank:shard_size *\n                                               (tp_rank + 1)]\n                 param_slice = param.data[shard_size * stride_id:shard_size *\n", "issue": "Error while creating inference server with EleutherAI/gpt-j-6b model.\n I am trying to run GPTj-6b on an instance with 8 L4 GPU getting the below error. Using version 0.2. but tried with latest as well.\r\n\r\n python -m vllm.entrypoints.api_server     --model EleutherAI/gpt-j-6b     --tensor-parallel-size 8  --dtype float16 --host 0.0.0.0 --port 8000 --gpu-memory-utilization 0.95\r\n  File \"/opt/conda/envs/myenv/lib/python3.8/site-packages/ray/_private/auto_init_hook.py\", line 24, in auto_init_wrapper\r\n    return fn(args, **kwargs)\r\n  File \"/opt/conda/envs/myenv/lib/python3.8/site-packages/ray/_private/client_mode_hook.py\", line 103, in wrapper\r\n    return func(args, kwargs)\r\n  File \"/opt/conda/envs/myenv/lib/python3.8/site-packages/ray/_private/worker.py\", line 2547, in get\r\n    raise value.as_instanceof_cause()\r\nray.exceptions.RayTaskError(AssertionError): ray::RayWorker.execute_method() (pid=17501, ip=10.138.15.207, actor_id=6393f8b00b5b463275043b0b01000000, repr=<vllm.engine.ray_utils.RayWorker object at 0x7f1ac80b5550>)\r\n  File \"/opt/conda/envs/myenv/lib/python3.8/site-packages/vllm/engine/ray_utils.py\", line 32, in execute_method\r\n    return executor(*args, kwargs)\r\n  File \"/opt/conda/envs/myenv/lib/python3.8/site-packages/vllm/worker/worker.py\", line 68, in init_model\r\n    self.model = get_model(self.model_config)\r\n  File \"/opt/conda/envs/myenv/lib/python3.8/site-packages/vllm/model_executor/model_loader.py\", line 101, in get_model\r\n    model.load_weights(model_config.model, model_config.download_dir,\r\n  File \"/opt/conda/envs/myenv/lib/python3.8/site-packages/vllm/model_executor/models/gpt_j.py\", line 251, in load_weights\r\n    assert param_slice.shape == loaded_weight.shape\r\nAssertionError\n", "before_files": [{"content": "# coding=utf-8\n# Adapted from\n# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/gptj/modeling_gptj.py\n# Copyright 2023 The vLLM team.\n# Copyright 2021 The EleutherAI and HuggingFace Teams. All rights 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 GPT-J model compatible with HuggingFace weights.\n\nThe input of the model is flattened to a 1D tensor of tokens. The model uses\nInputMetadata to extract the original 2D shape of the input.\n\"\"\"\nfrom typing import List, Optional, Tuple\n\nimport torch\nfrom torch import nn\nfrom transformers import GPTJConfig\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 PagedAttentionWithRoPE\nfrom vllm.model_executor.layers.sampler import Sampler\nfrom vllm.model_executor.weight_utils import (hf_model_weights_iterator,\n                                              load_tensor_parallel_weights)\nfrom vllm.model_executor.parallel_utils.parallel_state import (\n    get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)\nfrom vllm.model_executor.parallel_utils.layers import (VocabParallelEmbedding,\n                                                       ColumnParallelLinear,\n                                                       RowParallelLinear)\nfrom vllm.sequence import SamplerOutput\n\nKVCache = Tuple[torch.Tensor, torch.Tensor]\n\n\nclass GPTJAttention(nn.Module):\n\n    def __init__(self, config: GPTJConfig):\n        super().__init__()\n        self.total_num_heads = config.num_attention_heads\n        self.hidden_size = config.hidden_size\n        self.head_size = self.hidden_size // self.total_num_heads\n\n        self.qkv_proj = ColumnParallelLinear(\n            config.hidden_size,\n            3 * config.hidden_size,\n            bias=False,\n            gather_output=False,\n        )\n        self.out_proj = RowParallelLinear(\n            config.hidden_size,\n            config.hidden_size,\n            bias=False,\n            input_is_parallel=True,\n        )\n\n        tp_world_size = get_tensor_model_parallel_world_size()\n        assert self.total_num_heads % tp_world_size == 0\n        self.num_heads = self.total_num_heads // tp_world_size\n\n        scaling = self.head_size**-0.5\n        assert getattr(config, \"rotary\", True)\n        assert config.rotary_dim % 2 == 0\n        rope_theta = getattr(config, \"rope_theta\", 10000)\n        max_position_embeddings = getattr(config, \"max_position_embeddings\",\n                                          8192)\n        self.attn = PagedAttentionWithRoPE(\n            self.num_heads,\n            self.head_size,\n            scaling,\n            config.rotary_dim,\n            base=rope_theta,\n            max_position=max_position_embeddings,\n            is_neox_style=False)\n        self.warmup = False\n\n    def forward(\n        self,\n        position_ids: torch.Tensor,\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        k_cache, v_cache = kv_cache\n        attn_output = self.attn(position_ids, q, k, v, k_cache, v_cache,\n                                input_metadata, cache_event)\n        attn_output, _ = self.out_proj(attn_output)\n        return attn_output\n\n\nclass GPTJMLP(nn.Module):\n\n    def __init__(self, intermediate_size: int, config: GPTJConfig):\n        super().__init__()\n        hidden_size = config.n_embd\n        self.fc_in = ColumnParallelLinear(\n            hidden_size,\n            intermediate_size,\n            gather_output=False,\n        )\n        self.fc_out = RowParallelLinear(\n            intermediate_size,\n            hidden_size,\n            input_is_parallel=True,\n        )\n        self.act = get_act_fn(config.activation_function)\n\n    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:\n        hidden_states, _ = self.fc_in(hidden_states)\n        hidden_states = self.act(hidden_states)\n        hidden_states, _ = self.fc_out(hidden_states)\n        return hidden_states\n\n\nclass GPTJBlock(nn.Module):\n\n    def __init__(self, config: GPTJConfig):\n        super().__init__()\n        if config.n_inner is None:\n            inner_dim = 4 * config.n_embd\n        else:\n            inner_dim = config.n_inner\n        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)\n        self.attn = GPTJAttention(config)\n        self.mlp = GPTJMLP(inner_dim, config)\n\n    def forward(\n        self,\n        position_ids: torch.Tensor,\n        hidden_states: torch.Tensor,\n        kv_cache: KVCache,\n        input_metadata: InputMetadata,\n        cache_event: Optional[torch.cuda.Event],\n    ) -> torch.Tensor:\n        residual = hidden_states\n        hidden_states = self.ln_1(hidden_states)\n        attn_output = self.attn(\n            position_ids=position_ids,\n            hidden_states=hidden_states,\n            kv_cache=kv_cache,\n            input_metadata=input_metadata,\n            cache_event=cache_event,\n        )\n        mlp_output = self.mlp(hidden_states)\n        hidden_states = attn_output + mlp_output + residual\n        return hidden_states\n\n\nclass GPTJModel(nn.Module):\n\n    def __init__(self, config: GPTJConfig):\n        super().__init__()\n        self.config = config\n        self.embed_dim = config.n_embd\n        self.wte = VocabParallelEmbedding(\n            config.vocab_size,\n            self.embed_dim,\n        )\n        self.h = nn.ModuleList(\n            [GPTJBlock(config) for _ in range(config.n_layer)])\n        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        position_ids: torch.Tensor,\n        kv_caches: List[KVCache],\n        input_metadata: InputMetadata,\n        cache_events: Optional[List[torch.cuda.Event]],\n    ) -> torch.Tensor:\n        hidden_states = self.wte(input_ids)\n        for i in range(len(self.h)):\n            if cache_events is None:\n                cache_event = None\n            else:\n                cache_event = cache_events[i]\n            layer = self.h[i]\n            hidden_states = layer(\n                position_ids,\n                hidden_states,\n                kv_caches[i],\n                input_metadata,\n                cache_event,\n            )\n        hidden_states = self.ln_f(hidden_states)\n        return hidden_states\n\n\nclass GPTJForCausalLM(nn.Module):\n\n    def __init__(self, config: GPTJConfig):\n        super().__init__()\n        self.config = config\n        assert not config.tie_word_embeddings\n        self.transformer = GPTJModel(config)\n        self.lm_head = ColumnParallelLinear(\n            config.n_embd,\n            config.vocab_size,\n            gather_output=False,\n        )\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.transformer(input_ids, positions, kv_caches,\n                                         input_metadata, cache_events)\n        next_tokens = self.sampler(self.lm_head.weight, hidden_states,\n                                   input_metadata, self.lm_head.bias)\n        return next_tokens\n\n    _column_parallel_weights = [\n        \"wte.weight\", \"fc_in.weight\", \"fc_in.bias\", \"lm_head.weight\",\n        \"lm_head.bias\"\n    ]\n    _row_parallel_weights = [\"out_proj.weight\", \"fc_out.weight\"]\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        tp_rank = get_tensor_model_parallel_rank()\n        state_dict = self.state_dict()\n        for name, loaded_weight in hf_model_weights_iterator(\n                model_name_or_path, cache_dir, load_format, revision):\n            if \"attn.bias\" in name or \"attn.masked_bias\" in name:\n                continue\n\n            is_attention_weight = False\n            for stride_id, att_weight_name in enumerate(\n                [\"q_proj\", \"k_proj\", \"v_proj\"]):\n                if att_weight_name not in name:\n                    continue\n                param = state_dict[name.replace(att_weight_name, \"qkv_proj\")]\n                shard_size = param.shape[1]\n                loaded_weight = loaded_weight[shard_size * tp_rank:shard_size *\n                                              (tp_rank + 1)]\n                param_slice = param.data[shard_size * stride_id:shard_size *\n                                         (stride_id + 1)]\n                assert param_slice.shape == loaded_weight.shape\n                param_slice.copy_(loaded_weight)\n                is_attention_weight = True\n                break\n            if is_attention_weight:\n                continue\n\n            param = state_dict[name]\n            load_tensor_parallel_weights(param, loaded_weight, name,\n                                         self._column_parallel_weights,\n                                         self._row_parallel_weights, tp_rank)\n", "path": "vllm/model_executor/models/gpt_j.py"}], "after_files": [{"content": "# coding=utf-8\n# Adapted from\n# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/gptj/modeling_gptj.py\n# Copyright 2023 The vLLM team.\n# Copyright 2021 The EleutherAI and HuggingFace Teams. All rights 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 GPT-J model compatible with HuggingFace weights.\n\nThe input of the model is flattened to a 1D tensor of tokens. The model uses\nInputMetadata to extract the original 2D shape of the input.\n\"\"\"\nfrom typing import List, Optional, Tuple\n\nimport torch\nfrom torch import nn\nfrom transformers import GPTJConfig\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 PagedAttentionWithRoPE\nfrom vllm.model_executor.layers.sampler import Sampler\nfrom vllm.model_executor.weight_utils import (hf_model_weights_iterator,\n                                              load_tensor_parallel_weights)\nfrom vllm.model_executor.parallel_utils.parallel_state import (\n    get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)\nfrom vllm.model_executor.parallel_utils.layers import (VocabParallelEmbedding,\n                                                       ColumnParallelLinear,\n                                                       RowParallelLinear)\nfrom vllm.sequence import SamplerOutput\n\nKVCache = Tuple[torch.Tensor, torch.Tensor]\n\n\nclass GPTJAttention(nn.Module):\n\n    def __init__(self, config: GPTJConfig):\n        super().__init__()\n        self.total_num_heads = config.num_attention_heads\n        self.hidden_size = config.hidden_size\n        self.head_size = self.hidden_size // self.total_num_heads\n\n        self.qkv_proj = ColumnParallelLinear(\n            config.hidden_size,\n            3 * config.hidden_size,\n            bias=False,\n            gather_output=False,\n        )\n        self.out_proj = RowParallelLinear(\n            config.hidden_size,\n            config.hidden_size,\n            bias=False,\n            input_is_parallel=True,\n        )\n\n        tp_world_size = get_tensor_model_parallel_world_size()\n        assert self.total_num_heads % tp_world_size == 0\n        self.num_heads = self.total_num_heads // tp_world_size\n\n        scaling = self.head_size**-0.5\n        assert getattr(config, \"rotary\", True)\n        assert config.rotary_dim % 2 == 0\n        rope_theta = getattr(config, \"rope_theta\", 10000)\n        max_position_embeddings = getattr(config, \"max_position_embeddings\",\n                                          8192)\n        self.attn = PagedAttentionWithRoPE(\n            self.num_heads,\n            self.head_size,\n            scaling,\n            config.rotary_dim,\n            base=rope_theta,\n            max_position=max_position_embeddings,\n            is_neox_style=False)\n        self.warmup = False\n\n    def forward(\n        self,\n        position_ids: torch.Tensor,\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        k_cache, v_cache = kv_cache\n        attn_output = self.attn(position_ids, q, k, v, k_cache, v_cache,\n                                input_metadata, cache_event)\n        attn_output, _ = self.out_proj(attn_output)\n        return attn_output\n\n\nclass GPTJMLP(nn.Module):\n\n    def __init__(self, intermediate_size: int, config: GPTJConfig):\n        super().__init__()\n        hidden_size = config.n_embd\n        self.fc_in = ColumnParallelLinear(\n            hidden_size,\n            intermediate_size,\n            gather_output=False,\n        )\n        self.fc_out = RowParallelLinear(\n            intermediate_size,\n            hidden_size,\n            input_is_parallel=True,\n        )\n        self.act = get_act_fn(config.activation_function)\n\n    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:\n        hidden_states, _ = self.fc_in(hidden_states)\n        hidden_states = self.act(hidden_states)\n        hidden_states, _ = self.fc_out(hidden_states)\n        return hidden_states\n\n\nclass GPTJBlock(nn.Module):\n\n    def __init__(self, config: GPTJConfig):\n        super().__init__()\n        if config.n_inner is None:\n            inner_dim = 4 * config.n_embd\n        else:\n            inner_dim = config.n_inner\n        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)\n        self.attn = GPTJAttention(config)\n        self.mlp = GPTJMLP(inner_dim, config)\n\n    def forward(\n        self,\n        position_ids: torch.Tensor,\n        hidden_states: torch.Tensor,\n        kv_cache: KVCache,\n        input_metadata: InputMetadata,\n        cache_event: Optional[torch.cuda.Event],\n    ) -> torch.Tensor:\n        residual = hidden_states\n        hidden_states = self.ln_1(hidden_states)\n        attn_output = self.attn(\n            position_ids=position_ids,\n            hidden_states=hidden_states,\n            kv_cache=kv_cache,\n            input_metadata=input_metadata,\n            cache_event=cache_event,\n        )\n        mlp_output = self.mlp(hidden_states)\n        hidden_states = attn_output + mlp_output + residual\n        return hidden_states\n\n\nclass GPTJModel(nn.Module):\n\n    def __init__(self, config: GPTJConfig):\n        super().__init__()\n        self.config = config\n        self.embed_dim = config.n_embd\n        self.wte = VocabParallelEmbedding(\n            config.vocab_size,\n            self.embed_dim,\n        )\n        self.h = nn.ModuleList(\n            [GPTJBlock(config) for _ in range(config.n_layer)])\n        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        position_ids: torch.Tensor,\n        kv_caches: List[KVCache],\n        input_metadata: InputMetadata,\n        cache_events: Optional[List[torch.cuda.Event]],\n    ) -> torch.Tensor:\n        hidden_states = self.wte(input_ids)\n        for i in range(len(self.h)):\n            if cache_events is None:\n                cache_event = None\n            else:\n                cache_event = cache_events[i]\n            layer = self.h[i]\n            hidden_states = layer(\n                position_ids,\n                hidden_states,\n                kv_caches[i],\n                input_metadata,\n                cache_event,\n            )\n        hidden_states = self.ln_f(hidden_states)\n        return hidden_states\n\n\nclass GPTJForCausalLM(nn.Module):\n\n    def __init__(self, config: GPTJConfig):\n        super().__init__()\n        self.config = config\n        assert not config.tie_word_embeddings\n        self.transformer = GPTJModel(config)\n        self.lm_head = ColumnParallelLinear(\n            config.n_embd,\n            config.vocab_size,\n            gather_output=False,\n        )\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.transformer(input_ids, positions, kv_caches,\n                                         input_metadata, cache_events)\n        next_tokens = self.sampler(self.lm_head.weight, hidden_states,\n                                   input_metadata, self.lm_head.bias)\n        return next_tokens\n\n    _column_parallel_weights = [\n        \"wte.weight\", \"fc_in.weight\", \"fc_in.bias\", \"lm_head.weight\",\n        \"lm_head.bias\"\n    ]\n    _row_parallel_weights = [\"out_proj.weight\", \"fc_out.weight\"]\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        tp_rank = get_tensor_model_parallel_rank()\n        state_dict = self.state_dict()\n        for name, loaded_weight in hf_model_weights_iterator(\n                model_name_or_path, cache_dir, load_format, revision):\n            if \"attn.bias\" in name or \"attn.masked_bias\" in name:\n                continue\n\n            is_attention_weight = False\n            for stride_id, att_weight_name in enumerate(\n                [\"q_proj\", \"k_proj\", \"v_proj\"]):\n                if att_weight_name not in name:\n                    continue\n                param = state_dict[name.replace(att_weight_name, \"qkv_proj\")]\n                shard_size = param.shape[0] // 3\n                loaded_weight = loaded_weight[shard_size * tp_rank:shard_size *\n                                              (tp_rank + 1)]\n                param_slice = param.data[shard_size * stride_id:shard_size *\n                                         (stride_id + 1)]\n                assert param_slice.shape == loaded_weight.shape\n                param_slice.copy_(loaded_weight)\n                is_attention_weight = True\n                break\n            if is_attention_weight:\n                continue\n\n            param = state_dict[name]\n            load_tensor_parallel_weights(param, loaded_weight, name,\n                                         self._column_parallel_weights,\n                                         self._row_parallel_weights, tp_rank)\n", "path": "vllm/model_executor/models/gpt_j.py"}]}