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/shard/shard_config.py`
Content:
```
1 from dataclasses import dataclass, field
2 from typing import Any, Dict, Optional
3 
4 import torch.distributed as dist
5 from torch.distributed import ProcessGroup
6 
7 from colossalai.pipeline.stage_manager import PipelineStageManager
8 
9 __all__ = ["ShardConfig"]
10 
11 
12 @dataclass
13 class ShardConfig:
14     r"""
15     The config for sharding the huggingface model
16 
17     Args:
18         tensor_parallel_process_group (Optional[ProcessGroup]): The process group of tensor parallelism, it's necessary when using tensor parallel. Defaults to None, which is the global process group.
19         pipeline_stage_manager (Optional[PipelineStageManager]): If using pipeline parallelism, it's necessary to specify a pipeline stage manager for inter-process communication in pipeline parallelism. Defaults to None, which means not using pipeline parallelism.
20         enable_tensor_parallelism (bool): Whether to use tensor parallelism. Defaults to True.
21         enable_fused_normalization (bool): Whether to use fused layernorm. Defaults to False.
22         enable_flash_attention (bool, optional): Whether to switch on flash attention. Defaults to False.
23         enable_jit_fused (bool, optional): Whether to switch on JIT fused operators. Defaults to False.
24         enable_sequence_parallelism (bool): Whether to turn on sequence parallelism, which partitions non-tensor-parallel regions along the sequence dimension. Defaults to False.
25         enable_sequence_overlap (bool): Whether to turn on sequence overlap, which overlap the computation and communication in sequence parallelism. It can only be used when enable_sequence_parallelism is True. Defaults to False.
26         enable_all_optimization (bool): Whether to turn on all optimization tools including 'fused normalization', 'flash attention', 'JIT fused operators', 'sequence parallelism' and 'sequence overlap'. Defaults to False.
27     """
28     tensor_parallel_process_group: Optional[ProcessGroup] = None
29     pipeline_stage_manager: Optional[PipelineStageManager] = None
30     enable_tensor_parallelism: bool = True
31     enable_fused_normalization: bool = False
32     enable_flash_attention: bool = False
33     enable_jit_fused: bool = False
34     enable_all_optimization: bool = False
35     enable_sequence_parallelism: bool = False
36     enable_sequence_overlap: bool = False
37     parallel_output = True
38     extra_kwargs: Dict[str, Any] = field(default_factory=dict)
39     # pipeline_parallel_size: int
40     # data_parallel_size: int
41     # tensor_parallel_mode: Literal['1d', '2d', '2.5d', '3d']
42 
43     @property
44     def tensor_parallel_size(self):
45         return self._tensor_parallel_size
46 
47     def __post_init__(self):
48         if not self.enable_tensor_parallelism and self.enable_sequence_parallelism:
49             raise ValueError(
50                 "enable_sequence_parallelism can only be set to True when enable_tensor_parallelism is True"
51             )
52         if not self.enable_sequence_parallelism and self.enable_sequence_overlap:
53             raise ValueError("enable_sequence_overlap can only be set to True when enable_sequence_parallelism is True")
54         if not self.enable_tensor_parallelism:
55             self._tensor_parallel_size = 1
56         else:
57             # get the parallel size
58             self._tensor_parallel_size = dist.get_world_size(self.tensor_parallel_process_group)
59         # turn on all optimization if all_optimization is set to True
60         if self.enable_all_optimization:
61             self._turn_on_all_optimization()
62 
63     def _turn_on_all_optimization(self):
64         """
65         Turn on all optimization.
66         """
67         # you can add all the optimization flag here
68         self.enable_fused_normalization = True
69         self.enable_flash_attention = True
70         self.enable_jit_fused = True
71         self.enable_sequence_parallelism = True
72         self.enable_sequence_overlap = True
73 
74     def _infer(self):
75         """
76         Set default params for inference.
77         """
78         # assert self.pipeline_stage_manager is None, "pipeline parallelism is not supported in inference for now"
79
```

Path: `colossalai/inference/modeling/policy/nopadding_llama.py`
Content:
```
1 from functools import partial
2 
3 from torch.nn import Parameter
4 from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, LlamaRMSNorm
5 
6 from colossalai.inference.modeling.models.nopadding_llama import (
7     NopadLlamaAttention,
8     NopadLlamaMLP,
9     llama_causal_lm_forward,
10     llama_decoder_layer_forward,
11     llama_model_forward,
12     llama_rmsnorm_forward,
13 )
14 from colossalai.inference.utils import init_to_get_rotary
15 from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription
16 
17 # import colossalai
18 from colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy
19 
20 
21 class NoPaddingLlamaModelInferPolicy(LlamaForCausalLMPolicy):
22     def __init__(self) -> None:
23         super().__init__()
24 
25     def module_policy(self):
26         policy = super().module_policy()
27 
28         decoder_attribute_replacement = {
29             "lm_head.weight": Parameter(self.model.lm_head.weight.transpose(0, 1), requires_grad=False),
30         }
31         policy[LlamaForCausalLM] = ModulePolicyDescription(
32             attribute_replacement=decoder_attribute_replacement,
33         )
34 
35         policy[LlamaDecoderLayer] = ModulePolicyDescription(
36             sub_module_replacement=[
37                 SubModuleReplacementDescription(
38                     suffix="mlp",
39                     target_module=NopadLlamaMLP,
40                 ),
41                 SubModuleReplacementDescription(
42                     suffix="self_attn",
43                     target_module=NopadLlamaAttention,
44                 ),
45             ]
46         )
47 
48         self.shard_config._infer()
49 
50         infer_forward = llama_causal_lm_forward
51         method_replacement = {"forward": partial(infer_forward)}
52         self.append_or_create_method_replacement(
53             description=method_replacement, policy=policy, target_key=LlamaForCausalLM
54         )
55 
56         infer_forward = llama_model_forward
57         method_replacement = {"forward": partial(infer_forward)}
58         self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=LlamaModel)
59 
60         infer_forward = llama_decoder_layer_forward
61         method_replacement = {"forward": partial(infer_forward)}
62         self.append_or_create_method_replacement(
63             description=method_replacement, policy=policy, target_key=LlamaDecoderLayer
64         )
65 
66         infer_forward = llama_rmsnorm_forward
67         method_replacement = {"forward": partial(infer_forward)}
68         self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=LlamaRMSNorm)
69 
70         return policy
71 
72     def postprocess(self):
73         init_to_get_rotary(self.model.model)
74         return self.model
75
```
--- 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/inference/modeling/policy/nopadding_llama.py b/colossalai/inference/modeling/policy/nopadding_llama.py
--- a/colossalai/inference/modeling/policy/nopadding_llama.py
+++ b/colossalai/inference/modeling/policy/nopadding_llama.py
@@ -1,5 +1,3 @@
-from functools import partial
-
 from torch.nn import Parameter
 from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, LlamaRMSNorm
 
@@ -13,8 +11,6 @@
 )
 from colossalai.inference.utils import init_to_get_rotary
 from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription
-
-# import colossalai
 from colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy
 
 
@@ -45,27 +41,18 @@
             ]
         )
 
-        self.shard_config._infer()
-
-        infer_forward = llama_causal_lm_forward
-        method_replacement = {"forward": partial(infer_forward)}
         self.append_or_create_method_replacement(
-            description=method_replacement, policy=policy, target_key=LlamaForCausalLM
+            description={"forward": llama_causal_lm_forward}, policy=policy, target_key=LlamaForCausalLM
         )
-
-        infer_forward = llama_model_forward
-        method_replacement = {"forward": partial(infer_forward)}
-        self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=LlamaModel)
-
-        infer_forward = llama_decoder_layer_forward
-        method_replacement = {"forward": partial(infer_forward)}
         self.append_or_create_method_replacement(
-            description=method_replacement, policy=policy, target_key=LlamaDecoderLayer
+            description={"forward": llama_model_forward}, policy=policy, target_key=LlamaModel
+        )
+        self.append_or_create_method_replacement(
+            description={"forward": llama_decoder_layer_forward}, policy=policy, target_key=LlamaDecoderLayer
+        )
+        self.append_or_create_method_replacement(
+            description={"forward": llama_rmsnorm_forward}, policy=policy, target_key=LlamaRMSNorm
         )
-
-        infer_forward = llama_rmsnorm_forward
-        method_replacement = {"forward": partial(infer_forward)}
-        self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=LlamaRMSNorm)
 
         return policy
 
diff --git a/colossalai/shardformer/shard/shard_config.py b/colossalai/shardformer/shard/shard_config.py
--- a/colossalai/shardformer/shard/shard_config.py
+++ b/colossalai/shardformer/shard/shard_config.py
@@ -36,8 +36,6 @@
     enable_sequence_overlap: bool = False
     parallel_output = True
     extra_kwargs: Dict[str, Any] = field(default_factory=dict)
-    # pipeline_parallel_size: int
-    # data_parallel_size: int
     # tensor_parallel_mode: Literal['1d', '2d', '2.5d', '3d']
 
     @property
@@ -70,9 +68,3 @@
         self.enable_jit_fused = True
         self.enable_sequence_parallelism = True
         self.enable_sequence_overlap = True
-
-    def _infer(self):
-        """
-        Set default params for inference.
-        """
-        # assert self.pipeline_stage_manager is None, "pipeline parallelism is not supported in inference for now"

{"golden_diff": "diff --git a/colossalai/inference/modeling/policy/nopadding_llama.py b/colossalai/inference/modeling/policy/nopadding_llama.py\n--- a/colossalai/inference/modeling/policy/nopadding_llama.py\n+++ b/colossalai/inference/modeling/policy/nopadding_llama.py\n@@ -1,5 +1,3 @@\n-from functools import partial\n-\n from torch.nn import Parameter\n from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, LlamaRMSNorm\n \n@@ -13,8 +11,6 @@\n )\n from colossalai.inference.utils import init_to_get_rotary\n from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription\n-\n-# import colossalai\n from colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy\n \n \n@@ -45,27 +41,18 @@\n             ]\n         )\n \n-        self.shard_config._infer()\n-\n-        infer_forward = llama_causal_lm_forward\n-        method_replacement = {\"forward\": partial(infer_forward)}\n         self.append_or_create_method_replacement(\n-            description=method_replacement, policy=policy, target_key=LlamaForCausalLM\n+            description={\"forward\": llama_causal_lm_forward}, policy=policy, target_key=LlamaForCausalLM\n         )\n-\n-        infer_forward = llama_model_forward\n-        method_replacement = {\"forward\": partial(infer_forward)}\n-        self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=LlamaModel)\n-\n-        infer_forward = llama_decoder_layer_forward\n-        method_replacement = {\"forward\": partial(infer_forward)}\n         self.append_or_create_method_replacement(\n-            description=method_replacement, policy=policy, target_key=LlamaDecoderLayer\n+            description={\"forward\": llama_model_forward}, policy=policy, target_key=LlamaModel\n+        )\n+        self.append_or_create_method_replacement(\n+            description={\"forward\": llama_decoder_layer_forward}, policy=policy, target_key=LlamaDecoderLayer\n+        )\n+        self.append_or_create_method_replacement(\n+            description={\"forward\": llama_rmsnorm_forward}, policy=policy, target_key=LlamaRMSNorm\n         )\n-\n-        infer_forward = llama_rmsnorm_forward\n-        method_replacement = {\"forward\": partial(infer_forward)}\n-        self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=LlamaRMSNorm)\n \n         return policy\n \ndiff --git a/colossalai/shardformer/shard/shard_config.py b/colossalai/shardformer/shard/shard_config.py\n--- a/colossalai/shardformer/shard/shard_config.py\n+++ b/colossalai/shardformer/shard/shard_config.py\n@@ -36,8 +36,6 @@\n     enable_sequence_overlap: bool = False\n     parallel_output = True\n     extra_kwargs: Dict[str, Any] = field(default_factory=dict)\n-    # pipeline_parallel_size: int\n-    # data_parallel_size: int\n     # tensor_parallel_mode: Literal['1d', '2d', '2.5d', '3d']\n \n     @property\n@@ -70,9 +68,3 @@\n         self.enable_jit_fused = True\n         self.enable_sequence_parallelism = True\n         self.enable_sequence_overlap = True\n-\n-    def _infer(self):\n-        \"\"\"\n-        Set default params for inference.\n-        \"\"\"\n-        # assert self.pipeline_stage_manager is None, \"pipeline parallelism is not supported in inference for now\"\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "from dataclasses import dataclass, field\nfrom typing import Any, Dict, Optional\n\nimport torch.distributed as dist\nfrom torch.distributed import ProcessGroup\n\nfrom colossalai.pipeline.stage_manager import PipelineStageManager\n\n__all__ = [\"ShardConfig\"]\n\n\n@dataclass\nclass ShardConfig:\n    r\"\"\"\n    The config for sharding the huggingface model\n\n    Args:\n        tensor_parallel_process_group (Optional[ProcessGroup]): The process group of tensor parallelism, it's necessary when using tensor parallel. Defaults to None, which is the global process group.\n        pipeline_stage_manager (Optional[PipelineStageManager]): If using pipeline parallelism, it's necessary to specify a pipeline stage manager for inter-process communication in pipeline parallelism. Defaults to None, which means not using pipeline parallelism.\n        enable_tensor_parallelism (bool): Whether to use tensor parallelism. Defaults to True.\n        enable_fused_normalization (bool): Whether to use fused layernorm. Defaults to False.\n        enable_flash_attention (bool, optional): Whether to switch on flash attention. Defaults to False.\n        enable_jit_fused (bool, optional): Whether to switch on JIT fused operators. Defaults to False.\n        enable_sequence_parallelism (bool): Whether to turn on sequence parallelism, which partitions non-tensor-parallel regions along the sequence dimension. Defaults to False.\n        enable_sequence_overlap (bool): Whether to turn on sequence overlap, which overlap the computation and communication in sequence parallelism. It can only be used when enable_sequence_parallelism is True. Defaults to False.\n        enable_all_optimization (bool): Whether to turn on all optimization tools including 'fused normalization', 'flash attention', 'JIT fused operators', 'sequence parallelism' and 'sequence overlap'. Defaults to False.\n    \"\"\"\n    tensor_parallel_process_group: Optional[ProcessGroup] = None\n    pipeline_stage_manager: Optional[PipelineStageManager] = None\n    enable_tensor_parallelism: bool = True\n    enable_fused_normalization: bool = False\n    enable_flash_attention: bool = False\n    enable_jit_fused: bool = False\n    enable_all_optimization: bool = False\n    enable_sequence_parallelism: bool = False\n    enable_sequence_overlap: bool = False\n    parallel_output = True\n    extra_kwargs: Dict[str, Any] = field(default_factory=dict)\n    # pipeline_parallel_size: int\n    # data_parallel_size: int\n    # tensor_parallel_mode: Literal['1d', '2d', '2.5d', '3d']\n\n    @property\n    def tensor_parallel_size(self):\n        return self._tensor_parallel_size\n\n    def __post_init__(self):\n        if not self.enable_tensor_parallelism and self.enable_sequence_parallelism:\n            raise ValueError(\n                \"enable_sequence_parallelism can only be set to True when enable_tensor_parallelism is True\"\n            )\n        if not self.enable_sequence_parallelism and self.enable_sequence_overlap:\n            raise ValueError(\"enable_sequence_overlap can only be set to True when enable_sequence_parallelism is True\")\n        if not self.enable_tensor_parallelism:\n            self._tensor_parallel_size = 1\n        else:\n            # get the parallel size\n            self._tensor_parallel_size = dist.get_world_size(self.tensor_parallel_process_group)\n        # turn on all optimization if all_optimization is set to True\n        if self.enable_all_optimization:\n            self._turn_on_all_optimization()\n\n    def _turn_on_all_optimization(self):\n        \"\"\"\n        Turn on all optimization.\n        \"\"\"\n        # you can add all the optimization flag here\n        self.enable_fused_normalization = True\n        self.enable_flash_attention = True\n        self.enable_jit_fused = True\n        self.enable_sequence_parallelism = True\n        self.enable_sequence_overlap = True\n\n    def _infer(self):\n        \"\"\"\n        Set default params for inference.\n        \"\"\"\n        # assert self.pipeline_stage_manager is None, \"pipeline parallelism is not supported in inference for now\"\n", "path": "colossalai/shardformer/shard/shard_config.py"}, {"content": "from functools import partial\n\nfrom torch.nn import Parameter\nfrom transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, LlamaRMSNorm\n\nfrom colossalai.inference.modeling.models.nopadding_llama import (\n    NopadLlamaAttention,\n    NopadLlamaMLP,\n    llama_causal_lm_forward,\n    llama_decoder_layer_forward,\n    llama_model_forward,\n    llama_rmsnorm_forward,\n)\nfrom colossalai.inference.utils import init_to_get_rotary\nfrom colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription\n\n# import colossalai\nfrom colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy\n\n\nclass NoPaddingLlamaModelInferPolicy(LlamaForCausalLMPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        policy = super().module_policy()\n\n        decoder_attribute_replacement = {\n            \"lm_head.weight\": Parameter(self.model.lm_head.weight.transpose(0, 1), requires_grad=False),\n        }\n        policy[LlamaForCausalLM] = ModulePolicyDescription(\n            attribute_replacement=decoder_attribute_replacement,\n        )\n\n        policy[LlamaDecoderLayer] = ModulePolicyDescription(\n            sub_module_replacement=[\n                SubModuleReplacementDescription(\n                    suffix=\"mlp\",\n                    target_module=NopadLlamaMLP,\n                ),\n                SubModuleReplacementDescription(\n                    suffix=\"self_attn\",\n                    target_module=NopadLlamaAttention,\n                ),\n            ]\n        )\n\n        self.shard_config._infer()\n\n        infer_forward = llama_causal_lm_forward\n        method_replacement = {\"forward\": partial(infer_forward)}\n        self.append_or_create_method_replacement(\n            description=method_replacement, policy=policy, target_key=LlamaForCausalLM\n        )\n\n        infer_forward = llama_model_forward\n        method_replacement = {\"forward\": partial(infer_forward)}\n        self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=LlamaModel)\n\n        infer_forward = llama_decoder_layer_forward\n        method_replacement = {\"forward\": partial(infer_forward)}\n        self.append_or_create_method_replacement(\n            description=method_replacement, policy=policy, target_key=LlamaDecoderLayer\n        )\n\n        infer_forward = llama_rmsnorm_forward\n        method_replacement = {\"forward\": partial(infer_forward)}\n        self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=LlamaRMSNorm)\n\n        return policy\n\n    def postprocess(self):\n        init_to_get_rotary(self.model.model)\n        return self.model\n", "path": "colossalai/inference/modeling/policy/nopadding_llama.py"}], "after_files": [{"content": "from dataclasses import dataclass, field\nfrom typing import Any, Dict, Optional\n\nimport torch.distributed as dist\nfrom torch.distributed import ProcessGroup\n\nfrom colossalai.pipeline.stage_manager import PipelineStageManager\n\n__all__ = [\"ShardConfig\"]\n\n\n@dataclass\nclass ShardConfig:\n    r\"\"\"\n    The config for sharding the huggingface model\n\n    Args:\n        tensor_parallel_process_group (Optional[ProcessGroup]): The process group of tensor parallelism, it's necessary when using tensor parallel. Defaults to None, which is the global process group.\n        pipeline_stage_manager (Optional[PipelineStageManager]): If using pipeline parallelism, it's necessary to specify a pipeline stage manager for inter-process communication in pipeline parallelism. Defaults to None, which means not using pipeline parallelism.\n        enable_tensor_parallelism (bool): Whether to use tensor parallelism. Defaults to True.\n        enable_fused_normalization (bool): Whether to use fused layernorm. Defaults to False.\n        enable_flash_attention (bool, optional): Whether to switch on flash attention. Defaults to False.\n        enable_jit_fused (bool, optional): Whether to switch on JIT fused operators. Defaults to False.\n        enable_sequence_parallelism (bool): Whether to turn on sequence parallelism, which partitions non-tensor-parallel regions along the sequence dimension. Defaults to False.\n        enable_sequence_overlap (bool): Whether to turn on sequence overlap, which overlap the computation and communication in sequence parallelism. It can only be used when enable_sequence_parallelism is True. Defaults to False.\n        enable_all_optimization (bool): Whether to turn on all optimization tools including 'fused normalization', 'flash attention', 'JIT fused operators', 'sequence parallelism' and 'sequence overlap'. Defaults to False.\n    \"\"\"\n    tensor_parallel_process_group: Optional[ProcessGroup] = None\n    pipeline_stage_manager: Optional[PipelineStageManager] = None\n    enable_tensor_parallelism: bool = True\n    enable_fused_normalization: bool = False\n    enable_flash_attention: bool = False\n    enable_jit_fused: bool = False\n    enable_all_optimization: bool = False\n    enable_sequence_parallelism: bool = False\n    enable_sequence_overlap: bool = False\n    parallel_output = True\n    extra_kwargs: Dict[str, Any] = field(default_factory=dict)\n    # tensor_parallel_mode: Literal['1d', '2d', '2.5d', '3d']\n\n    @property\n    def tensor_parallel_size(self):\n        return self._tensor_parallel_size\n\n    def __post_init__(self):\n        if not self.enable_tensor_parallelism and self.enable_sequence_parallelism:\n            raise ValueError(\n                \"enable_sequence_parallelism can only be set to True when enable_tensor_parallelism is True\"\n            )\n        if not self.enable_sequence_parallelism and self.enable_sequence_overlap:\n            raise ValueError(\"enable_sequence_overlap can only be set to True when enable_sequence_parallelism is True\")\n        if not self.enable_tensor_parallelism:\n            self._tensor_parallel_size = 1\n        else:\n            # get the parallel size\n            self._tensor_parallel_size = dist.get_world_size(self.tensor_parallel_process_group)\n        # turn on all optimization if all_optimization is set to True\n        if self.enable_all_optimization:\n            self._turn_on_all_optimization()\n\n    def _turn_on_all_optimization(self):\n        \"\"\"\n        Turn on all optimization.\n        \"\"\"\n        # you can add all the optimization flag here\n        self.enable_fused_normalization = True\n        self.enable_flash_attention = True\n        self.enable_jit_fused = True\n        self.enable_sequence_parallelism = True\n        self.enable_sequence_overlap = True\n", "path": "colossalai/shardformer/shard/shard_config.py"}, {"content": "from torch.nn import Parameter\nfrom transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, LlamaRMSNorm\n\nfrom colossalai.inference.modeling.models.nopadding_llama import (\n    NopadLlamaAttention,\n    NopadLlamaMLP,\n    llama_causal_lm_forward,\n    llama_decoder_layer_forward,\n    llama_model_forward,\n    llama_rmsnorm_forward,\n)\nfrom colossalai.inference.utils import init_to_get_rotary\nfrom colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription\nfrom colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy\n\n\nclass NoPaddingLlamaModelInferPolicy(LlamaForCausalLMPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        policy = super().module_policy()\n\n        decoder_attribute_replacement = {\n            \"lm_head.weight\": Parameter(self.model.lm_head.weight.transpose(0, 1), requires_grad=False),\n        }\n        policy[LlamaForCausalLM] = ModulePolicyDescription(\n            attribute_replacement=decoder_attribute_replacement,\n        )\n\n        policy[LlamaDecoderLayer] = ModulePolicyDescription(\n            sub_module_replacement=[\n                SubModuleReplacementDescription(\n                    suffix=\"mlp\",\n                    target_module=NopadLlamaMLP,\n                ),\n                SubModuleReplacementDescription(\n                    suffix=\"self_attn\",\n                    target_module=NopadLlamaAttention,\n                ),\n            ]\n        )\n\n        self.append_or_create_method_replacement(\n            description={\"forward\": llama_causal_lm_forward}, policy=policy, target_key=LlamaForCausalLM\n        )\n        self.append_or_create_method_replacement(\n            description={\"forward\": llama_model_forward}, policy=policy, target_key=LlamaModel\n        )\n        self.append_or_create_method_replacement(\n            description={\"forward\": llama_decoder_layer_forward}, policy=policy, target_key=LlamaDecoderLayer\n        )\n        self.append_or_create_method_replacement(\n            description={\"forward\": llama_rmsnorm_forward}, policy=policy, target_key=LlamaRMSNorm\n        )\n\n        return policy\n\n    def postprocess(self):\n        init_to_get_rotary(self.model.model)\n        return self.model\n", "path": "colossalai/inference/modeling/policy/nopadding_llama.py"}]}