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/inference/tensor_parallel/policies/llama.py`
  Content:
  ```
  1 from functools import partial
  2 
  3 import torch
  4 from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel, LlamaRMSNorm
  5 
  6 from colossalai.shardformer.layer import VocabParallelEmbedding1D
  7 from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
  8 # import colossalai
  9 from colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy
  10 
  11 from ..modeling.llama import LlamaInferenceForwards, get_llama_vllm_rmsnorm_forward
  12 
  13 try:
  14     from colossalai.kernel.triton import rmsnorm_forward
  15 
  16     HAS_TRITON_RMSNORM = True
  17 except:
  18     print("you should install triton from https://github.com/openai/triton")
  19     HAS_TRITON_RMSNORM = False
  20 
  21 
  22 def get_triton_rmsnorm_forward():
  23     if HAS_TRITON_RMSNORM:
  24 
  25         def _triton_rmsnorm_forward(self: LlamaRMSNorm, hidden_states: torch.Tensor):
  26             return rmsnorm_forward(hidden_states, self.weight.data, self.variance_epsilon)
  27 
  28         return _triton_rmsnorm_forward
  29     else:
  30         return None
  31 
  32 
  33 class LlamaModelInferPolicy(LlamaForCausalLMPolicy):
  34     def __init__(self) -> None:
  35         super().__init__()
  36 
  37     def module_policy(self):
  38         policy = super().module_policy()
  39 
  40         if self.shard_config.inference_gptq:
  41             from colossalai.inference.quant.gptq.cai_gptq import ColCaiQuantLinear, RowCaiQuantLinear
  42 
  43             decoder_attribute_replacement = {
  44                 "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
  45                 "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size,
  46             }
  47             policy[LlamaDecoderLayer] = ModulePolicyDescription(
  48                 attribute_replacement=decoder_attribute_replacement,
  49                 sub_module_replacement=[
  50                     SubModuleReplacementDescription(
  51                         suffix="self_attn.q_proj",
  52                         target_module=ColCaiQuantLinear,
  53                         kwargs={'split_num': 1},
  54                     ),
  55                     SubModuleReplacementDescription(
  56                         suffix="self_attn.k_proj",
  57                         target_module=ColCaiQuantLinear,
  58                         kwargs={'split_num': 1},
  59                     ),
  60                     SubModuleReplacementDescription(
  61                         suffix="self_attn.v_proj",
  62                         target_module=ColCaiQuantLinear,
  63                         kwargs={'split_num': 1},
  64                     ),
  65                     SubModuleReplacementDescription(
  66                         suffix="self_attn.o_proj",
  67                         target_module=RowCaiQuantLinear,
  68                         kwargs={'split_num': 1},
  69                     ),
  70                     SubModuleReplacementDescription(
  71                         suffix="mlp.gate_proj",
  72                         target_module=ColCaiQuantLinear,
  73                         kwargs={'split_num': 1},
  74                     ),
  75                     SubModuleReplacementDescription(
  76                         suffix="mlp.up_proj",
  77                         target_module=ColCaiQuantLinear,
  78                         kwargs={'split_num': 1},
  79                     ),
  80                     SubModuleReplacementDescription(
  81                         suffix="mlp.down_proj",
  82                         target_module=RowCaiQuantLinear,
  83                         kwargs={'split_num': 1},
  84                     )
  85                 ],
  86             )
  87 
  88         self.shard_config._infer()
  89 
  90         infer_forward = LlamaInferenceForwards.llama_model_forward
  91         method_replacement = {"forward": partial(infer_forward)}
  92         self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=LlamaModel)
  93 
  94         infer_forward = LlamaInferenceForwards.llama_decoder_layer_forward
  95         method_replacement = {"forward": partial(infer_forward)}
  96         self.append_or_create_method_replacement(
  97             description=method_replacement, policy=policy, target_key=LlamaDecoderLayer
  98         )
  99 
  100         infer_forward = LlamaInferenceForwards.llama_flash_attn_kvcache_forward
  101         method_replacement = {"forward": partial(infer_forward)}
  102         self.append_or_create_method_replacement(
  103             description=method_replacement, policy=policy, target_key=LlamaAttention
  104         )
  105 
  106         infer_forward = None
  107         if HAS_TRITON_RMSNORM:
  108             infer_forward = get_triton_rmsnorm_forward()
  109         else:
  110             # NOTE: adding rms_norm from cuda kernels caused precision issue, fix @tiandiao123
  111             infer_forward = get_llama_vllm_rmsnorm_forward()
  112 
  113         if infer_forward is not None:
  114             method_replacement = {"forward": partial(infer_forward)}
  115             self.append_or_create_method_replacement(
  116                 description=method_replacement, policy=policy, target_key=LlamaRMSNorm
  117             )
  118 
  119         return policy
  120
  ```
  --- 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/tensor_parallel/policies/llama.py b/colossalai/inference/tensor_parallel/policies/llama.py
  --- a/colossalai/inference/tensor_parallel/policies/llama.py
  +++ b/colossalai/inference/tensor_parallel/policies/llama.py
  @@ -3,8 +3,8 @@
   import torch
   from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel, LlamaRMSNorm
   
  -from colossalai.shardformer.layer import VocabParallelEmbedding1D
  -from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
  +from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription
  +
   # import colossalai
   from colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy
   
  @@ -50,38 +50,38 @@
                       SubModuleReplacementDescription(
                           suffix="self_attn.q_proj",
                           target_module=ColCaiQuantLinear,
  -                        kwargs={'split_num': 1},
  +                        kwargs={"split_num": 1},
                       ),
                       SubModuleReplacementDescription(
                           suffix="self_attn.k_proj",
                           target_module=ColCaiQuantLinear,
  -                        kwargs={'split_num': 1},
  +                        kwargs={"split_num": 1},
                       ),
                       SubModuleReplacementDescription(
                           suffix="self_attn.v_proj",
                           target_module=ColCaiQuantLinear,
  -                        kwargs={'split_num': 1},
  +                        kwargs={"split_num": 1},
                       ),
                       SubModuleReplacementDescription(
                           suffix="self_attn.o_proj",
                           target_module=RowCaiQuantLinear,
  -                        kwargs={'split_num': 1},
  +                        kwargs={"split_num": 1},
                       ),
                       SubModuleReplacementDescription(
                           suffix="mlp.gate_proj",
                           target_module=ColCaiQuantLinear,
  -                        kwargs={'split_num': 1},
  +                        kwargs={"split_num": 1},
                       ),
                       SubModuleReplacementDescription(
                           suffix="mlp.up_proj",
                           target_module=ColCaiQuantLinear,
  -                        kwargs={'split_num': 1},
  +                        kwargs={"split_num": 1},
                       ),
                       SubModuleReplacementDescription(
                           suffix="mlp.down_proj",
                           target_module=RowCaiQuantLinear,
  -                        kwargs={'split_num': 1},
  -                    )
  +                        kwargs={"split_num": 1},
  +                    ),
                   ],
               )
  
verification_info: {"golden_diff": "diff --git a/colossalai/inference/tensor_parallel/policies/llama.py b/colossalai/inference/tensor_parallel/policies/llama.py\n--- a/colossalai/inference/tensor_parallel/policies/llama.py\n+++ b/colossalai/inference/tensor_parallel/policies/llama.py\n@@ -3,8 +3,8 @@\n import torch\n from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel, LlamaRMSNorm\n \n-from colossalai.shardformer.layer import VocabParallelEmbedding1D\n-from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription\n+from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription\n+\n # import colossalai\n from colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy\n \n@@ -50,38 +50,38 @@\n                     SubModuleReplacementDescription(\n                         suffix=\"self_attn.q_proj\",\n                         target_module=ColCaiQuantLinear,\n-                        kwargs={'split_num': 1},\n+                        kwargs={\"split_num\": 1},\n                     ),\n                     SubModuleReplacementDescription(\n                         suffix=\"self_attn.k_proj\",\n                         target_module=ColCaiQuantLinear,\n-                        kwargs={'split_num': 1},\n+                        kwargs={\"split_num\": 1},\n                     ),\n                     SubModuleReplacementDescription(\n                         suffix=\"self_attn.v_proj\",\n                         target_module=ColCaiQuantLinear,\n-                        kwargs={'split_num': 1},\n+                        kwargs={\"split_num\": 1},\n                     ),\n                     SubModuleReplacementDescription(\n                         suffix=\"self_attn.o_proj\",\n                         target_module=RowCaiQuantLinear,\n-                        kwargs={'split_num': 1},\n+                        kwargs={\"split_num\": 1},\n                     ),\n                     SubModuleReplacementDescription(\n                         suffix=\"mlp.gate_proj\",\n                         target_module=ColCaiQuantLinear,\n-                        kwargs={'split_num': 1},\n+                        kwargs={\"split_num\": 1},\n                     ),\n                     SubModuleReplacementDescription(\n                         suffix=\"mlp.up_proj\",\n                         target_module=ColCaiQuantLinear,\n-                        kwargs={'split_num': 1},\n+                        kwargs={\"split_num\": 1},\n                     ),\n                     SubModuleReplacementDescription(\n                         suffix=\"mlp.down_proj\",\n                         target_module=RowCaiQuantLinear,\n-                        kwargs={'split_num': 1},\n-                    )\n+                        kwargs={\"split_num\": 1},\n+                    ),\n                 ],\n             )\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "from functools import partial\n\nimport torch\nfrom transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel, LlamaRMSNorm\n\nfrom colossalai.shardformer.layer import VocabParallelEmbedding1D\nfrom colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription\n# import colossalai\nfrom colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy\n\nfrom ..modeling.llama import LlamaInferenceForwards, get_llama_vllm_rmsnorm_forward\n\ntry:\n    from colossalai.kernel.triton import rmsnorm_forward\n\n    HAS_TRITON_RMSNORM = True\nexcept:\n    print(\"you should install triton from https://github.com/openai/triton\")\n    HAS_TRITON_RMSNORM = False\n\n\ndef get_triton_rmsnorm_forward():\n    if HAS_TRITON_RMSNORM:\n\n        def _triton_rmsnorm_forward(self: LlamaRMSNorm, hidden_states: torch.Tensor):\n            return rmsnorm_forward(hidden_states, self.weight.data, self.variance_epsilon)\n\n        return _triton_rmsnorm_forward\n    else:\n        return None\n\n\nclass LlamaModelInferPolicy(LlamaForCausalLMPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        policy = super().module_policy()\n\n        if self.shard_config.inference_gptq:\n            from colossalai.inference.quant.gptq.cai_gptq import ColCaiQuantLinear, RowCaiQuantLinear\n\n            decoder_attribute_replacement = {\n                \"self_attn.hidden_size\": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,\n                \"self_attn.num_heads\": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size,\n            }\n            policy[LlamaDecoderLayer] = ModulePolicyDescription(\n                attribute_replacement=decoder_attribute_replacement,\n                sub_module_replacement=[\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.q_proj\",\n                        target_module=ColCaiQuantLinear,\n                        kwargs={'split_num': 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.k_proj\",\n                        target_module=ColCaiQuantLinear,\n                        kwargs={'split_num': 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.v_proj\",\n                        target_module=ColCaiQuantLinear,\n                        kwargs={'split_num': 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.o_proj\",\n                        target_module=RowCaiQuantLinear,\n                        kwargs={'split_num': 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.gate_proj\",\n                        target_module=ColCaiQuantLinear,\n                        kwargs={'split_num': 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.up_proj\",\n                        target_module=ColCaiQuantLinear,\n                        kwargs={'split_num': 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.down_proj\",\n                        target_module=RowCaiQuantLinear,\n                        kwargs={'split_num': 1},\n                    )\n                ],\n            )\n\n        self.shard_config._infer()\n\n        infer_forward = LlamaInferenceForwards.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 = LlamaInferenceForwards.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 = LlamaInferenceForwards.llama_flash_attn_kvcache_forward\n        method_replacement = {\"forward\": partial(infer_forward)}\n        self.append_or_create_method_replacement(\n            description=method_replacement, policy=policy, target_key=LlamaAttention\n        )\n\n        infer_forward = None\n        if HAS_TRITON_RMSNORM:\n            infer_forward = get_triton_rmsnorm_forward()\n        else:\n            # NOTE: adding rms_norm from cuda kernels caused precision issue, fix @tiandiao123\n            infer_forward = get_llama_vllm_rmsnorm_forward()\n\n        if infer_forward is not None:\n            method_replacement = {\"forward\": partial(infer_forward)}\n            self.append_or_create_method_replacement(\n                description=method_replacement, policy=policy, target_key=LlamaRMSNorm\n            )\n\n        return policy\n", "path": "colossalai/inference/tensor_parallel/policies/llama.py"}], "after_files": [{"content": "from functools import partial\n\nimport torch\nfrom transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel, LlamaRMSNorm\n\nfrom colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription\n\n# import colossalai\nfrom colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy\n\nfrom ..modeling.llama import LlamaInferenceForwards, get_llama_vllm_rmsnorm_forward\n\ntry:\n    from colossalai.kernel.triton import rmsnorm_forward\n\n    HAS_TRITON_RMSNORM = True\nexcept:\n    print(\"you should install triton from https://github.com/openai/triton\")\n    HAS_TRITON_RMSNORM = False\n\n\ndef get_triton_rmsnorm_forward():\n    if HAS_TRITON_RMSNORM:\n\n        def _triton_rmsnorm_forward(self: LlamaRMSNorm, hidden_states: torch.Tensor):\n            return rmsnorm_forward(hidden_states, self.weight.data, self.variance_epsilon)\n\n        return _triton_rmsnorm_forward\n    else:\n        return None\n\n\nclass LlamaModelInferPolicy(LlamaForCausalLMPolicy):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def module_policy(self):\n        policy = super().module_policy()\n\n        if self.shard_config.inference_gptq:\n            from colossalai.inference.quant.gptq.cai_gptq import ColCaiQuantLinear, RowCaiQuantLinear\n\n            decoder_attribute_replacement = {\n                \"self_attn.hidden_size\": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,\n                \"self_attn.num_heads\": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size,\n            }\n            policy[LlamaDecoderLayer] = ModulePolicyDescription(\n                attribute_replacement=decoder_attribute_replacement,\n                sub_module_replacement=[\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.q_proj\",\n                        target_module=ColCaiQuantLinear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.k_proj\",\n                        target_module=ColCaiQuantLinear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.v_proj\",\n                        target_module=ColCaiQuantLinear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.o_proj\",\n                        target_module=RowCaiQuantLinear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.gate_proj\",\n                        target_module=ColCaiQuantLinear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.up_proj\",\n                        target_module=ColCaiQuantLinear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.down_proj\",\n                        target_module=RowCaiQuantLinear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                ],\n            )\n\n        self.shard_config._infer()\n\n        infer_forward = LlamaInferenceForwards.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 = LlamaInferenceForwards.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 = LlamaInferenceForwards.llama_flash_attn_kvcache_forward\n        method_replacement = {\"forward\": partial(infer_forward)}\n        self.append_or_create_method_replacement(\n            description=method_replacement, policy=policy, target_key=LlamaAttention\n        )\n\n        infer_forward = None\n        if HAS_TRITON_RMSNORM:\n            infer_forward = get_triton_rmsnorm_forward()\n        else:\n            # NOTE: adding rms_norm from cuda kernels caused precision issue, fix @tiandiao123\n            infer_forward = get_llama_vllm_rmsnorm_forward()\n\n        if infer_forward is not None:\n            method_replacement = {\"forward\": partial(infer_forward)}\n            self.append_or_create_method_replacement(\n                description=method_replacement, policy=policy, target_key=LlamaRMSNorm\n            )\n\n        return policy\n", "path": "colossalai/inference/tensor_parallel/policies/llama.py"}]}