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/sharder.py`
Content:
```
1 from types import MethodType
2 from typing import Any, Callable, Dict, List, Optional, Set, Union
3 
4 import torch.nn as nn
5 from torch import Tensor
6 
7 from colossalai.lazy import LazyInitContext
8 
9 from .._utils import getattr_, setattr_
10 from ..policies.auto_policy import get_autopolicy
11 from ..policies.base_policy import Policy, SubModuleReplacementDescription
12 from .shard_config import ShardConfig
13 from .utils import set_tensors_to_none
14 
15 __all__ = ['ModelSharder', 'shard_model']
16 
17 
18 class ModelSharder(object):
19     r"""
20     Shard the original huggingface model according to the policy
21 
22     Args:
23         policy (:class:`Policy`): The policy to shard the model
24         model (:class:`torch.Module`): The model to shard
25         shard_config: The setting of distributed model
26     """
27 
28     def __init__(self, model: nn.Module, policy: Policy, shard_config: ShardConfig = None) -> None:
29         self.model = model
30         self.policy = get_autopolicy(self.model) if policy is None else policy
31         self.shard_config = shard_config
32 
33     def shard(self) -> List[Dict[int, Tensor]]:
34         r"""
35         Shard the model according to the policy
36         """
37         self.policy.set_model(self.model)
38         self.policy.set_shard_config(self.shard_config)
39         self._preprocess()
40         # get shared params before release unheld layers, this avoid misjudgement of shared params (None is None)
41         shared_params = self.policy.get_shared_params()
42         held_layers = self._release_unheld_layers()
43         self._replace_module(include=held_layers)
44         self._materialize()
45         self._postprocess()
46         return shared_params
47 
48     def _preprocess(self) -> None:
49         self.model = self.policy.preprocess()
50 
51     def _postprocess(self) -> None:
52         self.model = self.policy.postprocess()
53 
54     def _replace_module(self, include: Optional[Set[nn.Module]] = None) -> None:
55         r"""
56         Replace the module according to the policy, and replace the module one by one
57 
58         Args:
59             model (:class:`torch.nn.Module`): The model to shard
60         """
61         module_descriptions = self.policy.module_policy()
62         for layer_cls, module_description in module_descriptions.items():
63             attr_replacement = module_description.attribute_replacement
64             param_replacement = module_description.param_replacement
65             sub_module_replacement = module_description.sub_module_replacement
66             method_replacement = module_description.method_replacement
67             self._recursive_replace_layer(self.model,
68                                           layer_cls,
69                                           attr_replacement,
70                                           param_replacement,
71                                           method_replacement,
72                                           sub_module_replacement,
73                                           include=include)
74 
75     def _recursive_replace_layer(
76         self,
77         module: nn.Module,
78         origin_cls: Union[str, nn.Module],
79         attr_replacement: Dict[str, Any],
80         param_replacement: List[Callable],
81         method_replacement: Dict[str, Callable],
82         sub_module_replacement: List[SubModuleReplacementDescription],
83         include: Optional[Set[nn.Module]] = None,
84     ) -> None:
85         r"""
86         Reverse the replace layer operation
87 
88         Args:
89             module (torch.nn.Module): The object of layer to shard
90             origin_cls (Union[str, torch.nn.Module]): The origin layer class or a string of layer class name
91             attr_replacement (Dict[str, Any]): The attribute dict to modify
92             param_replacement (List[Callable]): The function list to get parameter shard information in policy
93             method_replacement (Dict[str, Callable]):  Key is the method name, value is the method for replacement
94             sub_module_replacement ((List[SubModuleReplacementDescription]): The function list to get sub module shard information in policy
95         """
96         # released layers are not shardable
97         can_replace_param_or_layer = include is None or module in include
98         if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or \
99            (module.__class__ == origin_cls):
100             if attr_replacement is not None:
101                 self._replace_attr(module, attr_replacement)
102 
103             if param_replacement is not None and can_replace_param_or_layer:
104                 self._replace_param(module, param_replacement)
105 
106             if method_replacement is not None:
107                 self._replace_method(module, method_replacement)
108 
109             if sub_module_replacement is not None and can_replace_param_or_layer:
110                 self._replace_sub_module(module, sub_module_replacement)
111 
112         for name, child in module.named_children():
113             self._recursive_replace_layer(child,
114                                           origin_cls,
115                                           attr_replacement,
116                                           param_replacement,
117                                           method_replacement,
118                                           sub_module_replacement,
119                                           include=include)
120 
121     def _replace_attr(
122         self,
123         module: nn.Module,
124         attr_replacement: Dict[str, Any],
125     ) -> None:
126         r"""
127         Replace the attribute of the layer
128 
129         Args:
130             module (:class:`torch.nn.Module`): The object of layer to shard
131             attr_replacement (Dict): The attribute dict to modify
132         """
133         for k, v in attr_replacement.items():
134             setattr_(module, k, v, ignore=True)
135 
136     def _replace_param(
137         self,
138         module: nn.Module,
139         param_replacement: List[Callable],
140     ) -> None:
141         r"""
142         Replace the parameter of the layer
143 
144         Args:
145             module (:class:`torch.nn.Module`): The object of layer to shard
146             param_replacement (List[Callable]): The function list to get parameter shard information in policy
147         """
148         for param_func in param_replacement:
149             param_func(module)
150 
151     def _replace_method(self, module: nn.Module, method_replacement: Dict[str, Callable]):
152         for method_name, new_method in method_replacement.items():
153             # bind the new method to the module
154             bound_method = MethodType(new_method, module)
155             setattr(module, method_name, bound_method)
156 
157     def _replace_sub_module(
158         self,
159         org_layer: nn.Module,
160         sub_module_replacement: List[SubModuleReplacementDescription],
161     ) -> None:
162         r"""
163         Shard one layer according to the policy, the layer should be the same class as the key in policy's argument_policy return dict
164 
165         Args:
166             org_layer (torch.nn.Module): The origin layer object to shard
167             sub_module_replacement (List[SubModuleReplacementDescription]): The sub module replacement description list
168 
169         """
170         for description in sub_module_replacement:
171             suffix = description.suffix
172             target_module = description.target_module
173             kwargs = {} if description.kwargs is None else description.kwargs
174 
175             assert target_module is not None, 'target_module should not be None'
176 
177             # TODO: support different parallel mode
178             native_sub_module = getattr_(org_layer, suffix, ignore=True)
179 
180             assert not isinstance(native_sub_module, target_module), \
181                 f"The module with suffix {suffix} has been replaced, please check the policy"
182 
183             # if it is None and we are allowed to ignore this module
184             # just skip
185             if description.ignore_if_not_exist and native_sub_module is None:
186                 continue
187 
188             try:
189                 replace_layer = target_module.from_native_module(native_sub_module,
190                                                                  self.shard_config.tensor_parallel_process_group,
191                                                                  **kwargs)
192             except Exception as e:
193                 raise RuntimeError(
194                     f"Failed to replace {suffix} of type {native_sub_module.__class__.__qualname__}"
195                     f" with {target_module.__qualname__} with the exception: {e}. "
196                     "Please check your model configuration or sharding policy, you can set up an issue for us to help you as well."
197                 )
198 
199             setattr_(org_layer, suffix, replace_layer)
200 
201     def _get_recursive_held_layers(self, held_layers: Optional[List[nn.Module]]) -> Optional[List[nn.Module]]:
202 
203         def collect_sub_modules(module: nn.Module):
204             if module is None:
205                 return
206             recursive_held_layers.append(module)
207             for name, child in module.named_children():
208                 collect_sub_modules(child)
209 
210         recursive_held_layers = []
211         for module in held_layers:
212             collect_sub_modules(module)
213         return recursive_held_layers
214 
215     def _release_unheld_layers(self) -> Optional[Set[nn.Module]]:
216         r"""
217         Release the unheld layers in the model
218         """
219         if self.shard_config and self.shard_config.pipeline_stage_manager:
220             held_layers = self.policy.get_held_layers()
221             set_tensors_to_none(self.model, exclude=set(held_layers))
222             return set(self._get_recursive_held_layers(held_layers))
223         return None
224 
225     def _materialize(self) -> None:
226         r"""
227         Materialize the model if lazy initialization is used
228         """
229         LazyInitContext.materialize(self.model)
230
```
--- 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/shardformer/shard/sharder.py b/colossalai/shardformer/shard/sharder.py
--- a/colossalai/shardformer/shard/sharder.py
+++ b/colossalai/shardformer/shard/sharder.py
@@ -92,22 +92,21 @@
             param_replacement (List[Callable]): The function list to get parameter shard information in policy
             method_replacement (Dict[str, Callable]):  Key is the method name, value is the method for replacement
             sub_module_replacement ((List[SubModuleReplacementDescription]): The function list to get sub module shard information in policy
+            include (Set[nn.Module], optional): The set of modules to keep on current device when pipeline parallel is enabled. Defaults to None
         """
-        # released layers are not shardable
-        can_replace_param_or_layer = include is None or module in include
         if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or \
            (module.__class__ == origin_cls):
             if attr_replacement is not None:
                 self._replace_attr(module, attr_replacement)
 
-            if param_replacement is not None and can_replace_param_or_layer:
+            if param_replacement is not None and (include is None or module in include):
                 self._replace_param(module, param_replacement)
 
             if method_replacement is not None:
                 self._replace_method(module, method_replacement)
 
-            if sub_module_replacement is not None and can_replace_param_or_layer:
-                self._replace_sub_module(module, sub_module_replacement)
+            if sub_module_replacement is not None:
+                self._replace_sub_module(module, sub_module_replacement, include)
 
         for name, child in module.named_children():
             self._recursive_replace_layer(child,
@@ -154,18 +153,17 @@
             bound_method = MethodType(new_method, module)
             setattr(module, method_name, bound_method)
 
-    def _replace_sub_module(
-        self,
-        org_layer: nn.Module,
-        sub_module_replacement: List[SubModuleReplacementDescription],
-    ) -> None:
+    def _replace_sub_module(self,
+                            org_layer: nn.Module,
+                            sub_module_replacement: List[SubModuleReplacementDescription],
+                            include: Optional[Set[nn.Module]] = None) -> None:
         r"""
         Shard one layer according to the policy, the layer should be the same class as the key in policy's argument_policy return dict
 
         Args:
             org_layer (torch.nn.Module): The origin layer object to shard
             sub_module_replacement (List[SubModuleReplacementDescription]): The sub module replacement description list
-
+            include (Set[nn.Module], optional): The set of modules to keep on current device when pipeline parallel is enabled. Defaults to None
         """
         for description in sub_module_replacement:
             suffix = description.suffix
@@ -174,9 +172,12 @@
 
             assert target_module is not None, 'target_module should not be None'
 
-            # TODO: support different parallel mode
             native_sub_module = getattr_(org_layer, suffix, ignore=True)
 
+            # Skip replacement if submodule is not kept by current device when pipeline parallel is enabled.
+            if (include is not None) and (native_sub_module is not None) and (native_sub_module not in include):
+                continue
+
             assert not isinstance(native_sub_module, target_module), \
                 f"The module with suffix {suffix} has been replaced, please check the policy"

{"golden_diff": "diff --git a/colossalai/shardformer/shard/sharder.py b/colossalai/shardformer/shard/sharder.py\n--- a/colossalai/shardformer/shard/sharder.py\n+++ b/colossalai/shardformer/shard/sharder.py\n@@ -92,22 +92,21 @@\n             param_replacement (List[Callable]): The function list to get parameter shard information in policy\n             method_replacement (Dict[str, Callable]):  Key is the method name, value is the method for replacement\n             sub_module_replacement ((List[SubModuleReplacementDescription]): The function list to get sub module shard information in policy\n+            include (Set[nn.Module], optional): The set of modules to keep on current device when pipeline parallel is enabled. Defaults to None\n         \"\"\"\n-        # released layers are not shardable\n-        can_replace_param_or_layer = include is None or module in include\n         if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or \\\n            (module.__class__ == origin_cls):\n             if attr_replacement is not None:\n                 self._replace_attr(module, attr_replacement)\n \n-            if param_replacement is not None and can_replace_param_or_layer:\n+            if param_replacement is not None and (include is None or module in include):\n                 self._replace_param(module, param_replacement)\n \n             if method_replacement is not None:\n                 self._replace_method(module, method_replacement)\n \n-            if sub_module_replacement is not None and can_replace_param_or_layer:\n-                self._replace_sub_module(module, sub_module_replacement)\n+            if sub_module_replacement is not None:\n+                self._replace_sub_module(module, sub_module_replacement, include)\n \n         for name, child in module.named_children():\n             self._recursive_replace_layer(child,\n@@ -154,18 +153,17 @@\n             bound_method = MethodType(new_method, module)\n             setattr(module, method_name, bound_method)\n \n-    def _replace_sub_module(\n-        self,\n-        org_layer: nn.Module,\n-        sub_module_replacement: List[SubModuleReplacementDescription],\n-    ) -> None:\n+    def _replace_sub_module(self,\n+                            org_layer: nn.Module,\n+                            sub_module_replacement: List[SubModuleReplacementDescription],\n+                            include: Optional[Set[nn.Module]] = None) -> None:\n         r\"\"\"\n         Shard one layer according to the policy, the layer should be the same class as the key in policy's argument_policy return dict\n \n         Args:\n             org_layer (torch.nn.Module): The origin layer object to shard\n             sub_module_replacement (List[SubModuleReplacementDescription]): The sub module replacement description list\n-\n+            include (Set[nn.Module], optional): The set of modules to keep on current device when pipeline parallel is enabled. Defaults to None\n         \"\"\"\n         for description in sub_module_replacement:\n             suffix = description.suffix\n@@ -174,9 +172,12 @@\n \n             assert target_module is not None, 'target_module should not be None'\n \n-            # TODO: support different parallel mode\n             native_sub_module = getattr_(org_layer, suffix, ignore=True)\n \n+            # Skip replacement if submodule is not kept by current device when pipeline parallel is enabled.\n+            if (include is not None) and (native_sub_module is not None) and (native_sub_module not in include):\n+                continue\n+\n             assert not isinstance(native_sub_module, target_module), \\\n                 f\"The module with suffix {suffix} has been replaced, please check the policy\"\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "from types import MethodType\nfrom typing import Any, Callable, Dict, List, Optional, Set, Union\n\nimport torch.nn as nn\nfrom torch import Tensor\n\nfrom colossalai.lazy import LazyInitContext\n\nfrom .._utils import getattr_, setattr_\nfrom ..policies.auto_policy import get_autopolicy\nfrom ..policies.base_policy import Policy, SubModuleReplacementDescription\nfrom .shard_config import ShardConfig\nfrom .utils import set_tensors_to_none\n\n__all__ = ['ModelSharder', 'shard_model']\n\n\nclass ModelSharder(object):\n    r\"\"\"\n    Shard the original huggingface model according to the policy\n\n    Args:\n        policy (:class:`Policy`): The policy to shard the model\n        model (:class:`torch.Module`): The model to shard\n        shard_config: The setting of distributed model\n    \"\"\"\n\n    def __init__(self, model: nn.Module, policy: Policy, shard_config: ShardConfig = None) -> None:\n        self.model = model\n        self.policy = get_autopolicy(self.model) if policy is None else policy\n        self.shard_config = shard_config\n\n    def shard(self) -> List[Dict[int, Tensor]]:\n        r\"\"\"\n        Shard the model according to the policy\n        \"\"\"\n        self.policy.set_model(self.model)\n        self.policy.set_shard_config(self.shard_config)\n        self._preprocess()\n        # get shared params before release unheld layers, this avoid misjudgement of shared params (None is None)\n        shared_params = self.policy.get_shared_params()\n        held_layers = self._release_unheld_layers()\n        self._replace_module(include=held_layers)\n        self._materialize()\n        self._postprocess()\n        return shared_params\n\n    def _preprocess(self) -> None:\n        self.model = self.policy.preprocess()\n\n    def _postprocess(self) -> None:\n        self.model = self.policy.postprocess()\n\n    def _replace_module(self, include: Optional[Set[nn.Module]] = None) -> None:\n        r\"\"\"\n        Replace the module according to the policy, and replace the module one by one\n\n        Args:\n            model (:class:`torch.nn.Module`): The model to shard\n        \"\"\"\n        module_descriptions = self.policy.module_policy()\n        for layer_cls, module_description in module_descriptions.items():\n            attr_replacement = module_description.attribute_replacement\n            param_replacement = module_description.param_replacement\n            sub_module_replacement = module_description.sub_module_replacement\n            method_replacement = module_description.method_replacement\n            self._recursive_replace_layer(self.model,\n                                          layer_cls,\n                                          attr_replacement,\n                                          param_replacement,\n                                          method_replacement,\n                                          sub_module_replacement,\n                                          include=include)\n\n    def _recursive_replace_layer(\n        self,\n        module: nn.Module,\n        origin_cls: Union[str, nn.Module],\n        attr_replacement: Dict[str, Any],\n        param_replacement: List[Callable],\n        method_replacement: Dict[str, Callable],\n        sub_module_replacement: List[SubModuleReplacementDescription],\n        include: Optional[Set[nn.Module]] = None,\n    ) -> None:\n        r\"\"\"\n        Reverse the replace layer operation\n\n        Args:\n            module (torch.nn.Module): The object of layer to shard\n            origin_cls (Union[str, torch.nn.Module]): The origin layer class or a string of layer class name\n            attr_replacement (Dict[str, Any]): The attribute dict to modify\n            param_replacement (List[Callable]): The function list to get parameter shard information in policy\n            method_replacement (Dict[str, Callable]):  Key is the method name, value is the method for replacement\n            sub_module_replacement ((List[SubModuleReplacementDescription]): The function list to get sub module shard information in policy\n        \"\"\"\n        # released layers are not shardable\n        can_replace_param_or_layer = include is None or module in include\n        if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or \\\n           (module.__class__ == origin_cls):\n            if attr_replacement is not None:\n                self._replace_attr(module, attr_replacement)\n\n            if param_replacement is not None and can_replace_param_or_layer:\n                self._replace_param(module, param_replacement)\n\n            if method_replacement is not None:\n                self._replace_method(module, method_replacement)\n\n            if sub_module_replacement is not None and can_replace_param_or_layer:\n                self._replace_sub_module(module, sub_module_replacement)\n\n        for name, child in module.named_children():\n            self._recursive_replace_layer(child,\n                                          origin_cls,\n                                          attr_replacement,\n                                          param_replacement,\n                                          method_replacement,\n                                          sub_module_replacement,\n                                          include=include)\n\n    def _replace_attr(\n        self,\n        module: nn.Module,\n        attr_replacement: Dict[str, Any],\n    ) -> None:\n        r\"\"\"\n        Replace the attribute of the layer\n\n        Args:\n            module (:class:`torch.nn.Module`): The object of layer to shard\n            attr_replacement (Dict): The attribute dict to modify\n        \"\"\"\n        for k, v in attr_replacement.items():\n            setattr_(module, k, v, ignore=True)\n\n    def _replace_param(\n        self,\n        module: nn.Module,\n        param_replacement: List[Callable],\n    ) -> None:\n        r\"\"\"\n        Replace the parameter of the layer\n\n        Args:\n            module (:class:`torch.nn.Module`): The object of layer to shard\n            param_replacement (List[Callable]): The function list to get parameter shard information in policy\n        \"\"\"\n        for param_func in param_replacement:\n            param_func(module)\n\n    def _replace_method(self, module: nn.Module, method_replacement: Dict[str, Callable]):\n        for method_name, new_method in method_replacement.items():\n            # bind the new method to the module\n            bound_method = MethodType(new_method, module)\n            setattr(module, method_name, bound_method)\n\n    def _replace_sub_module(\n        self,\n        org_layer: nn.Module,\n        sub_module_replacement: List[SubModuleReplacementDescription],\n    ) -> None:\n        r\"\"\"\n        Shard one layer according to the policy, the layer should be the same class as the key in policy's argument_policy return dict\n\n        Args:\n            org_layer (torch.nn.Module): The origin layer object to shard\n            sub_module_replacement (List[SubModuleReplacementDescription]): The sub module replacement description list\n\n        \"\"\"\n        for description in sub_module_replacement:\n            suffix = description.suffix\n            target_module = description.target_module\n            kwargs = {} if description.kwargs is None else description.kwargs\n\n            assert target_module is not None, 'target_module should not be None'\n\n            # TODO: support different parallel mode\n            native_sub_module = getattr_(org_layer, suffix, ignore=True)\n\n            assert not isinstance(native_sub_module, target_module), \\\n                f\"The module with suffix {suffix} has been replaced, please check the policy\"\n\n            # if it is None and we are allowed to ignore this module\n            # just skip\n            if description.ignore_if_not_exist and native_sub_module is None:\n                continue\n\n            try:\n                replace_layer = target_module.from_native_module(native_sub_module,\n                                                                 self.shard_config.tensor_parallel_process_group,\n                                                                 **kwargs)\n            except Exception as e:\n                raise RuntimeError(\n                    f\"Failed to replace {suffix} of type {native_sub_module.__class__.__qualname__}\"\n                    f\" with {target_module.__qualname__} with the exception: {e}. \"\n                    \"Please check your model configuration or sharding policy, you can set up an issue for us to help you as well.\"\n                )\n\n            setattr_(org_layer, suffix, replace_layer)\n\n    def _get_recursive_held_layers(self, held_layers: Optional[List[nn.Module]]) -> Optional[List[nn.Module]]:\n\n        def collect_sub_modules(module: nn.Module):\n            if module is None:\n                return\n            recursive_held_layers.append(module)\n            for name, child in module.named_children():\n                collect_sub_modules(child)\n\n        recursive_held_layers = []\n        for module in held_layers:\n            collect_sub_modules(module)\n        return recursive_held_layers\n\n    def _release_unheld_layers(self) -> Optional[Set[nn.Module]]:\n        r\"\"\"\n        Release the unheld layers in the model\n        \"\"\"\n        if self.shard_config and self.shard_config.pipeline_stage_manager:\n            held_layers = self.policy.get_held_layers()\n            set_tensors_to_none(self.model, exclude=set(held_layers))\n            return set(self._get_recursive_held_layers(held_layers))\n        return None\n\n    def _materialize(self) -> None:\n        r\"\"\"\n        Materialize the model if lazy initialization is used\n        \"\"\"\n        LazyInitContext.materialize(self.model)\n", "path": "colossalai/shardformer/shard/sharder.py"}], "after_files": [{"content": "from types import MethodType\nfrom typing import Any, Callable, Dict, List, Optional, Set, Union\n\nimport torch.nn as nn\nfrom torch import Tensor\n\nfrom colossalai.lazy import LazyInitContext\n\nfrom .._utils import getattr_, setattr_\nfrom ..policies.auto_policy import get_autopolicy\nfrom ..policies.base_policy import Policy, SubModuleReplacementDescription\nfrom .shard_config import ShardConfig\nfrom .utils import set_tensors_to_none\n\n__all__ = ['ModelSharder', 'shard_model']\n\n\nclass ModelSharder(object):\n    r\"\"\"\n    Shard the original huggingface model according to the policy\n\n    Args:\n        policy (:class:`Policy`): The policy to shard the model\n        model (:class:`torch.Module`): The model to shard\n        shard_config: The setting of distributed model\n    \"\"\"\n\n    def __init__(self, model: nn.Module, policy: Policy, shard_config: ShardConfig = None) -> None:\n        self.model = model\n        self.policy = get_autopolicy(self.model) if policy is None else policy\n        self.shard_config = shard_config\n\n    def shard(self) -> List[Dict[int, Tensor]]:\n        r\"\"\"\n        Shard the model according to the policy\n        \"\"\"\n        self.policy.set_model(self.model)\n        self.policy.set_shard_config(self.shard_config)\n        self._preprocess()\n        # get shared params before release unheld layers, this avoid misjudgement of shared params (None is None)\n        shared_params = self.policy.get_shared_params()\n        held_layers = self._release_unheld_layers()\n        self._replace_module(include=held_layers)\n        self._materialize()\n        self._postprocess()\n        return shared_params\n\n    def _preprocess(self) -> None:\n        self.model = self.policy.preprocess()\n\n    def _postprocess(self) -> None:\n        self.model = self.policy.postprocess()\n\n    def _replace_module(self, include: Optional[Set[nn.Module]] = None) -> None:\n        r\"\"\"\n        Replace the module according to the policy, and replace the module one by one\n\n        Args:\n            model (:class:`torch.nn.Module`): The model to shard\n        \"\"\"\n        module_descriptions = self.policy.module_policy()\n        for layer_cls, module_description in module_descriptions.items():\n            attr_replacement = module_description.attribute_replacement\n            param_replacement = module_description.param_replacement\n            sub_module_replacement = module_description.sub_module_replacement\n            method_replacement = module_description.method_replacement\n            self._recursive_replace_layer(self.model,\n                                          layer_cls,\n                                          attr_replacement,\n                                          param_replacement,\n                                          method_replacement,\n                                          sub_module_replacement,\n                                          include=include)\n\n    def _recursive_replace_layer(\n        self,\n        module: nn.Module,\n        origin_cls: Union[str, nn.Module],\n        attr_replacement: Dict[str, Any],\n        param_replacement: List[Callable],\n        method_replacement: Dict[str, Callable],\n        sub_module_replacement: List[SubModuleReplacementDescription],\n        include: Optional[Set[nn.Module]] = None,\n    ) -> None:\n        r\"\"\"\n        Reverse the replace layer operation\n\n        Args:\n            module (torch.nn.Module): The object of layer to shard\n            origin_cls (Union[str, torch.nn.Module]): The origin layer class or a string of layer class name\n            attr_replacement (Dict[str, Any]): The attribute dict to modify\n            param_replacement (List[Callable]): The function list to get parameter shard information in policy\n            method_replacement (Dict[str, Callable]):  Key is the method name, value is the method for replacement\n            sub_module_replacement ((List[SubModuleReplacementDescription]): The function list to get sub module shard information in policy\n            include (Set[nn.Module], optional): The set of modules to keep on current device when pipeline parallel is enabled. Defaults to None\n        \"\"\"\n        if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or \\\n           (module.__class__ == origin_cls):\n            if attr_replacement is not None:\n                self._replace_attr(module, attr_replacement)\n\n            if param_replacement is not None and (include is None or module in include):\n                self._replace_param(module, param_replacement)\n\n            if method_replacement is not None:\n                self._replace_method(module, method_replacement)\n\n            if sub_module_replacement is not None:\n                self._replace_sub_module(module, sub_module_replacement, include)\n\n        for name, child in module.named_children():\n            self._recursive_replace_layer(child,\n                                          origin_cls,\n                                          attr_replacement,\n                                          param_replacement,\n                                          method_replacement,\n                                          sub_module_replacement,\n                                          include=include)\n\n    def _replace_attr(\n        self,\n        module: nn.Module,\n        attr_replacement: Dict[str, Any],\n    ) -> None:\n        r\"\"\"\n        Replace the attribute of the layer\n\n        Args:\n            module (:class:`torch.nn.Module`): The object of layer to shard\n            attr_replacement (Dict): The attribute dict to modify\n        \"\"\"\n        for k, v in attr_replacement.items():\n            setattr_(module, k, v, ignore=True)\n\n    def _replace_param(\n        self,\n        module: nn.Module,\n        param_replacement: List[Callable],\n    ) -> None:\n        r\"\"\"\n        Replace the parameter of the layer\n\n        Args:\n            module (:class:`torch.nn.Module`): The object of layer to shard\n            param_replacement (List[Callable]): The function list to get parameter shard information in policy\n        \"\"\"\n        for param_func in param_replacement:\n            param_func(module)\n\n    def _replace_method(self, module: nn.Module, method_replacement: Dict[str, Callable]):\n        for method_name, new_method in method_replacement.items():\n            # bind the new method to the module\n            bound_method = MethodType(new_method, module)\n            setattr(module, method_name, bound_method)\n\n    def _replace_sub_module(self,\n                            org_layer: nn.Module,\n                            sub_module_replacement: List[SubModuleReplacementDescription],\n                            include: Optional[Set[nn.Module]] = None) -> None:\n        r\"\"\"\n        Shard one layer according to the policy, the layer should be the same class as the key in policy's argument_policy return dict\n\n        Args:\n            org_layer (torch.nn.Module): The origin layer object to shard\n            sub_module_replacement (List[SubModuleReplacementDescription]): The sub module replacement description list\n            include (Set[nn.Module], optional): The set of modules to keep on current device when pipeline parallel is enabled. Defaults to None\n        \"\"\"\n        for description in sub_module_replacement:\n            suffix = description.suffix\n            target_module = description.target_module\n            kwargs = {} if description.kwargs is None else description.kwargs\n\n            assert target_module is not None, 'target_module should not be None'\n\n            native_sub_module = getattr_(org_layer, suffix, ignore=True)\n\n            # Skip replacement if submodule is not kept by current device when pipeline parallel is enabled.\n            if (include is not None) and (native_sub_module is not None) and (native_sub_module not in include):\n                continue\n\n            assert not isinstance(native_sub_module, target_module), \\\n                f\"The module with suffix {suffix} has been replaced, please check the policy\"\n\n            # if it is None and we are allowed to ignore this module\n            # just skip\n            if description.ignore_if_not_exist and native_sub_module is None:\n                continue\n\n            try:\n                replace_layer = target_module.from_native_module(native_sub_module,\n                                                                 self.shard_config.tensor_parallel_process_group,\n                                                                 **kwargs)\n            except Exception as e:\n                raise RuntimeError(\n                    f\"Failed to replace {suffix} of type {native_sub_module.__class__.__qualname__}\"\n                    f\" with {target_module.__qualname__} with the exception: {e}. \"\n                    \"Please check your model configuration or sharding policy, you can set up an issue for us to help you as well.\"\n                )\n\n            setattr_(org_layer, suffix, replace_layer)\n\n    def _get_recursive_held_layers(self, held_layers: Optional[List[nn.Module]]) -> Optional[List[nn.Module]]:\n\n        def collect_sub_modules(module: nn.Module):\n            if module is None:\n                return\n            recursive_held_layers.append(module)\n            for name, child in module.named_children():\n                collect_sub_modules(child)\n\n        recursive_held_layers = []\n        for module in held_layers:\n            collect_sub_modules(module)\n        return recursive_held_layers\n\n    def _release_unheld_layers(self) -> Optional[Set[nn.Module]]:\n        r\"\"\"\n        Release the unheld layers in the model\n        \"\"\"\n        if self.shard_config and self.shard_config.pipeline_stage_manager:\n            held_layers = self.policy.get_held_layers()\n            set_tensors_to_none(self.model, exclude=set(held_layers))\n            return set(self._get_recursive_held_layers(held_layers))\n        return None\n\n    def _materialize(self) -> None:\n        r\"\"\"\n        Materialize the model if lazy initialization is used\n        \"\"\"\n        LazyInitContext.materialize(self.model)\n", "path": "colossalai/shardformer/shard/sharder.py"}]}