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/engine/policies/llama.py`
Content:
```
1 from functools import partial
2 from typing import List
3 
4 import torch
5 from torch.nn import Module
6 from transformers.models.llama.modeling_llama import (
7     LlamaAttention,
8     LlamaDecoderLayer,
9     LlamaForCausalLM,
10     LlamaModel,
11     LlamaRMSNorm,
12 )
13 
14 from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription
15 
16 # import colossalai
17 from colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy
18 
19 from ..modeling._utils import init_to_get_rotary
20 from ..modeling.llama import LlamaInferenceForwards
21 
22 try:
23     from colossalai.kernel.triton import rmsnorm_forward
24 
25     HAS_TRITON_RMSNORM = True
26 except:
27     print("you should install triton from https://github.com/openai/triton")
28     HAS_TRITON_RMSNORM = False
29 
30 
31 def get_triton_rmsnorm_forward():
32     if HAS_TRITON_RMSNORM:
33 
34         def _triton_rmsnorm_forward(self: LlamaRMSNorm, hidden_states: torch.Tensor):
35             return rmsnorm_forward(hidden_states, self.weight.data, self.variance_epsilon)
36 
37         return _triton_rmsnorm_forward
38     else:
39         return None
40 
41 
42 class LlamaModelInferPolicy(LlamaForCausalLMPolicy):
43     def __init__(self) -> None:
44         super().__init__()
45 
46     def module_policy(self):
47         policy = super().module_policy()
48         decoder_attribute_replacement = {
49             "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
50             "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size,
51             "self_attn.num_key_value_heads": self.model.config.num_key_value_heads
52             // self.shard_config.tensor_parallel_size,
53         }
54         if self.shard_config.extra_kwargs.get("quant", None) == "gptq":
55             from colossalai.inference.quant.gptq.cai_gptq import ColCaiQuantLinear, RowCaiQuantLinear
56 
57             policy[LlamaDecoderLayer] = ModulePolicyDescription(
58                 attribute_replacement=decoder_attribute_replacement,
59                 sub_module_replacement=[
60                     SubModuleReplacementDescription(
61                         suffix="self_attn.q_proj",
62                         target_module=ColCaiQuantLinear,
63                         kwargs={"split_num": 1},
64                     ),
65                     SubModuleReplacementDescription(
66                         suffix="self_attn.k_proj",
67                         target_module=ColCaiQuantLinear,
68                         kwargs={"split_num": 1},
69                     ),
70                     SubModuleReplacementDescription(
71                         suffix="self_attn.v_proj",
72                         target_module=ColCaiQuantLinear,
73                         kwargs={"split_num": 1},
74                     ),
75                     SubModuleReplacementDescription(
76                         suffix="self_attn.o_proj",
77                         target_module=RowCaiQuantLinear,
78                         kwargs={"split_num": 1},
79                     ),
80                     SubModuleReplacementDescription(
81                         suffix="mlp.gate_proj",
82                         target_module=ColCaiQuantLinear,
83                         kwargs={"split_num": 1},
84                     ),
85                     SubModuleReplacementDescription(
86                         suffix="mlp.up_proj",
87                         target_module=ColCaiQuantLinear,
88                         kwargs={"split_num": 1},
89                     ),
90                     SubModuleReplacementDescription(
91                         suffix="mlp.down_proj",
92                         target_module=RowCaiQuantLinear,
93                         kwargs={"split_num": 1},
94                     ),
95                 ],
96             )
97 
98         elif self.shard_config.extra_kwargs.get("quant", None) == "smoothquant":
99             from colossalai.inference.quant.smoothquant.models.llama import LlamaSmoothquantDecoderLayer
100             from colossalai.inference.quant.smoothquant.models.parallel_linear import (
101                 ColW8A8BFP32OFP32Linear,
102                 RowW8A8B8O8Linear,
103                 RowW8A8BFP32O32LinearSiLU,
104                 RowW8A8BFP32OFP32Linear,
105             )
106 
107             policy[LlamaSmoothquantDecoderLayer] = ModulePolicyDescription(
108                 attribute_replacement=decoder_attribute_replacement,
109                 sub_module_replacement=[
110                     SubModuleReplacementDescription(
111                         suffix="self_attn.q_proj",
112                         target_module=RowW8A8B8O8Linear,
113                         kwargs={"split_num": 1},
114                     ),
115                     SubModuleReplacementDescription(
116                         suffix="self_attn.k_proj",
117                         target_module=RowW8A8B8O8Linear,
118                         kwargs={"split_num": 1},
119                     ),
120                     SubModuleReplacementDescription(
121                         suffix="self_attn.v_proj",
122                         target_module=RowW8A8B8O8Linear,
123                         kwargs={"split_num": 1},
124                     ),
125                     SubModuleReplacementDescription(
126                         suffix="self_attn.o_proj",
127                         target_module=ColW8A8BFP32OFP32Linear,
128                         kwargs={"split_num": 1},
129                     ),
130                     SubModuleReplacementDescription(
131                         suffix="mlp.gate_proj",
132                         target_module=RowW8A8BFP32O32LinearSiLU,
133                         kwargs={"split_num": 1},
134                     ),
135                     SubModuleReplacementDescription(
136                         suffix="mlp.up_proj",
137                         target_module=RowW8A8BFP32OFP32Linear,
138                         kwargs={"split_num": 1},
139                     ),
140                     SubModuleReplacementDescription(
141                         suffix="mlp.down_proj",
142                         target_module=ColW8A8BFP32OFP32Linear,
143                         kwargs={"split_num": 1},
144                     ),
145                 ],
146             )
147         self.shard_config._infer()
148 
149         infer_forward = LlamaInferenceForwards.llama_model_forward
150         method_replacement = {"forward": partial(infer_forward)}
151         self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=LlamaModel)
152 
153         infer_forward = LlamaInferenceForwards.llama_decoder_layer_forward
154         method_replacement = {"forward": partial(infer_forward)}
155         self.append_or_create_method_replacement(
156             description=method_replacement, policy=policy, target_key=LlamaDecoderLayer
157         )
158 
159         infer_forward = LlamaInferenceForwards.llama_flash_attn_kvcache_forward
160         method_replacement = {"forward": partial(infer_forward)}
161         self.append_or_create_method_replacement(
162             description=method_replacement, policy=policy, target_key=LlamaAttention
163         )
164 
165         # set as default, in inference we also use pipeline style forward, just setting stage as 1
166         self.set_pipeline_forward(
167             model_cls=LlamaForCausalLM, new_forward=LlamaInferenceForwards.llama_causal_lm_forward, policy=policy
168         )
169 
170         infer_forward = None
171         if HAS_TRITON_RMSNORM:
172             infer_forward = get_triton_rmsnorm_forward()
173 
174         if infer_forward is not None:
175             method_replacement = {"forward": partial(infer_forward)}
176             self.append_or_create_method_replacement(
177                 description=method_replacement, policy=policy, target_key=LlamaRMSNorm
178             )
179 
180         return policy
181 
182     def postprocess(self):
183         init_to_get_rotary(self.model.model)
184         return self.model
185 
186     def get_held_layers(self) -> List[Module]:
187         """Get pipeline layers for current stage."""
188         assert self.pipeline_stage_manager is not None
189 
190         if self.model.__class__.__name__ == "LlamaModel":
191             module = self.model
192         else:
193             module = self.model.model
194         stage_manager = self.pipeline_stage_manager
195 
196         held_layers = []
197         layers_per_stage = self.distribute_layers(len(module.layers), stage_manager.num_stages)
198         if stage_manager.is_first_stage():
199             held_layers.append(module.embed_tokens)
200             held_layers.append(self.model.lm_head)
201         start_idx, end_idx = self.get_stage_index(layers_per_stage, stage_manager.stage)
202         held_layers.extend(module.layers[start_idx:end_idx])
203         if stage_manager.is_last_stage():
204             held_layers.append(module.norm)
205 
206         return held_layers
207
```
--- 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/engine/policies/llama.py b/colossalai/inference/engine/policies/llama.py
--- a/colossalai/inference/engine/policies/llama.py
+++ b/colossalai/inference/engine/policies/llama.py
@@ -20,23 +20,17 @@
 from ..modeling.llama import LlamaInferenceForwards
 
 try:
-    from colossalai.kernel.triton import rmsnorm_forward
-
+    from lightllm.models.llama.triton_kernel.rmsnorm import rmsnorm_forward as lightllm_rmsnorm_forward
     HAS_TRITON_RMSNORM = True
 except:
     print("you should install triton from https://github.com/openai/triton")
     HAS_TRITON_RMSNORM = False
 
-
-def get_triton_rmsnorm_forward():
-    if HAS_TRITON_RMSNORM:
-
+if HAS_TRITON_RMSNORM:
+    def get_triton_rmsnorm_forward():
         def _triton_rmsnorm_forward(self: LlamaRMSNorm, hidden_states: torch.Tensor):
-            return rmsnorm_forward(hidden_states, self.weight.data, self.variance_epsilon)
-
+            return lightllm_rmsnorm_forward(hidden_states, self.weight.data, self.variance_epsilon)
         return _triton_rmsnorm_forward
-    else:
-        return None
 
 
 class LlamaModelInferPolicy(LlamaForCausalLMPolicy):

{"golden_diff": "diff --git a/colossalai/inference/engine/policies/llama.py b/colossalai/inference/engine/policies/llama.py\n--- a/colossalai/inference/engine/policies/llama.py\n+++ b/colossalai/inference/engine/policies/llama.py\n@@ -20,23 +20,17 @@\n from ..modeling.llama import LlamaInferenceForwards\n \n try:\n-    from colossalai.kernel.triton import rmsnorm_forward\n-\n+    from lightllm.models.llama.triton_kernel.rmsnorm import rmsnorm_forward as lightllm_rmsnorm_forward\n     HAS_TRITON_RMSNORM = True\n except:\n     print(\"you should install triton from https://github.com/openai/triton\")\n     HAS_TRITON_RMSNORM = False\n \n-\n-def get_triton_rmsnorm_forward():\n-    if HAS_TRITON_RMSNORM:\n-\n+if HAS_TRITON_RMSNORM:\n+    def get_triton_rmsnorm_forward():\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 lightllm_rmsnorm_forward(hidden_states, self.weight.data, self.variance_epsilon)\n         return _triton_rmsnorm_forward\n-    else:\n-        return None\n \n \n class LlamaModelInferPolicy(LlamaForCausalLMPolicy):\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\nfrom typing import List\n\nimport torch\nfrom torch.nn import Module\nfrom transformers.models.llama.modeling_llama import (\n    LlamaAttention,\n    LlamaDecoderLayer,\n    LlamaForCausalLM,\n    LlamaModel,\n    LlamaRMSNorm,\n)\n\nfrom colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription\n\n# import colossalai\nfrom colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy\n\nfrom ..modeling._utils import init_to_get_rotary\nfrom ..modeling.llama import LlamaInferenceForwards\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        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            \"self_attn.num_key_value_heads\": self.model.config.num_key_value_heads\n            // self.shard_config.tensor_parallel_size,\n        }\n        if self.shard_config.extra_kwargs.get(\"quant\", None) == \"gptq\":\n            from colossalai.inference.quant.gptq.cai_gptq import ColCaiQuantLinear, RowCaiQuantLinear\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        elif self.shard_config.extra_kwargs.get(\"quant\", None) == \"smoothquant\":\n            from colossalai.inference.quant.smoothquant.models.llama import LlamaSmoothquantDecoderLayer\n            from colossalai.inference.quant.smoothquant.models.parallel_linear import (\n                ColW8A8BFP32OFP32Linear,\n                RowW8A8B8O8Linear,\n                RowW8A8BFP32O32LinearSiLU,\n                RowW8A8BFP32OFP32Linear,\n            )\n\n            policy[LlamaSmoothquantDecoderLayer] = ModulePolicyDescription(\n                attribute_replacement=decoder_attribute_replacement,\n                sub_module_replacement=[\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.q_proj\",\n                        target_module=RowW8A8B8O8Linear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.k_proj\",\n                        target_module=RowW8A8B8O8Linear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.v_proj\",\n                        target_module=RowW8A8B8O8Linear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.o_proj\",\n                        target_module=ColW8A8BFP32OFP32Linear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.gate_proj\",\n                        target_module=RowW8A8BFP32O32LinearSiLU,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.up_proj\",\n                        target_module=RowW8A8BFP32OFP32Linear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.down_proj\",\n                        target_module=ColW8A8BFP32OFP32Linear,\n                        kwargs={\"split_num\": 1},\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        # set as default, in inference we also use pipeline style forward, just setting stage as 1\n        self.set_pipeline_forward(\n            model_cls=LlamaForCausalLM, new_forward=LlamaInferenceForwards.llama_causal_lm_forward, policy=policy\n        )\n\n        infer_forward = None\n        if HAS_TRITON_RMSNORM:\n            infer_forward = get_triton_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\n    def postprocess(self):\n        init_to_get_rotary(self.model.model)\n        return self.model\n\n    def get_held_layers(self) -> List[Module]:\n        \"\"\"Get pipeline layers for current stage.\"\"\"\n        assert self.pipeline_stage_manager is not None\n\n        if self.model.__class__.__name__ == \"LlamaModel\":\n            module = self.model\n        else:\n            module = self.model.model\n        stage_manager = self.pipeline_stage_manager\n\n        held_layers = []\n        layers_per_stage = self.distribute_layers(len(module.layers), stage_manager.num_stages)\n        if stage_manager.is_first_stage():\n            held_layers.append(module.embed_tokens)\n            held_layers.append(self.model.lm_head)\n        start_idx, end_idx = self.get_stage_index(layers_per_stage, stage_manager.stage)\n        held_layers.extend(module.layers[start_idx:end_idx])\n        if stage_manager.is_last_stage():\n            held_layers.append(module.norm)\n\n        return held_layers\n", "path": "colossalai/inference/engine/policies/llama.py"}], "after_files": [{"content": "from functools import partial\nfrom typing import List\n\nimport torch\nfrom torch.nn import Module\nfrom transformers.models.llama.modeling_llama import (\n    LlamaAttention,\n    LlamaDecoderLayer,\n    LlamaForCausalLM,\n    LlamaModel,\n    LlamaRMSNorm,\n)\n\nfrom colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription\n\n# import colossalai\nfrom colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy\n\nfrom ..modeling._utils import init_to_get_rotary\nfrom ..modeling.llama import LlamaInferenceForwards\n\ntry:\n    from lightllm.models.llama.triton_kernel.rmsnorm import rmsnorm_forward as lightllm_rmsnorm_forward\n    HAS_TRITON_RMSNORM = True\nexcept:\n    print(\"you should install triton from https://github.com/openai/triton\")\n    HAS_TRITON_RMSNORM = False\n\nif HAS_TRITON_RMSNORM:\n    def get_triton_rmsnorm_forward():\n        def _triton_rmsnorm_forward(self: LlamaRMSNorm, hidden_states: torch.Tensor):\n            return lightllm_rmsnorm_forward(hidden_states, self.weight.data, self.variance_epsilon)\n        return _triton_rmsnorm_forward\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        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            \"self_attn.num_key_value_heads\": self.model.config.num_key_value_heads\n            // self.shard_config.tensor_parallel_size,\n        }\n        if self.shard_config.extra_kwargs.get(\"quant\", None) == \"gptq\":\n            from colossalai.inference.quant.gptq.cai_gptq import ColCaiQuantLinear, RowCaiQuantLinear\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        elif self.shard_config.extra_kwargs.get(\"quant\", None) == \"smoothquant\":\n            from colossalai.inference.quant.smoothquant.models.llama import LlamaSmoothquantDecoderLayer\n            from colossalai.inference.quant.smoothquant.models.parallel_linear import (\n                ColW8A8BFP32OFP32Linear,\n                RowW8A8B8O8Linear,\n                RowW8A8BFP32O32LinearSiLU,\n                RowW8A8BFP32OFP32Linear,\n            )\n\n            policy[LlamaSmoothquantDecoderLayer] = ModulePolicyDescription(\n                attribute_replacement=decoder_attribute_replacement,\n                sub_module_replacement=[\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.q_proj\",\n                        target_module=RowW8A8B8O8Linear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.k_proj\",\n                        target_module=RowW8A8B8O8Linear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.v_proj\",\n                        target_module=RowW8A8B8O8Linear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"self_attn.o_proj\",\n                        target_module=ColW8A8BFP32OFP32Linear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.gate_proj\",\n                        target_module=RowW8A8BFP32O32LinearSiLU,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.up_proj\",\n                        target_module=RowW8A8BFP32OFP32Linear,\n                        kwargs={\"split_num\": 1},\n                    ),\n                    SubModuleReplacementDescription(\n                        suffix=\"mlp.down_proj\",\n                        target_module=ColW8A8BFP32OFP32Linear,\n                        kwargs={\"split_num\": 1},\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        # set as default, in inference we also use pipeline style forward, just setting stage as 1\n        self.set_pipeline_forward(\n            model_cls=LlamaForCausalLM, new_forward=LlamaInferenceForwards.llama_causal_lm_forward, policy=policy\n        )\n\n        infer_forward = None\n        if HAS_TRITON_RMSNORM:\n            infer_forward = get_triton_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\n    def postprocess(self):\n        init_to_get_rotary(self.model.model)\n        return self.model\n\n    def get_held_layers(self) -> List[Module]:\n        \"\"\"Get pipeline layers for current stage.\"\"\"\n        assert self.pipeline_stage_manager is not None\n\n        if self.model.__class__.__name__ == \"LlamaModel\":\n            module = self.model\n        else:\n            module = self.model.model\n        stage_manager = self.pipeline_stage_manager\n\n        held_layers = []\n        layers_per_stage = self.distribute_layers(len(module.layers), stage_manager.num_stages)\n        if stage_manager.is_first_stage():\n            held_layers.append(module.embed_tokens)\n            held_layers.append(self.model.lm_head)\n        start_idx, end_idx = self.get_stage_index(layers_per_stage, stage_manager.stage)\n        held_layers.extend(module.layers[start_idx:end_idx])\n        if stage_manager.is_last_stage():\n            held_layers.append(module.norm)\n\n        return held_layers\n", "path": "colossalai/inference/engine/policies/llama.py"}]}