We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `colossalai/shardformer/policies/gptj.py`
Content:
```
1 import warnings
2 from functools import partial
3 from typing import Callable, Dict, List
4 
5 from torch import Tensor, nn
6 
7 import colossalai.shardformer.layer as col_nn
8 
9 from ..modeling.gptj import (
10     GPTJPipelineForwards,
11     get_gptj_flash_attention_forward,
12     gptj_model_forward_for_flash_attention,
13 )
14 from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
15 
16 __all__ = [
17     "GPTJPolicy",
18     "GPTJModelPolicy",
19     "GPTJForCausalLMPolicy",
20     "GPTJForSequenceClassificationPolicy",
21     "GPTJForQuestionAnsweringPolicy",
22     "FlaxGPTJPolicy",
23     "FlaxGPTJForCausalLMPolicy",
24 ]
25 
26 
27 class GPTJPolicy(Policy):
28     def config_sanity_check(self):
29         pass
30 
31     def preprocess(self):
32         self.tie_weight = self.tie_weight_check()
33         self.origin_attn_implement = self.model.config._attn_implementation
34         return self.model
35 
36     def module_policy(self):
37         from transformers.models.gptj.modeling_gptj import GPTJAttention, GPTJBlock, GPTJModel
38 
39         ATTN_IMPLEMENTATION = {
40             "eager": GPTJAttention,
41         }
42 
43         policy = {}
44 
45         attn_cls = ATTN_IMPLEMENTATION[self.origin_attn_implement]
46 
47         embedding_cls = None
48         if self.shard_config.enable_tensor_parallelism:
49             embedding_cls = col_nn.VocabParallelEmbedding1D
50         else:
51             if self.tie_weight:
52                 embedding_cls = col_nn.PaddingEmbedding
53 
54         if self.shard_config.enable_sequence_parallelism:
55             self.shard_config.enable_sequence_parallelism = False
56             warnings.warn("GPTJ doesn't support sequence parallelism now, will ignore the sequence parallelism flag.")
57         use_sequence_parallel = self.shard_config.enable_sequence_parallelism
58 
59         overlap = self.shard_config.enable_sequence_overlap
60         if self.shard_config.enable_tensor_parallelism:
61             policy[GPTJModel] = ModulePolicyDescription(
62                 sub_module_replacement=[
63                     SubModuleReplacementDescription(
64                         suffix="drop",
65                         target_module=col_nn.DropoutForParallelInput,
66                     ),
67                 ]
68             )
69 
70             policy[GPTJBlock] = ModulePolicyDescription(
71                 attribute_replacement={
72                     "attn.embed_dim": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
73                     "attn.num_attention_heads": self.model.config.num_attention_heads
74                     // self.shard_config.tensor_parallel_size,
75                 },
76                 sub_module_replacement=[
77                     SubModuleReplacementDescription(
78                         suffix="attn.k_proj",
79                         target_module=col_nn.Linear1D_Col,
80                         kwargs={
81                             "seq_parallel": use_sequence_parallel,
82                             "overlap": overlap,
83                         },
84                     ),
85                     SubModuleReplacementDescription(
86                         suffix="attn.q_proj",
87                         target_module=col_nn.Linear1D_Col,
88                         kwargs={
89                             "seq_parallel": use_sequence_parallel,
90                             "overlap": overlap,
91                         },
92                     ),
93                     SubModuleReplacementDescription(
94                         suffix="attn.v_proj",
95                         target_module=col_nn.Linear1D_Col,
96                         kwargs={
97                             "seq_parallel": use_sequence_parallel,
98                             "overlap": overlap,
99                         },
100                     ),
101                     SubModuleReplacementDescription(
102                         suffix="attn.out_proj",
103                         target_module=col_nn.Linear1D_Row,
104                         kwargs={"seq_parallel": use_sequence_parallel},
105                     ),
106                     SubModuleReplacementDescription(
107                         suffix="mlp.fc_in",
108                         target_module=col_nn.Linear1D_Col,
109                         kwargs={"seq_parallel": use_sequence_parallel},
110                     ),
111                     SubModuleReplacementDescription(
112                         suffix="mlp.fc_out",
113                         target_module=col_nn.Linear1D_Row,
114                         kwargs={"seq_parallel": use_sequence_parallel},
115                     ),
116                     SubModuleReplacementDescription(
117                         suffix="attn.attn_dropout",
118                         target_module=col_nn.DropoutForParallelInput,
119                     ),
120                     SubModuleReplacementDescription(
121                         suffix="attn.resid_dropout",
122                         target_module=col_nn.DropoutForParallelInput,
123                     ),
124                     SubModuleReplacementDescription(
125                         suffix="mlp.dropout",
126                         target_module=col_nn.DropoutForParallelInput,
127                     ),
128                 ],
129             )
130 
131         if embedding_cls is not None:
132             self.append_or_create_submodule_replacement(
133                 description=SubModuleReplacementDescription(
134                     suffix="wte",
135                     target_module=embedding_cls,
136                     kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
137                 ),
138                 policy=policy,
139                 target_key=GPTJModel,
140             )
141 
142         # optimization configuration
143         if self.shard_config.enable_fused_normalization:
144             self.append_or_create_submodule_replacement(
145                 description=SubModuleReplacementDescription(
146                     suffix="ln_f",
147                     target_module=col_nn.FusedLayerNorm,
148                 ),
149                 policy=policy,
150                 target_key=GPTJModel,
151             )
152 
153             self.append_or_create_submodule_replacement(
154                 description=[
155                     SubModuleReplacementDescription(
156                         suffix="ln_1",
157                         target_module=col_nn.FusedLayerNorm,
158                     )
159                 ],
160                 policy=policy,
161                 target_key=GPTJBlock,
162             )
163 
164         if self.shard_config.enable_flash_attention:
165             self.append_or_create_method_replacement(
166                 description={
167                     "forward": get_gptj_flash_attention_forward(),
168                 },
169                 policy=policy,
170                 target_key=attn_cls,
171             )
172             if not self.shard_config.pipeline_stage_manager:
173                 self.append_or_create_method_replacement(
174                     description={"forward": gptj_model_forward_for_flash_attention(self.shard_config)},
175                     policy=policy,
176                     target_key=GPTJModel,
177                 )
178 
179         return policy
180 
181     def postprocess(self):
182         return self.model
183 
184     def get_held_layers(self) -> List[nn.Module]:
185         """Get pipeline layers for current stage."""
186         assert self.pipeline_stage_manager is not None
187 
188         if self.model.__class__.__name__ == "GPTJModel":
189             module = self.model
190         else:
191             module = self.model.transformer
192         stage_manager = self.pipeline_stage_manager
193 
194         held_layers = []
195         layers_per_stage = stage_manager.distribute_layers(len(module.h))
196         if stage_manager.is_first_stage():
197             held_layers.append(module.wte)
198             held_layers.append(module.drop)
199         start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage)
200         held_layers.extend(module.h[start_idx:end_idx])
201         if stage_manager.is_last_stage():
202             held_layers.append(module.ln_f)
203         return held_layers
204 
205     def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None:
206         """If under pipeline parallel setting, replacing the original forward method of huggingface
207         to customized forward method, and add this changing to policy."""
208         if not self.pipeline_stage_manager:
209             raise ValueError("set_pipeline_forward method can only be called when pipeline parallel is enabled.")
210         stage_manager = self.pipeline_stage_manager
211         if self.model.__class__.__name__ == "GPTJModel":
212             module = self.model
213         else:
214             module = self.model.transformer
215 
216         layers_per_stage = stage_manager.distribute_layers(len(module.h))
217         stage_index = stage_manager.get_stage_index(layers_per_stage)
218         method_replacement = {
219             "forward": partial(
220                 new_forward,
221                 stage_manager=stage_manager,
222                 stage_index=stage_index,
223                 shard_config=self.shard_config,
224             )
225         }
226         self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls)
227 
228 
229 # GPTJModel
230 class GPTJModelPolicy(GPTJPolicy):
231     def __init__(self) -> None:
232         super().__init__()
233 
234     def module_policy(self):
235         from transformers.models.gptj.modeling_gptj import GPTJModel
236 
237         policy = super().module_policy()
238 
239         if self.pipeline_stage_manager is not None:
240             self.set_pipeline_forward(
241                 model_cls=GPTJModel,
242                 new_forward=GPTJPipelineForwards.gptj_model_forward,
243                 policy=policy,
244             )
245         return policy
246 
247     def get_held_layers(self) -> List[nn.Module]:
248         return super().get_held_layers()
249 
250     def get_shared_params(self) -> List[Dict[int, Tensor]]:
251         """No shared params in GPT2Model."""
252         return []
253 
254 
255 # GPTJForCausalLM
256 class GPTJForCausalLMPolicy(GPTJPolicy):
257     def __init__(self) -> None:
258         super().__init__()
259 
260     def module_policy(self):
261         from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
262 
263         policy = super().module_policy()
264 
265         if self.shard_config.enable_tensor_parallelism:
266             addon_module = {
267                 GPTJForCausalLM: ModulePolicyDescription(
268                     sub_module_replacement=[
269                         SubModuleReplacementDescription(
270                             suffix="lm_head",
271                             target_module=col_nn.VocabParallelLMHead1D,
272                             kwargs={
273                                 "gather_output": True,
274                                 "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by,
275                             },
276                         )
277                     ]
278                 )
279             }
280         else:
281             addon_module = {
282                 GPTJForCausalLM: ModulePolicyDescription(
283                     sub_module_replacement=[
284                         SubModuleReplacementDescription(
285                             suffix="lm_head",
286                             target_module=col_nn.PaddingLMHead,
287                             kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
288                         )
289                     ]
290                 )
291             }
292         policy.update(addon_module)
293 
294         if self.pipeline_stage_manager is not None:
295             self.set_pipeline_forward(
296                 model_cls=GPTJForCausalLM,
297                 new_forward=GPTJPipelineForwards.gptj_causallm_model_forward,
298                 policy=policy,
299             )
300         return policy
301 
302     def get_held_layers(self) -> List[nn.Module]:
303         held_layers = super().get_held_layers()
304         if self.pipeline_stage_manager.is_last_stage():
305             held_layers.append(self.model.lm_head)
306         return held_layers
307 
308     def get_shared_params(self) -> List[Dict[int, Tensor]]:
309         """The weights of wte and lm_head are shared."""
310         module = self.model
311         stage_manager = self.pipeline_stage_manager
312         if stage_manager is not None:
313             if stage_manager.num_stages > 1 and id(module.transformer.wte.weight) == id(module.lm_head.weight):
314                 first_stage, last_stage = 0, stage_manager.num_stages - 1
315                 return [
316                     {
317                         first_stage: module.transformer.wte.weight,
318                         last_stage: module.lm_head.weight,
319                     }
320                 ]
321         return []
322 
323 
324 # GPTJForSequenceClassification
325 class GPTJForSequenceClassificationPolicy(GPTJPolicy):
326     def __init__(self) -> None:
327         super().__init__()
328 
329     def module_policy(self):
330         from transformers.models.gptj.modeling_gptj import GPTJForSequenceClassification
331 
332         policy = super().module_policy()
333 
334         if self.pipeline_stage_manager is not None:
335             self.set_pipeline_forward(
336                 model_cls=GPTJForSequenceClassification,
337                 new_forward=GPTJPipelineForwards.gptj_for_sequence_classification_forward,
338                 policy=policy,
339             )
340         return policy
341 
342     def get_held_layers(self) -> List[nn.Module]:
343         held_layers = super().get_held_layers()
344         if self.pipeline_stage_manager.is_last_stage():
345             held_layers.append(self.model.score)
346         return held_layers
347 
348     def get_shared_params(self) -> List[Dict[int, Tensor]]:
349         """No shared params in GPTJForSequenceClassification."""
350         return []
351 
352 
353 # GPTJForQuestionAnswering
354 class GPTJForQuestionAnsweringPolicy(GPTJPolicy):
355     def __init__(self) -> None:
356         super().__init__()
357 
358     def module_policy(self):
359         from transformers.models.gptj.modeling_gptj import GPTJForQuestionAnswering
360 
361         policy = super().module_policy()
362 
363         if self.pipeline_stage_manager is not None:
364             self.set_pipeline_forward(
365                 model_cls=GPTJForQuestionAnswering,
366                 new_forward=GPTJPipelineForwards.gptj_for_question_answering_forward,
367                 policy=policy,
368             )
369         return policy
370 
371     def get_held_layers(self) -> List[nn.Module]:
372         held_layers = super().get_held_layers()
373         if self.pipeline_stage_manager.is_last_stage():
374             held_layers.append(self.model.qa_outputs)
375         return held_layers
376 
377     def get_shared_params(self) -> List[Dict[int, Tensor]]:
378         """No shared params in GPT2ForQuestionAnswering."""
379         return []
380
```
--- 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/colossalai/shardformer/policies/gptj.py b/colossalai/shardformer/policies/gptj.py
--- a/colossalai/shardformer/policies/gptj.py
+++ b/colossalai/shardformer/policies/gptj.py
@@ -54,7 +54,6 @@
         if self.shard_config.enable_sequence_parallelism:
             self.shard_config.enable_sequence_parallelism = False
             warnings.warn("GPTJ doesn't support sequence parallelism now, will ignore the sequence parallelism flag.")
-        use_sequence_parallel = self.shard_config.enable_sequence_parallelism
 
         overlap = self.shard_config.enable_sequence_overlap
         if self.shard_config.enable_tensor_parallelism:
@@ -78,7 +77,6 @@
                         suffix="attn.k_proj",
                         target_module=col_nn.Linear1D_Col,
                         kwargs={
-                            "seq_parallel": use_sequence_parallel,
                             "overlap": overlap,
                         },
                     ),
@@ -86,7 +84,6 @@
                         suffix="attn.q_proj",
                         target_module=col_nn.Linear1D_Col,
                         kwargs={
-                            "seq_parallel": use_sequence_parallel,
                             "overlap": overlap,
                         },
                     ),
@@ -94,24 +91,20 @@
                         suffix="attn.v_proj",
                         target_module=col_nn.Linear1D_Col,
                         kwargs={
-                            "seq_parallel": use_sequence_parallel,
                             "overlap": overlap,
                         },
                     ),
                     SubModuleReplacementDescription(
                         suffix="attn.out_proj",
                         target_module=col_nn.Linear1D_Row,
-                        kwargs={"seq_parallel": use_sequence_parallel},
                     ),
                     SubModuleReplacementDescription(
                         suffix="mlp.fc_in",
                         target_module=col_nn.Linear1D_Col,
-                        kwargs={"seq_parallel": use_sequence_parallel},
                     ),
                     SubModuleReplacementDescription(
                         suffix="mlp.fc_out",
                         target_module=col_nn.Linear1D_Row,
-                        kwargs={"seq_parallel": use_sequence_parallel},
                     ),
                     SubModuleReplacementDescription(
                         suffix="attn.attn_dropout",

{"golden_diff": "diff --git a/colossalai/shardformer/policies/gptj.py b/colossalai/shardformer/policies/gptj.py\n--- a/colossalai/shardformer/policies/gptj.py\n+++ b/colossalai/shardformer/policies/gptj.py\n@@ -54,7 +54,6 @@\n         if self.shard_config.enable_sequence_parallelism:\n             self.shard_config.enable_sequence_parallelism = False\n             warnings.warn(\"GPTJ doesn't support sequence parallelism now, will ignore the sequence parallelism flag.\")\n-        use_sequence_parallel = self.shard_config.enable_sequence_parallelism\n \n         overlap = self.shard_config.enable_sequence_overlap\n         if self.shard_config.enable_tensor_parallelism:\n@@ -78,7 +77,6 @@\n                         suffix=\"attn.k_proj\",\n                         target_module=col_nn.Linear1D_Col,\n                         kwargs={\n-                            \"seq_parallel\": use_sequence_parallel,\n                             \"overlap\": overlap,\n                         },\n                     ),\n@@ -86,7 +84,6 @@\n                         suffix=\"attn.q_proj\",\n                         target_module=col_nn.Linear1D_Col,\n                         kwargs={\n-                            \"seq_parallel\": use_sequence_parallel,\n                             \"overlap\": overlap,\n                         },\n                     ),\n@@ -94,24 +91,20 @@\n                         suffix=\"attn.v_proj\",\n                         target_module=col_nn.Linear1D_Col,\n                         kwargs={\n-                            \"seq_parallel\": use_sequence_parallel,\n                             \"overlap\": overlap,\n                         },\n                     ),\n                     SubModuleReplacementDescription(\n                         suffix=\"attn.out_proj\",\n                         target_module=col_nn.Linear1D_Row,\n-                        kwargs={\"seq_parallel\": use_sequence_parallel},\n                     ),\n                     SubModuleReplacementDescription(\n                         suffix=\"mlp.fc_in\",\n                         target_module=col_nn.Linear1D_Col,\n-                        kwargs={\"seq_parallel\": use_sequence_parallel},\n                     ),\n                     SubModuleReplacementDescription(\n                         suffix=\"mlp.fc_out\",\n                         target_module=col_nn.Linear1D_Row,\n-                        kwargs={\"seq_parallel\": use_sequence_parallel},\n                     ),\n                     SubModuleReplacementDescription(\n                         suffix=\"attn.attn_dropout\",\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import warnings\nfrom functools import partial\nfrom typing import Callable, Dict, List\n\nfrom torch import Tensor, nn\n\nimport colossalai.shardformer.layer as col_nn\n\nfrom ..modeling.gptj import (\n    GPTJPipelineForwards,\n    get_gptj_flash_attention_forward,\n    gptj_model_forward_for_flash_attention,\n)\nfrom .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription\n\n__all__ = [\n    \"GPTJPolicy\",\n    \"GPTJModelPolicy\",\n    \"GPTJForCausalLMPolicy\",\n    \"GPTJForSequenceClassificationPolicy\",\n    \"GPTJForQuestionAnsweringPolicy\",\n    \"FlaxGPTJPolicy\",\n    \"FlaxGPTJForCausalLMPolicy\",\n]\n\n\nclass GPTJPolicy(Policy):\n    def config_sanity_check(self):\n        pass\n\n    def preprocess(self):\n        self.tie_weight = self.tie_weight_check()\n        self.origin_attn_implement = self.model.config._attn_implementation\n        return self.model\n\n    def module_policy(self):\n        from transformers.models.gptj.modeling_gptj import GPTJAttention, GPTJBlock, GPTJModel\n\n        ATTN_IMPLEMENTATION = {\n            \"eager\": GPTJAttention,\n        }\n\n        policy = {}\n\n        attn_cls = ATTN_IMPLEMENTATION[self.origin_attn_implement]\n\n        embedding_cls = None\n        if self.shard_config.enable_tensor_parallelism:\n            embedding_cls = col_nn.VocabParallelEmbedding1D\n        else:\n            if self.tie_weight:\n                embedding_cls = col_nn.PaddingEmbedding\n\n        if self.shard_config.enable_sequence_parallelism:\n            self.shard_config.enable_sequence_parallelism = False\n            warnings.warn(\"GPTJ doesn't support sequence parallelism now, will ignore the sequence parallelism flag.\")\n        use_sequence_parallel = self.shard_config.enable_sequence_parallelism\n\n        overlap = self.shard_config.enable_sequence_overlap\n        if self.shard_config.enable_tensor_parallelism:\n            policy[GPTJModel] = ModulePolicyDescription(\n                sub_module_replacement=[\n                    SubModuleReplacementDescription(\n                        suffix=\"drop\",\n                        target_module=col_nn.DropoutForParallelInput,\n                    ),\n                ]\n            )\n\n            policy[GPTJBlock] = ModulePolicyDescription(\n                attribute_replacement={\n                    \"attn.embed_dim\": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,\n                    \"attn.num_attention_heads\": self.model.config.num_attention_heads\n                    // self.shard_config.tensor_parallel_size,\n                },\n                sub_module_replacement=[\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.k_proj\",\n                        target_module=col_nn.Linear1D_Col,\n                        kwargs={\n                            \"seq_parallel\": use_sequence_parallel,\n                            \"overlap\": overlap,\n                        },\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.q_proj\",\n                        target_module=col_nn.Linear1D_Col,\n                        kwargs={\n                            \"seq_parallel\": use_sequence_parallel,\n                            \"overlap\": overlap,\n                        },\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.v_proj\",\n                        target_module=col_nn.Linear1D_Col,\n                        kwargs={\n                            \"seq_parallel\": use_sequence_parallel,\n                            \"overlap\": overlap,\n                        },\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.out_proj\",\n                        target_module=col_nn.Linear1D_Row,\n                        kwargs={\"seq_parallel\": use_sequence_parallel},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.fc_in\",\n                        target_module=col_nn.Linear1D_Col,\n                        kwargs={\"seq_parallel\": use_sequence_parallel},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.fc_out\",\n                        target_module=col_nn.Linear1D_Row,\n                        kwargs={\"seq_parallel\": use_sequence_parallel},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.attn_dropout\",\n                        target_module=col_nn.DropoutForParallelInput,\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.resid_dropout\",\n                        target_module=col_nn.DropoutForParallelInput,\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.dropout\",\n                        target_module=col_nn.DropoutForParallelInput,\n                    ),\n                ],\n            )\n\n        if embedding_cls is not None:\n            self.append_or_create_submodule_replacement(\n                description=SubModuleReplacementDescription(\n                    suffix=\"wte\",\n                    target_module=embedding_cls,\n                    kwargs={\"make_vocab_size_divisible_by\": self.shard_config.make_vocab_size_divisible_by},\n                ),\n                policy=policy,\n                target_key=GPTJModel,\n            )\n\n        # optimization configuration\n        if self.shard_config.enable_fused_normalization:\n            self.append_or_create_submodule_replacement(\n                description=SubModuleReplacementDescription(\n                    suffix=\"ln_f\",\n                    target_module=col_nn.FusedLayerNorm,\n                ),\n                policy=policy,\n                target_key=GPTJModel,\n            )\n\n            self.append_or_create_submodule_replacement(\n                description=[\n                    SubModuleReplacementDescription(\n                        suffix=\"ln_1\",\n                        target_module=col_nn.FusedLayerNorm,\n                    )\n                ],\n                policy=policy,\n                target_key=GPTJBlock,\n            )\n\n        if self.shard_config.enable_flash_attention:\n            self.append_or_create_method_replacement(\n                description={\n                    \"forward\": get_gptj_flash_attention_forward(),\n                },\n                policy=policy,\n                target_key=attn_cls,\n            )\n            if not self.shard_config.pipeline_stage_manager:\n                self.append_or_create_method_replacement(\n                    description={\"forward\": gptj_model_forward_for_flash_attention(self.shard_config)},\n                    policy=policy,\n                    target_key=GPTJModel,\n                )\n\n        return policy\n\n    def postprocess(self):\n        return self.model\n\n    def get_held_layers(self) -> List[nn.Module]:\n        \"\"\"Get pipeline layers for current stage.\"\"\"\n        assert self.pipeline_stage_manager is not None\n\n        if self.model.__class__.__name__ == \"GPTJModel\":\n            module = self.model\n        else:\n            module = self.model.transformer\n        stage_manager = self.pipeline_stage_manager\n\n        held_layers = []\n        layers_per_stage = stage_manager.distribute_layers(len(module.h))\n        if stage_manager.is_first_stage():\n            held_layers.append(module.wte)\n            held_layers.append(module.drop)\n        start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage)\n        held_layers.extend(module.h[start_idx:end_idx])\n        if stage_manager.is_last_stage():\n            held_layers.append(module.ln_f)\n        return held_layers\n\n    def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None:\n        \"\"\"If under pipeline parallel setting, replacing the original forward method of huggingface\n        to customized forward method, and add this changing to policy.\"\"\"\n        if not self.pipeline_stage_manager:\n            raise ValueError(\"set_pipeline_forward method can only be called when pipeline parallel is enabled.\")\n        stage_manager = self.pipeline_stage_manager\n        if self.model.__class__.__name__ == \"GPTJModel\":\n            module = self.model\n        else:\n            module = self.model.transformer\n\n        layers_per_stage = stage_manager.distribute_layers(len(module.h))\n        stage_index = stage_manager.get_stage_index(layers_per_stage)\n        method_replacement = {\n            \"forward\": partial(\n                new_forward,\n                stage_manager=stage_manager,\n                stage_index=stage_index,\n                shard_config=self.shard_config,\n            )\n        }\n        self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls)\n\n\n# GPTJModel\nclass GPTJModelPolicy(GPTJPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        from transformers.models.gptj.modeling_gptj import GPTJModel\n\n        policy = super().module_policy()\n\n        if self.pipeline_stage_manager is not None:\n            self.set_pipeline_forward(\n                model_cls=GPTJModel,\n                new_forward=GPTJPipelineForwards.gptj_model_forward,\n                policy=policy,\n            )\n        return policy\n\n    def get_held_layers(self) -> List[nn.Module]:\n        return super().get_held_layers()\n\n    def get_shared_params(self) -> List[Dict[int, Tensor]]:\n        \"\"\"No shared params in GPT2Model.\"\"\"\n        return []\n\n\n# GPTJForCausalLM\nclass GPTJForCausalLMPolicy(GPTJPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        from transformers.models.gptj.modeling_gptj import GPTJForCausalLM\n\n        policy = super().module_policy()\n\n        if self.shard_config.enable_tensor_parallelism:\n            addon_module = {\n                GPTJForCausalLM: ModulePolicyDescription(\n                    sub_module_replacement=[\n                        SubModuleReplacementDescription(\n                            suffix=\"lm_head\",\n                            target_module=col_nn.VocabParallelLMHead1D,\n                            kwargs={\n                                \"gather_output\": True,\n                                \"make_vocab_size_divisible_by\": self.shard_config.make_vocab_size_divisible_by,\n                            },\n                        )\n                    ]\n                )\n            }\n        else:\n            addon_module = {\n                GPTJForCausalLM: ModulePolicyDescription(\n                    sub_module_replacement=[\n                        SubModuleReplacementDescription(\n                            suffix=\"lm_head\",\n                            target_module=col_nn.PaddingLMHead,\n                            kwargs={\"make_vocab_size_divisible_by\": self.shard_config.make_vocab_size_divisible_by},\n                        )\n                    ]\n                )\n            }\n        policy.update(addon_module)\n\n        if self.pipeline_stage_manager is not None:\n            self.set_pipeline_forward(\n                model_cls=GPTJForCausalLM,\n                new_forward=GPTJPipelineForwards.gptj_causallm_model_forward,\n                policy=policy,\n            )\n        return policy\n\n    def get_held_layers(self) -> List[nn.Module]:\n        held_layers = super().get_held_layers()\n        if self.pipeline_stage_manager.is_last_stage():\n            held_layers.append(self.model.lm_head)\n        return held_layers\n\n    def get_shared_params(self) -> List[Dict[int, Tensor]]:\n        \"\"\"The weights of wte and lm_head are shared.\"\"\"\n        module = self.model\n        stage_manager = self.pipeline_stage_manager\n        if stage_manager is not None:\n            if stage_manager.num_stages > 1 and id(module.transformer.wte.weight) == id(module.lm_head.weight):\n                first_stage, last_stage = 0, stage_manager.num_stages - 1\n                return [\n                    {\n                        first_stage: module.transformer.wte.weight,\n                        last_stage: module.lm_head.weight,\n                    }\n                ]\n        return []\n\n\n# GPTJForSequenceClassification\nclass GPTJForSequenceClassificationPolicy(GPTJPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        from transformers.models.gptj.modeling_gptj import GPTJForSequenceClassification\n\n        policy = super().module_policy()\n\n        if self.pipeline_stage_manager is not None:\n            self.set_pipeline_forward(\n                model_cls=GPTJForSequenceClassification,\n                new_forward=GPTJPipelineForwards.gptj_for_sequence_classification_forward,\n                policy=policy,\n            )\n        return policy\n\n    def get_held_layers(self) -> List[nn.Module]:\n        held_layers = super().get_held_layers()\n        if self.pipeline_stage_manager.is_last_stage():\n            held_layers.append(self.model.score)\n        return held_layers\n\n    def get_shared_params(self) -> List[Dict[int, Tensor]]:\n        \"\"\"No shared params in GPTJForSequenceClassification.\"\"\"\n        return []\n\n\n# GPTJForQuestionAnswering\nclass GPTJForQuestionAnsweringPolicy(GPTJPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        from transformers.models.gptj.modeling_gptj import GPTJForQuestionAnswering\n\n        policy = super().module_policy()\n\n        if self.pipeline_stage_manager is not None:\n            self.set_pipeline_forward(\n                model_cls=GPTJForQuestionAnswering,\n                new_forward=GPTJPipelineForwards.gptj_for_question_answering_forward,\n                policy=policy,\n            )\n        return policy\n\n    def get_held_layers(self) -> List[nn.Module]:\n        held_layers = super().get_held_layers()\n        if self.pipeline_stage_manager.is_last_stage():\n            held_layers.append(self.model.qa_outputs)\n        return held_layers\n\n    def get_shared_params(self) -> List[Dict[int, Tensor]]:\n        \"\"\"No shared params in GPT2ForQuestionAnswering.\"\"\"\n        return []\n", "path": "colossalai/shardformer/policies/gptj.py"}], "after_files": [{"content": "import warnings\nfrom functools import partial\nfrom typing import Callable, Dict, List\n\nfrom torch import Tensor, nn\n\nimport colossalai.shardformer.layer as col_nn\n\nfrom ..modeling.gptj import (\n    GPTJPipelineForwards,\n    get_gptj_flash_attention_forward,\n    gptj_model_forward_for_flash_attention,\n)\nfrom .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription\n\n__all__ = [\n    \"GPTJPolicy\",\n    \"GPTJModelPolicy\",\n    \"GPTJForCausalLMPolicy\",\n    \"GPTJForSequenceClassificationPolicy\",\n    \"GPTJForQuestionAnsweringPolicy\",\n    \"FlaxGPTJPolicy\",\n    \"FlaxGPTJForCausalLMPolicy\",\n]\n\n\nclass GPTJPolicy(Policy):\n    def config_sanity_check(self):\n        pass\n\n    def preprocess(self):\n        self.tie_weight = self.tie_weight_check()\n        self.origin_attn_implement = self.model.config._attn_implementation\n        return self.model\n\n    def module_policy(self):\n        from transformers.models.gptj.modeling_gptj import GPTJAttention, GPTJBlock, GPTJModel\n\n        ATTN_IMPLEMENTATION = {\n            \"eager\": GPTJAttention,\n        }\n\n        policy = {}\n\n        attn_cls = ATTN_IMPLEMENTATION[self.origin_attn_implement]\n\n        embedding_cls = None\n        if self.shard_config.enable_tensor_parallelism:\n            embedding_cls = col_nn.VocabParallelEmbedding1D\n        else:\n            if self.tie_weight:\n                embedding_cls = col_nn.PaddingEmbedding\n\n        if self.shard_config.enable_sequence_parallelism:\n            self.shard_config.enable_sequence_parallelism = False\n            warnings.warn(\"GPTJ doesn't support sequence parallelism now, will ignore the sequence parallelism flag.\")\n\n        overlap = self.shard_config.enable_sequence_overlap\n        if self.shard_config.enable_tensor_parallelism:\n            policy[GPTJModel] = ModulePolicyDescription(\n                sub_module_replacement=[\n                    SubModuleReplacementDescription(\n                        suffix=\"drop\",\n                        target_module=col_nn.DropoutForParallelInput,\n                    ),\n                ]\n            )\n\n            policy[GPTJBlock] = ModulePolicyDescription(\n                attribute_replacement={\n                    \"attn.embed_dim\": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,\n                    \"attn.num_attention_heads\": self.model.config.num_attention_heads\n                    // self.shard_config.tensor_parallel_size,\n                },\n                sub_module_replacement=[\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.k_proj\",\n                        target_module=col_nn.Linear1D_Col,\n                        kwargs={\n                            \"overlap\": overlap,\n                        },\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.q_proj\",\n                        target_module=col_nn.Linear1D_Col,\n                        kwargs={\n                            \"overlap\": overlap,\n                        },\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.v_proj\",\n                        target_module=col_nn.Linear1D_Col,\n                        kwargs={\n                            \"overlap\": overlap,\n                        },\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.out_proj\",\n                        target_module=col_nn.Linear1D_Row,\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.fc_in\",\n                        target_module=col_nn.Linear1D_Col,\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.fc_out\",\n                        target_module=col_nn.Linear1D_Row,\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.attn_dropout\",\n                        target_module=col_nn.DropoutForParallelInput,\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"attn.resid_dropout\",\n                        target_module=col_nn.DropoutForParallelInput,\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.dropout\",\n                        target_module=col_nn.DropoutForParallelInput,\n                    ),\n                ],\n            )\n\n        if embedding_cls is not None:\n            self.append_or_create_submodule_replacement(\n                description=SubModuleReplacementDescription(\n                    suffix=\"wte\",\n                    target_module=embedding_cls,\n                    kwargs={\"make_vocab_size_divisible_by\": self.shard_config.make_vocab_size_divisible_by},\n                ),\n                policy=policy,\n                target_key=GPTJModel,\n            )\n\n        # optimization configuration\n        if self.shard_config.enable_fused_normalization:\n            self.append_or_create_submodule_replacement(\n                description=SubModuleReplacementDescription(\n                    suffix=\"ln_f\",\n                    target_module=col_nn.FusedLayerNorm,\n                ),\n                policy=policy,\n                target_key=GPTJModel,\n            )\n\n            self.append_or_create_submodule_replacement(\n                description=[\n                    SubModuleReplacementDescription(\n                        suffix=\"ln_1\",\n                        target_module=col_nn.FusedLayerNorm,\n                    )\n                ],\n                policy=policy,\n                target_key=GPTJBlock,\n            )\n\n        if self.shard_config.enable_flash_attention:\n            self.append_or_create_method_replacement(\n                description={\n                    \"forward\": get_gptj_flash_attention_forward(),\n                },\n                policy=policy,\n                target_key=attn_cls,\n            )\n            if not self.shard_config.pipeline_stage_manager:\n                self.append_or_create_method_replacement(\n                    description={\"forward\": gptj_model_forward_for_flash_attention(self.shard_config)},\n                    policy=policy,\n                    target_key=GPTJModel,\n                )\n\n        return policy\n\n    def postprocess(self):\n        return self.model\n\n    def get_held_layers(self) -> List[nn.Module]:\n        \"\"\"Get pipeline layers for current stage.\"\"\"\n        assert self.pipeline_stage_manager is not None\n\n        if self.model.__class__.__name__ == \"GPTJModel\":\n            module = self.model\n        else:\n            module = self.model.transformer\n        stage_manager = self.pipeline_stage_manager\n\n        held_layers = []\n        layers_per_stage = stage_manager.distribute_layers(len(module.h))\n        if stage_manager.is_first_stage():\n            held_layers.append(module.wte)\n            held_layers.append(module.drop)\n        start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage)\n        held_layers.extend(module.h[start_idx:end_idx])\n        if stage_manager.is_last_stage():\n            held_layers.append(module.ln_f)\n        return held_layers\n\n    def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None:\n        \"\"\"If under pipeline parallel setting, replacing the original forward method of huggingface\n        to customized forward method, and add this changing to policy.\"\"\"\n        if not self.pipeline_stage_manager:\n            raise ValueError(\"set_pipeline_forward method can only be called when pipeline parallel is enabled.\")\n        stage_manager = self.pipeline_stage_manager\n        if self.model.__class__.__name__ == \"GPTJModel\":\n            module = self.model\n        else:\n            module = self.model.transformer\n\n        layers_per_stage = stage_manager.distribute_layers(len(module.h))\n        stage_index = stage_manager.get_stage_index(layers_per_stage)\n        method_replacement = {\n            \"forward\": partial(\n                new_forward,\n                stage_manager=stage_manager,\n                stage_index=stage_index,\n                shard_config=self.shard_config,\n            )\n        }\n        self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls)\n\n\n# GPTJModel\nclass GPTJModelPolicy(GPTJPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        from transformers.models.gptj.modeling_gptj import GPTJModel\n\n        policy = super().module_policy()\n\n        if self.pipeline_stage_manager is not None:\n            self.set_pipeline_forward(\n                model_cls=GPTJModel,\n                new_forward=GPTJPipelineForwards.gptj_model_forward,\n                policy=policy,\n            )\n        return policy\n\n    def get_held_layers(self) -> List[nn.Module]:\n        return super().get_held_layers()\n\n    def get_shared_params(self) -> List[Dict[int, Tensor]]:\n        \"\"\"No shared params in GPT2Model.\"\"\"\n        return []\n\n\n# GPTJForCausalLM\nclass GPTJForCausalLMPolicy(GPTJPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        from transformers.models.gptj.modeling_gptj import GPTJForCausalLM\n\n        policy = super().module_policy()\n\n        if self.shard_config.enable_tensor_parallelism:\n            addon_module = {\n                GPTJForCausalLM: ModulePolicyDescription(\n                    sub_module_replacement=[\n                        SubModuleReplacementDescription(\n                            suffix=\"lm_head\",\n                            target_module=col_nn.VocabParallelLMHead1D,\n                            kwargs={\n                                \"gather_output\": True,\n                                \"make_vocab_size_divisible_by\": self.shard_config.make_vocab_size_divisible_by,\n                            },\n                        )\n                    ]\n                )\n            }\n        else:\n            addon_module = {\n                GPTJForCausalLM: ModulePolicyDescription(\n                    sub_module_replacement=[\n                        SubModuleReplacementDescription(\n                            suffix=\"lm_head\",\n                            target_module=col_nn.PaddingLMHead,\n                            kwargs={\"make_vocab_size_divisible_by\": self.shard_config.make_vocab_size_divisible_by},\n                        )\n                    ]\n                )\n            }\n        policy.update(addon_module)\n\n        if self.pipeline_stage_manager is not None:\n            self.set_pipeline_forward(\n                model_cls=GPTJForCausalLM,\n                new_forward=GPTJPipelineForwards.gptj_causallm_model_forward,\n                policy=policy,\n            )\n        return policy\n\n    def get_held_layers(self) -> List[nn.Module]:\n        held_layers = super().get_held_layers()\n        if self.pipeline_stage_manager.is_last_stage():\n            held_layers.append(self.model.lm_head)\n        return held_layers\n\n    def get_shared_params(self) -> List[Dict[int, Tensor]]:\n        \"\"\"The weights of wte and lm_head are shared.\"\"\"\n        module = self.model\n        stage_manager = self.pipeline_stage_manager\n        if stage_manager is not None:\n            if stage_manager.num_stages > 1 and id(module.transformer.wte.weight) == id(module.lm_head.weight):\n                first_stage, last_stage = 0, stage_manager.num_stages - 1\n                return [\n                    {\n                        first_stage: module.transformer.wte.weight,\n                        last_stage: module.lm_head.weight,\n                    }\n                ]\n        return []\n\n\n# GPTJForSequenceClassification\nclass GPTJForSequenceClassificationPolicy(GPTJPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        from transformers.models.gptj.modeling_gptj import GPTJForSequenceClassification\n\n        policy = super().module_policy()\n\n        if self.pipeline_stage_manager is not None:\n            self.set_pipeline_forward(\n                model_cls=GPTJForSequenceClassification,\n                new_forward=GPTJPipelineForwards.gptj_for_sequence_classification_forward,\n                policy=policy,\n            )\n        return policy\n\n    def get_held_layers(self) -> List[nn.Module]:\n        held_layers = super().get_held_layers()\n        if self.pipeline_stage_manager.is_last_stage():\n            held_layers.append(self.model.score)\n        return held_layers\n\n    def get_shared_params(self) -> List[Dict[int, Tensor]]:\n        \"\"\"No shared params in GPTJForSequenceClassification.\"\"\"\n        return []\n\n\n# GPTJForQuestionAnswering\nclass GPTJForQuestionAnsweringPolicy(GPTJPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        from transformers.models.gptj.modeling_gptj import GPTJForQuestionAnswering\n\n        policy = super().module_policy()\n\n        if self.pipeline_stage_manager is not None:\n            self.set_pipeline_forward(\n                model_cls=GPTJForQuestionAnswering,\n                new_forward=GPTJPipelineForwards.gptj_for_question_answering_forward,\n                policy=policy,\n            )\n        return policy\n\n    def get_held_layers(self) -> List[nn.Module]:\n        held_layers = super().get_held_layers()\n        if self.pipeline_stage_manager.is_last_stage():\n            held_layers.append(self.model.qa_outputs)\n        return held_layers\n\n    def get_shared_params(self) -> List[Dict[int, Tensor]]:\n        \"\"\"No shared params in GPT2ForQuestionAnswering.\"\"\"\n        return []\n", "path": "colossalai/shardformer/policies/gptj.py"}]}