[autoparallel] add layernorm handler (#1629)

colossalai/auto_parallel/solver/_utils.py

@@ -94,7 +94,43 @@ def exception_handler(func):
     def wrapper(*args, **kwargs):
         try:
             func(*args, **kwargs)
-        except Exception as e:
+        except AssertionError as e:
             warnings.warn(f'{e}')
 
     return wrapper
+
+
+def enumerate_all_possible_2d_sharding(mesh_dim_0, mesh_dim_1, dim_size):
+    dim_partition_list = []
+    # enumerate all the 2D sharding cases
+    for i in range(dim_size):
+        for j in range(i + 1, dim_size):
+            dim_partition_dict_0 = {i: [mesh_dim_0], j: [mesh_dim_1]}
+            dim_partition_dict_1 = {i: [mesh_dim_1], j: [mesh_dim_0]}
+            dim_partition_list.append(dim_partition_dict_0)
+            dim_partition_list.append(dim_partition_dict_1)
+    for i in range(dim_size):
+        dim_partition_dict_flatten = {i: [mesh_dim_0, mesh_dim_1]}
+        dim_partition_list.append(dim_partition_dict_flatten)
+
+    return dim_partition_list
+
+
+def enumerate_all_possible_1d_sharding(mesh_dim_0, dim_size):
+    dim_partition_list = []
+    # enumerate all the 1D sharding cases
+    for i in range(dim_size):
+        dim_partition_dict_0 = {i: [mesh_dim_0]}
+        dim_partition_list.append(dim_partition_dict_0)
+
+    return dim_partition_list
+
+
+def generate_sharding_size(dim_partition_dict, device_mesh):
+    total_sharding_size = 1
+    for mesh_dim_list in dim_partition_dict.values():
+        mesh_dim_sharding_size = [device_mesh.shape[mesh_dim] for mesh_dim in mesh_dim_list]
+        sharding_size = reduce(operator.mul, mesh_dim_sharding_size)
+        total_sharding_size *= sharding_size
+
+    return total_sharding_size

colossalai/auto_parallel/solver/constants.py

@@ -3,15 +3,27 @@
 
 __all__ = [
     'ELEMENTWISE_MODULE_OP', 'ELEMENTWISE_FUNC_OP', 'RESHAPE_FUNC_OP', 'CONV_MODULE_OP', 'CONV_FUNC_OP',
-    'LINEAR_MODULE_OP', 'LINEAR_FUNC_OP', 'BATCHNORM_MODULE_OP', 'POOL_MODULE_OP', 'NON_PARAM_FUNC_OP', 'BCAST_FUNC_OP'
+    'LINEAR_MODULE_OP', 'LINEAR_FUNC_OP', 'BATCHNORM_MODULE_OP', 'POOL_MODULE_OP', 'NON_PARAM_FUNC_OP', 'BCAST_FUNC_OP',
+    'EMBEDDING_MODULE_OP', 'LAYERNORM_MODULE_OP', 'ELEMENTWISE_METHOD_OP', 'RESHAPE_METHOD_OP'
 ]
 
 ELEMENTWISE_MODULE_OP = [torch.nn.Dropout, torch.nn.ReLU]
 ELEMENTWISE_FUNC_OP = [
     torch.abs, torch.cos, torch.exp, operator.neg, torch.multiply, torch.nn.functional.relu,
     torch.nn.functional.dropout, torch.flatten
 ]
-RESHAPE_FUNC_OP = [torch.flatten, torch.Tensor.view, torch.reshape]
+ELEMENTWISE_METHOD_OP = [
+    torch.Tensor.to,
+    torch.Tensor.type,
+]
+RESHAPE_FUNC_OP = [torch.flatten, torch.reshape]
+RESHAPE_METHOD_OP = [
+    torch.Tensor.view,
+    torch.Tensor.unsqueeze,
+    torch.Tensor.split,
+    torch.Tensor.permute,
+    torch.Tensor.transpose,
+]
 BCAST_FUNC_OP = [
     torch.add, torch.sub, torch.mul, torch.div, torch.floor_divide, torch.true_divide, operator.add, operator.sub,
     operator.mul, operator.floordiv, operator.truediv, torch.matmul
@@ -23,9 +35,11 @@
 CONV_FUNC_OP = [
     torch.conv1d, torch.conv2d, torch.conv3d, torch.conv_transpose1d, torch.conv_transpose2d, torch.conv_transpose3d
 ]
+EMBEDDING_MODULE_OP = [torch.nn.modules.sparse.Embedding]
 LINEAR_MODULE_OP = [torch.nn.Linear]
 LINEAR_FUNC_OP = [torch.nn.functional.linear, torch.matmul, torch.bmm]
 BATCHNORM_MODULE_OP = [torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, torch.nn.BatchNorm3d, torch.nn.SyncBatchNorm]
+LAYERNORM_MODULE_OP = [torch.nn.LayerNorm]
 POOL_MODULE_OP = [torch.nn.MaxPool1d, torch.nn.MaxPool2d, torch.nn.MaxPool3d, torch.nn.AdaptiveAvgPool2d]
 NON_PARAM_FUNC_OP = RESHAPE_FUNC_OP + ELEMENTWISE_FUNC_OP
 

colossalai/auto_parallel/solver/op_handler/bcast_op_handler.py

@@ -8,7 +8,7 @@
 from colossalai.tensor.sharding_spec import ShardingSpec
 from copy import deepcopy
 from typing import Dict, List
-from colossalai.auto_parallel.solver._utils import exception_handler
+from colossalai.auto_parallel.solver._utils import exception_handler, enumerate_all_possible_1d_sharding, enumerate_all_possible_2d_sharding
 
 __all__ = ['BcastOpHandler']
 
@@ -110,45 +110,19 @@ def _convert_partition_dict_to_sharding_spec(self, dim_partition_list):
 
         return sharding_spec_list
 
-    def _enumerate_all_possible_2d_sharding(self, mesh_dim_0, mesh_dim_1, dim_size):
-        dim_partition_list = []
-        # enumerate all the 2D sharding cases
-        for i in range(dim_size):
-            for j in range(i + 1, dim_size):
-                dim_partition_dict_0 = {i: [mesh_dim_0], j: [mesh_dim_1]}
-                dim_partition_dict_1 = {i: [mesh_dim_1], j: [mesh_dim_0]}
-                dim_partition_list.append(dim_partition_dict_0)
-                dim_partition_list.append(dim_partition_dict_1)
-        for i in range(dim_size):
-            dim_partition_dict_flatten = {i: [mesh_dim_0, mesh_dim_1]}
-            dim_partition_list.append(dim_partition_dict_flatten)
-
-        # sharding_spec_list = self._convert_partition_dict_to_sharding_spec(dim_partition_list)
-        return dim_partition_list
-
-    def _enumerate_all_possible_1d_sharding(self, mesh_dim_0, dim_size):
-        dim_partition_list = []
-        # enumerate all the 1D sharding cases
-        for i in range(dim_size):
-            dim_partition_dict_0 = {i: [mesh_dim_0]}
-            dim_partition_list.append(dim_partition_dict_0)
-
-        # sharding_spec_list = self._convert_partition_dict_to_sharding_spec(dim_partition_list)
-        return dim_partition_list
-
     def _enumerate_all_possible_output(self, mesh_dim_0, mesh_dim_1):
         # use mesh_dim_0, mesh_dim_1 instead of constant 0, 1 in here for N-D device mesh scaliablity.
 
         output_dim_partition_list = []
         dim_size = self.output_data.dim()
         # enumerate all the 2D sharding cases
-        sharding_list_2d = self._enumerate_all_possible_2d_sharding(mesh_dim_0, mesh_dim_1, dim_size)
+        sharding_list_2d = enumerate_all_possible_2d_sharding(mesh_dim_0, mesh_dim_1, dim_size)
         output_dim_partition_list.extend(sharding_list_2d)
 
         # enumerate all the 1D sharding cases
-        sharding_list_1d_on_dim_0 = self._enumerate_all_possible_1d_sharding(mesh_dim_0, dim_size)
+        sharding_list_1d_on_dim_0 = enumerate_all_possible_1d_sharding(mesh_dim_0, dim_size)
         output_dim_partition_list.extend(sharding_list_1d_on_dim_0)
-        sharding_list_1d_on_dim_1 = self._enumerate_all_possible_1d_sharding(mesh_dim_1, dim_size)
+        sharding_list_1d_on_dim_1 = enumerate_all_possible_1d_sharding(mesh_dim_1, dim_size)
         output_dim_partition_list.extend(sharding_list_1d_on_dim_1)
 
         # add empty dict for fully replicated case
@@ -545,15 +519,13 @@ def register_strategy(self) -> StrategiesVector:
             dim_size = self.output_data.dim() - 2
 
             # Both device mesh axises are uesd on batch dimensions
-            dim_partition_dicts_2d = self._enumerate_all_possible_2d_sharding(MESH_DIM_LIST[0], MESH_DIM_LIST[1],
-                                                                              dim_size)
+            dim_partition_dicts_2d = enumerate_all_possible_2d_sharding(MESH_DIM_LIST[0], MESH_DIM_LIST[1], dim_size)
             for dim_partition_dict in dim_partition_dicts_2d:
                 self._registry_no_split_strategies_for_matmul(dim_partition_dict)
 
             # Only one device mesh axis is uesd on batch dimensions
             for mesh_dim_index in [0, 1]:
-                dim_partition_dicts_1d = self._enumerate_all_possible_1d_sharding(MESH_DIM_LIST[mesh_dim_index],
-                                                                                  dim_size)
+                dim_partition_dicts_1d = enumerate_all_possible_1d_sharding(MESH_DIM_LIST[mesh_dim_index], dim_size)
                 for dim_partition_dict in dim_partition_dicts_1d:
                     self._registry_no_split_strategies_for_matmul(dim_partition_dict)
                     self._registry_1d_strategies_for_matmul(dim_partition_dict, [MESH_DIM_LIST[mesh_dim_index - 1]])