prompt: 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()
  
  ```
  
golden_diff: 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"
  
verification_info: {"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"}]}