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()

```

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},
+                    ),
                 ],
             )

{"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"}]}