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/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py`
Content:
```
1 import operator
2 import warnings
3 from copy import deepcopy
4 from functools import reduce
5 from typing import Dict, List
6 
7 import torch
8 from colossalai.auto_parallel.tensor_shard.deprecated._utils import \
9     ignore_sharding_exception
10 from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import (ShardingStrategy, StrategiesVector)
11 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
12 from colossalai.tensor.sharding_spec import ShardingSpec
13 
14 from .operator_handler import OperatorHandler
15 
16 __all__ = ['EmbeddingHandler']
17 
18 
19 class EmbeddingHandler(OperatorHandler):
20     """
21     An OperatorHandler which deals with the sharding strategies of Embedding operators(such as nn.embedding).
22     """
23 
24     def __init__(self, *args, **kwargs):
25         super().__init__(*args, **kwargs)
26         self.input_data = self.predecessor_node[0]._meta_data
27         self.weight = self.module_named_parameters['weight']
28         self.output_data = self.node._meta_data
29 
30     def _generate_compute_cost(self, total_sharding_size):
31         input_shape = self.input_data.shape
32         weight_shape = self.weight.shape
33         input_shape_product = reduce(operator.mul, input_shape, 1)
34         weight_shape_product = reduce(operator.mul, weight_shape, 1)
35         compute_cost = input_shape_product * weight_shape_product * 2 / total_sharding_size
36         return compute_cost
37 
38     def _generate_memory_cost(self, sharding_size_forward, sharding_size_backward_activation, sharding_size_weight):
39         '''
40         Compute the memory cost per device with this specific strategy.
41 
42         Argument:
43             sharding_size_forward(int): The forward activation will be divided
44                 into sharding_size_forward number partions.
45             sharding_size_backward_activation(int): The backward activation will 
46                 be divided into sharding_size_backward_activation number partions.
47             sharding_size_weight(int): The backward weight will be divided
48                 into sharding_size_weight number partions.
49 
50         Return:
51             memory_cost(Tuple[float]): Memory cost per device with this 
52                 specific strategy, the first element of this tuple is forward
53                 memory cost, and the second element of this tuple is backward
54                 memory cost.
55             memory_cost_forward(float): Memory cost of forward activation per 
56                 device with this specific strategy.
57             memory_cost_backward_activation(float): Memory cost of backward activation 
58                 per device with this specific strategy.
59         '''
60         # compute the memory cost of this strategy
61         dtype = self.input_data.dtype
62         numel_output = self.output_data.numel()
63         numel_input = self.input_data.numel()
64         numel_weight = self.weight.numel()
65         size_per_elem_bytes = torch.tensor([], dtype=dtype).element_size()
66 
67         # forward memory_cost
68         memory_cost_forward_activation = numel_output * size_per_elem_bytes / sharding_size_forward
69         memory_cost_forward_weight = numel_weight * size_per_elem_bytes / sharding_size_weight
70         memory_cost_forward = memory_cost_forward_activation + memory_cost_forward_weight
71 
72         # backward memory_cost
73         memory_cost_backward_activation = numel_input * size_per_elem_bytes / sharding_size_backward_activation
74         memory_cost_backward_weight = numel_weight * size_per_elem_bytes / sharding_size_weight
75         memory_cost_backward = memory_cost_backward_activation + memory_cost_backward_weight
76 
77         # memory_cost pair
78         memory_cost = (memory_cost_forward, memory_cost_backward)
79 
80         return memory_cost, memory_cost_forward_activation, memory_cost_backward_activation, memory_cost_backward_weight
81 
82     @ignore_sharding_exception
83     def split_weight_both_dim(self, mesh_dim_0, mesh_dim_1):
84         name = f'RRS{mesh_dim_1} = RR x S{mesh_dim_0}S{mesh_dim_1}'
85 
86         dim_partition_dict_for_input = {}
87         sharding_spec_for_input = self._generate_sharding_spec(self.input_data, dim_partition_dict_for_input)
88 
89         dim_partition_dict_for_weight = {0: [mesh_dim_0], 1: [mesh_dim_1]}
90         sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)
91 
92         dim_partition_dict_for_output = {2: [mesh_dim_1]}
93         sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
94 
95         # generate resharding cost for this strategy
96         resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])
97 
98         # compute the computation cost of this strategy
99         total_sharding_size = self.device_mesh.shape[0] * self.device_mesh.shape[1]
100         compute_cost = self._generate_compute_cost(total_sharding_size)
101 
102         # compute the memory cost of this strategy
103         sharding_size_forward = self.device_mesh.shape[mesh_dim_1]
104         sharding_size_backward_activation = 1
105         sharding_size_weight = self.device_mesh.shape[mesh_dim_0] * self.device_mesh.shape[mesh_dim_1]
106         memory_cost, memory_cost_forward_activation, memory_cost_backward_activation, _ = self._generate_memory_cost(
107             sharding_size_forward, sharding_size_backward_activation, sharding_size_weight)
108 
109         # compute the communication cost of this strategy during forward phase
110         communication_cost_forward = self.device_mesh.all_reduce_cost(memory_cost_forward_activation, mesh_dim_0)
111         # compute the communication cost of this strategy during backward phase
112         communication_cost_backward = self.device_mesh.all_reduce_cost(memory_cost_backward_activation, mesh_dim_1)
113         communication_cost = communication_cost_forward + communication_cost_backward
114         sharding_strategies = ShardingStrategy(name,
115                                                output_sharding_spec=sharding_spec_for_output,
116                                                compute_cost=compute_cost,
117                                                communication_cost=communication_cost,
118                                                memory_cost=memory_cost,
119                                                resharding_costs=resharding_costs,
120                                                input_shardings=(sharding_spec_for_input, sharding_spec_for_weight))
121         self.strategies_vector.append(sharding_strategies)
122 
123     @ignore_sharding_exception
124     def split_input_both_dim(self, mesh_dim_0, mesh_dim_1):
125         name = f'S{mesh_dim_0}S{mesh_dim_1}R = S{mesh_dim_0}S{mesh_dim_1} x RR'
126 
127         dim_partition_dict_for_input = {0: [mesh_dim_0], 1: [mesh_dim_1]}
128         sharding_spec_for_input = self._generate_sharding_spec(self.input_data, dim_partition_dict_for_input)
129 
130         dim_partition_dict_for_weight = {}
131         sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)
132 
133         dim_partition_dict_for_output = {0: [mesh_dim_0], 1: [mesh_dim_1]}
134         sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
135 
136         # generate resharding cost for this strategy
137         resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])
138 
139         # compute the computation cost of this strategy
140         total_sharding_size = self.device_mesh.shape[0] * self.device_mesh.shape[1]
141         compute_cost = self._generate_compute_cost(total_sharding_size)
142 
143         # compute the memory cost of this strategy
144         sharding_size_forward = self.device_mesh.shape[mesh_dim_0] * self.device_mesh.shape[mesh_dim_1]
145         sharding_size_backward_activation = self.device_mesh.shape[mesh_dim_0] * self.device_mesh.shape[mesh_dim_1]
146         sharding_size_weight = 1
147         memory_cost, memory_cost_forward_activation, memory_cost_backward_activation, memory_cost_backward_weight = self._generate_memory_cost(
148             sharding_size_forward, sharding_size_backward_activation, sharding_size_weight)
149 
150         # This strategy do not need to do all_reduce during forward phase
151         communication_cost_forward = 0
152         # compute the communication cost of this strategy during backward phase
153         communication_cost_backward_activation = 0
154         communication_cost_backward_weight = self.device_mesh.flatten_device_mesh.all_reduce_cost(
155             memory_cost_backward_weight, 0)
156         communication_cost_backward = communication_cost_backward_activation + communication_cost_backward_weight
157         communication_cost = communication_cost_forward + communication_cost_backward
158         sharding_strategies = ShardingStrategy(name,
159                                                output_sharding_spec=sharding_spec_for_output,
160                                                compute_cost=compute_cost,
161                                                communication_cost=communication_cost,
162                                                memory_cost=memory_cost,
163                                                resharding_costs=resharding_costs,
164                                                input_shardings=(sharding_spec_for_input, sharding_spec_for_weight))
165         self.strategies_vector.append(sharding_strategies)
166 
167     def register_strategy(self) -> StrategiesVector:
168         '''
169         Generate every possible strategies for a Conv node, and record all strategies into the strategies_vector.
170         '''
171         # RRS = RR x SS
172         self.split_weight_both_dim(0, 1)
173         self.split_weight_both_dim(1, 0)
174 
175         # SSR = SS x RR
176         self.split_input_both_dim(0, 1)
177         self.split_input_both_dim(1, 0)
178 
179         return self.strategies_vector
180
```
--- 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/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py b/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py
--- a/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py
+++ b/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py
@@ -5,9 +5,9 @@
 from typing import Dict, List
 
 import torch
-from colossalai.auto_parallel.tensor_shard.deprecated._utils import \
-    ignore_sharding_exception
-from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import (ShardingStrategy, StrategiesVector)
+
+from colossalai.auto_parallel.tensor_shard.deprecated._utils import ignore_sharding_exception
+from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import ShardingStrategy, StrategiesVector
 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
 from colossalai.tensor.sharding_spec import ShardingSpec
 
@@ -42,19 +42,19 @@
         Argument:
             sharding_size_forward(int): The forward activation will be divided
                 into sharding_size_forward number partions.
-            sharding_size_backward_activation(int): The backward activation will 
+            sharding_size_backward_activation(int): The backward activation will
                 be divided into sharding_size_backward_activation number partions.
             sharding_size_weight(int): The backward weight will be divided
                 into sharding_size_weight number partions.
 
         Return:
-            memory_cost(Tuple[float]): Memory cost per device with this 
+            memory_cost(Tuple[float]): Memory cost per device with this
                 specific strategy, the first element of this tuple is forward
                 memory cost, and the second element of this tuple is backward
                 memory cost.
-            memory_cost_forward(float): Memory cost of forward activation per 
+            memory_cost_forward(float): Memory cost of forward activation per
                 device with this specific strategy.
-            memory_cost_backward_activation(float): Memory cost of backward activation 
+            memory_cost_backward_activation(float): Memory cost of backward activation
                 per device with this specific strategy.
         '''
         # compute the memory cost of this strategy

{"golden_diff": "diff --git a/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py b/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py\n--- a/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py\n+++ b/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py\n@@ -5,9 +5,9 @@\n from typing import Dict, List\n \n import torch\n-from colossalai.auto_parallel.tensor_shard.deprecated._utils import \\\n-    ignore_sharding_exception\n-from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import (ShardingStrategy, StrategiesVector)\n+\n+from colossalai.auto_parallel.tensor_shard.deprecated._utils import ignore_sharding_exception\n+from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import ShardingStrategy, StrategiesVector\n from colossalai.tensor.shape_consistency import ShapeConsistencyManager\n from colossalai.tensor.sharding_spec import ShardingSpec\n \n@@ -42,19 +42,19 @@\n         Argument:\n             sharding_size_forward(int): The forward activation will be divided\n                 into sharding_size_forward number partions.\n-            sharding_size_backward_activation(int): The backward activation will \n+            sharding_size_backward_activation(int): The backward activation will\n                 be divided into sharding_size_backward_activation number partions.\n             sharding_size_weight(int): The backward weight will be divided\n                 into sharding_size_weight number partions.\n \n         Return:\n-            memory_cost(Tuple[float]): Memory cost per device with this \n+            memory_cost(Tuple[float]): Memory cost per device with this\n                 specific strategy, the first element of this tuple is forward\n                 memory cost, and the second element of this tuple is backward\n                 memory cost.\n-            memory_cost_forward(float): Memory cost of forward activation per \n+            memory_cost_forward(float): Memory cost of forward activation per\n                 device with this specific strategy.\n-            memory_cost_backward_activation(float): Memory cost of backward activation \n+            memory_cost_backward_activation(float): Memory cost of backward activation\n                 per device with this specific strategy.\n         '''\n         # compute the memory cost of this strategy\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import operator\nimport warnings\nfrom copy import deepcopy\nfrom functools import reduce\nfrom typing import Dict, List\n\nimport torch\nfrom colossalai.auto_parallel.tensor_shard.deprecated._utils import \\\n    ignore_sharding_exception\nfrom colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import (ShardingStrategy, StrategiesVector)\nfrom colossalai.tensor.shape_consistency import ShapeConsistencyManager\nfrom colossalai.tensor.sharding_spec import ShardingSpec\n\nfrom .operator_handler import OperatorHandler\n\n__all__ = ['EmbeddingHandler']\n\n\nclass EmbeddingHandler(OperatorHandler):\n    \"\"\"\n    An OperatorHandler which deals with the sharding strategies of Embedding operators(such as nn.embedding).\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        super().__init__(*args, **kwargs)\n        self.input_data = self.predecessor_node[0]._meta_data\n        self.weight = self.module_named_parameters['weight']\n        self.output_data = self.node._meta_data\n\n    def _generate_compute_cost(self, total_sharding_size):\n        input_shape = self.input_data.shape\n        weight_shape = self.weight.shape\n        input_shape_product = reduce(operator.mul, input_shape, 1)\n        weight_shape_product = reduce(operator.mul, weight_shape, 1)\n        compute_cost = input_shape_product * weight_shape_product * 2 / total_sharding_size\n        return compute_cost\n\n    def _generate_memory_cost(self, sharding_size_forward, sharding_size_backward_activation, sharding_size_weight):\n        '''\n        Compute the memory cost per device with this specific strategy.\n\n        Argument:\n            sharding_size_forward(int): The forward activation will be divided\n                into sharding_size_forward number partions.\n            sharding_size_backward_activation(int): The backward activation will \n                be divided into sharding_size_backward_activation number partions.\n            sharding_size_weight(int): The backward weight will be divided\n                into sharding_size_weight number partions.\n\n        Return:\n            memory_cost(Tuple[float]): Memory cost per device with this \n                specific strategy, the first element of this tuple is forward\n                memory cost, and the second element of this tuple is backward\n                memory cost.\n            memory_cost_forward(float): Memory cost of forward activation per \n                device with this specific strategy.\n            memory_cost_backward_activation(float): Memory cost of backward activation \n                per device with this specific strategy.\n        '''\n        # compute the memory cost of this strategy\n        dtype = self.input_data.dtype\n        numel_output = self.output_data.numel()\n        numel_input = self.input_data.numel()\n        numel_weight = self.weight.numel()\n        size_per_elem_bytes = torch.tensor([], dtype=dtype).element_size()\n\n        # forward memory_cost\n        memory_cost_forward_activation = numel_output * size_per_elem_bytes / sharding_size_forward\n        memory_cost_forward_weight = numel_weight * size_per_elem_bytes / sharding_size_weight\n        memory_cost_forward = memory_cost_forward_activation + memory_cost_forward_weight\n\n        # backward memory_cost\n        memory_cost_backward_activation = numel_input * size_per_elem_bytes / sharding_size_backward_activation\n        memory_cost_backward_weight = numel_weight * size_per_elem_bytes / sharding_size_weight\n        memory_cost_backward = memory_cost_backward_activation + memory_cost_backward_weight\n\n        # memory_cost pair\n        memory_cost = (memory_cost_forward, memory_cost_backward)\n\n        return memory_cost, memory_cost_forward_activation, memory_cost_backward_activation, memory_cost_backward_weight\n\n    @ignore_sharding_exception\n    def split_weight_both_dim(self, mesh_dim_0, mesh_dim_1):\n        name = f'RRS{mesh_dim_1} = RR x S{mesh_dim_0}S{mesh_dim_1}'\n\n        dim_partition_dict_for_input = {}\n        sharding_spec_for_input = self._generate_sharding_spec(self.input_data, dim_partition_dict_for_input)\n\n        dim_partition_dict_for_weight = {0: [mesh_dim_0], 1: [mesh_dim_1]}\n        sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)\n\n        dim_partition_dict_for_output = {2: [mesh_dim_1]}\n        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)\n\n        # generate resharding cost for this strategy\n        resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])\n\n        # compute the computation cost of this strategy\n        total_sharding_size = self.device_mesh.shape[0] * self.device_mesh.shape[1]\n        compute_cost = self._generate_compute_cost(total_sharding_size)\n\n        # compute the memory cost of this strategy\n        sharding_size_forward = self.device_mesh.shape[mesh_dim_1]\n        sharding_size_backward_activation = 1\n        sharding_size_weight = self.device_mesh.shape[mesh_dim_0] * self.device_mesh.shape[mesh_dim_1]\n        memory_cost, memory_cost_forward_activation, memory_cost_backward_activation, _ = self._generate_memory_cost(\n            sharding_size_forward, sharding_size_backward_activation, sharding_size_weight)\n\n        # compute the communication cost of this strategy during forward phase\n        communication_cost_forward = self.device_mesh.all_reduce_cost(memory_cost_forward_activation, mesh_dim_0)\n        # compute the communication cost of this strategy during backward phase\n        communication_cost_backward = self.device_mesh.all_reduce_cost(memory_cost_backward_activation, mesh_dim_1)\n        communication_cost = communication_cost_forward + communication_cost_backward\n        sharding_strategies = ShardingStrategy(name,\n                                               output_sharding_spec=sharding_spec_for_output,\n                                               compute_cost=compute_cost,\n                                               communication_cost=communication_cost,\n                                               memory_cost=memory_cost,\n                                               resharding_costs=resharding_costs,\n                                               input_shardings=(sharding_spec_for_input, sharding_spec_for_weight))\n        self.strategies_vector.append(sharding_strategies)\n\n    @ignore_sharding_exception\n    def split_input_both_dim(self, mesh_dim_0, mesh_dim_1):\n        name = f'S{mesh_dim_0}S{mesh_dim_1}R = S{mesh_dim_0}S{mesh_dim_1} x RR'\n\n        dim_partition_dict_for_input = {0: [mesh_dim_0], 1: [mesh_dim_1]}\n        sharding_spec_for_input = self._generate_sharding_spec(self.input_data, dim_partition_dict_for_input)\n\n        dim_partition_dict_for_weight = {}\n        sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)\n\n        dim_partition_dict_for_output = {0: [mesh_dim_0], 1: [mesh_dim_1]}\n        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)\n\n        # generate resharding cost for this strategy\n        resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])\n\n        # compute the computation cost of this strategy\n        total_sharding_size = self.device_mesh.shape[0] * self.device_mesh.shape[1]\n        compute_cost = self._generate_compute_cost(total_sharding_size)\n\n        # compute the memory cost of this strategy\n        sharding_size_forward = self.device_mesh.shape[mesh_dim_0] * self.device_mesh.shape[mesh_dim_1]\n        sharding_size_backward_activation = self.device_mesh.shape[mesh_dim_0] * self.device_mesh.shape[mesh_dim_1]\n        sharding_size_weight = 1\n        memory_cost, memory_cost_forward_activation, memory_cost_backward_activation, memory_cost_backward_weight = self._generate_memory_cost(\n            sharding_size_forward, sharding_size_backward_activation, sharding_size_weight)\n\n        # This strategy do not need to do all_reduce during forward phase\n        communication_cost_forward = 0\n        # compute the communication cost of this strategy during backward phase\n        communication_cost_backward_activation = 0\n        communication_cost_backward_weight = self.device_mesh.flatten_device_mesh.all_reduce_cost(\n            memory_cost_backward_weight, 0)\n        communication_cost_backward = communication_cost_backward_activation + communication_cost_backward_weight\n        communication_cost = communication_cost_forward + communication_cost_backward\n        sharding_strategies = ShardingStrategy(name,\n                                               output_sharding_spec=sharding_spec_for_output,\n                                               compute_cost=compute_cost,\n                                               communication_cost=communication_cost,\n                                               memory_cost=memory_cost,\n                                               resharding_costs=resharding_costs,\n                                               input_shardings=(sharding_spec_for_input, sharding_spec_for_weight))\n        self.strategies_vector.append(sharding_strategies)\n\n    def register_strategy(self) -> StrategiesVector:\n        '''\n        Generate every possible strategies for a Conv node, and record all strategies into the strategies_vector.\n        '''\n        # RRS = RR x SS\n        self.split_weight_both_dim(0, 1)\n        self.split_weight_both_dim(1, 0)\n\n        # SSR = SS x RR\n        self.split_input_both_dim(0, 1)\n        self.split_input_both_dim(1, 0)\n\n        return self.strategies_vector\n", "path": "colossalai/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py"}], "after_files": [{"content": "import operator\nimport warnings\nfrom copy import deepcopy\nfrom functools import reduce\nfrom typing import Dict, List\n\nimport torch\n\nfrom colossalai.auto_parallel.tensor_shard.deprecated._utils import ignore_sharding_exception\nfrom colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import ShardingStrategy, StrategiesVector\nfrom colossalai.tensor.shape_consistency import ShapeConsistencyManager\nfrom colossalai.tensor.sharding_spec import ShardingSpec\n\nfrom .operator_handler import OperatorHandler\n\n__all__ = ['EmbeddingHandler']\n\n\nclass EmbeddingHandler(OperatorHandler):\n    \"\"\"\n    An OperatorHandler which deals with the sharding strategies of Embedding operators(such as nn.embedding).\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        super().__init__(*args, **kwargs)\n        self.input_data = self.predecessor_node[0]._meta_data\n        self.weight = self.module_named_parameters['weight']\n        self.output_data = self.node._meta_data\n\n    def _generate_compute_cost(self, total_sharding_size):\n        input_shape = self.input_data.shape\n        weight_shape = self.weight.shape\n        input_shape_product = reduce(operator.mul, input_shape, 1)\n        weight_shape_product = reduce(operator.mul, weight_shape, 1)\n        compute_cost = input_shape_product * weight_shape_product * 2 / total_sharding_size\n        return compute_cost\n\n    def _generate_memory_cost(self, sharding_size_forward, sharding_size_backward_activation, sharding_size_weight):\n        '''\n        Compute the memory cost per device with this specific strategy.\n\n        Argument:\n            sharding_size_forward(int): The forward activation will be divided\n                into sharding_size_forward number partions.\n            sharding_size_backward_activation(int): The backward activation will\n                be divided into sharding_size_backward_activation number partions.\n            sharding_size_weight(int): The backward weight will be divided\n                into sharding_size_weight number partions.\n\n        Return:\n            memory_cost(Tuple[float]): Memory cost per device with this\n                specific strategy, the first element of this tuple is forward\n                memory cost, and the second element of this tuple is backward\n                memory cost.\n            memory_cost_forward(float): Memory cost of forward activation per\n                device with this specific strategy.\n            memory_cost_backward_activation(float): Memory cost of backward activation\n                per device with this specific strategy.\n        '''\n        # compute the memory cost of this strategy\n        dtype = self.input_data.dtype\n        numel_output = self.output_data.numel()\n        numel_input = self.input_data.numel()\n        numel_weight = self.weight.numel()\n        size_per_elem_bytes = torch.tensor([], dtype=dtype).element_size()\n\n        # forward memory_cost\n        memory_cost_forward_activation = numel_output * size_per_elem_bytes / sharding_size_forward\n        memory_cost_forward_weight = numel_weight * size_per_elem_bytes / sharding_size_weight\n        memory_cost_forward = memory_cost_forward_activation + memory_cost_forward_weight\n\n        # backward memory_cost\n        memory_cost_backward_activation = numel_input * size_per_elem_bytes / sharding_size_backward_activation\n        memory_cost_backward_weight = numel_weight * size_per_elem_bytes / sharding_size_weight\n        memory_cost_backward = memory_cost_backward_activation + memory_cost_backward_weight\n\n        # memory_cost pair\n        memory_cost = (memory_cost_forward, memory_cost_backward)\n\n        return memory_cost, memory_cost_forward_activation, memory_cost_backward_activation, memory_cost_backward_weight\n\n    @ignore_sharding_exception\n    def split_weight_both_dim(self, mesh_dim_0, mesh_dim_1):\n        name = f'RRS{mesh_dim_1} = RR x S{mesh_dim_0}S{mesh_dim_1}'\n\n        dim_partition_dict_for_input = {}\n        sharding_spec_for_input = self._generate_sharding_spec(self.input_data, dim_partition_dict_for_input)\n\n        dim_partition_dict_for_weight = {0: [mesh_dim_0], 1: [mesh_dim_1]}\n        sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)\n\n        dim_partition_dict_for_output = {2: [mesh_dim_1]}\n        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)\n\n        # generate resharding cost for this strategy\n        resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])\n\n        # compute the computation cost of this strategy\n        total_sharding_size = self.device_mesh.shape[0] * self.device_mesh.shape[1]\n        compute_cost = self._generate_compute_cost(total_sharding_size)\n\n        # compute the memory cost of this strategy\n        sharding_size_forward = self.device_mesh.shape[mesh_dim_1]\n        sharding_size_backward_activation = 1\n        sharding_size_weight = self.device_mesh.shape[mesh_dim_0] * self.device_mesh.shape[mesh_dim_1]\n        memory_cost, memory_cost_forward_activation, memory_cost_backward_activation, _ = self._generate_memory_cost(\n            sharding_size_forward, sharding_size_backward_activation, sharding_size_weight)\n\n        # compute the communication cost of this strategy during forward phase\n        communication_cost_forward = self.device_mesh.all_reduce_cost(memory_cost_forward_activation, mesh_dim_0)\n        # compute the communication cost of this strategy during backward phase\n        communication_cost_backward = self.device_mesh.all_reduce_cost(memory_cost_backward_activation, mesh_dim_1)\n        communication_cost = communication_cost_forward + communication_cost_backward\n        sharding_strategies = ShardingStrategy(name,\n                                               output_sharding_spec=sharding_spec_for_output,\n                                               compute_cost=compute_cost,\n                                               communication_cost=communication_cost,\n                                               memory_cost=memory_cost,\n                                               resharding_costs=resharding_costs,\n                                               input_shardings=(sharding_spec_for_input, sharding_spec_for_weight))\n        self.strategies_vector.append(sharding_strategies)\n\n    @ignore_sharding_exception\n    def split_input_both_dim(self, mesh_dim_0, mesh_dim_1):\n        name = f'S{mesh_dim_0}S{mesh_dim_1}R = S{mesh_dim_0}S{mesh_dim_1} x RR'\n\n        dim_partition_dict_for_input = {0: [mesh_dim_0], 1: [mesh_dim_1]}\n        sharding_spec_for_input = self._generate_sharding_spec(self.input_data, dim_partition_dict_for_input)\n\n        dim_partition_dict_for_weight = {}\n        sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)\n\n        dim_partition_dict_for_output = {0: [mesh_dim_0], 1: [mesh_dim_1]}\n        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)\n\n        # generate resharding cost for this strategy\n        resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])\n\n        # compute the computation cost of this strategy\n        total_sharding_size = self.device_mesh.shape[0] * self.device_mesh.shape[1]\n        compute_cost = self._generate_compute_cost(total_sharding_size)\n\n        # compute the memory cost of this strategy\n        sharding_size_forward = self.device_mesh.shape[mesh_dim_0] * self.device_mesh.shape[mesh_dim_1]\n        sharding_size_backward_activation = self.device_mesh.shape[mesh_dim_0] * self.device_mesh.shape[mesh_dim_1]\n        sharding_size_weight = 1\n        memory_cost, memory_cost_forward_activation, memory_cost_backward_activation, memory_cost_backward_weight = self._generate_memory_cost(\n            sharding_size_forward, sharding_size_backward_activation, sharding_size_weight)\n\n        # This strategy do not need to do all_reduce during forward phase\n        communication_cost_forward = 0\n        # compute the communication cost of this strategy during backward phase\n        communication_cost_backward_activation = 0\n        communication_cost_backward_weight = self.device_mesh.flatten_device_mesh.all_reduce_cost(\n            memory_cost_backward_weight, 0)\n        communication_cost_backward = communication_cost_backward_activation + communication_cost_backward_weight\n        communication_cost = communication_cost_forward + communication_cost_backward\n        sharding_strategies = ShardingStrategy(name,\n                                               output_sharding_spec=sharding_spec_for_output,\n                                               compute_cost=compute_cost,\n                                               communication_cost=communication_cost,\n                                               memory_cost=memory_cost,\n                                               resharding_costs=resharding_costs,\n                                               input_shardings=(sharding_spec_for_input, sharding_spec_for_weight))\n        self.strategies_vector.append(sharding_strategies)\n\n    def register_strategy(self) -> StrategiesVector:\n        '''\n        Generate every possible strategies for a Conv node, and record all strategies into the strategies_vector.\n        '''\n        # RRS = RR x SS\n        self.split_weight_both_dim(0, 1)\n        self.split_weight_both_dim(1, 0)\n\n        # SSR = SS x RR\n        self.split_input_both_dim(0, 1)\n        self.split_input_both_dim(1, 0)\n\n        return self.strategies_vector\n", "path": "colossalai/auto_parallel/tensor_shard/deprecated/op_handler/embedding_handler.py"}]}