Merge branch 'main' of https://github.com/oahzxl/ColossalAI into chunk

colossalai/auto_parallel/passes/comm_metainfo_pass.py

+from typing import Dict
+
+import torch
+from torch.fx import GraphModule
+from torch.fx.node import Node
+
+from colossalai.auto_parallel.meta_profiler import MetaInfo
+from colossalai.auto_parallel.passes.runtime_apply_pass import runtime_apply, runtime_comm_spec_apply
+from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, TrainCycleItem
+from colossalai.tensor.comm_spec import CommSpec
+from colossalai.tensor.shape_consistency import ShapeConsistencyManager
+from colossalai.tensor.sharding_spec import ShardingSpec
+
+shape_consistency_manager = ShapeConsistencyManager()
+
+
+def _construct_meta_info(node: Node, origin_sharding_spec: ShardingSpec,
+                         target_sharding_spec: ShardingSpec) -> MetaInfo:
+    # get comm_action_sequence and total_cost from shape_consistency_manager
+    _, comm_action_sequence, total_cost = shape_consistency_manager.shape_consistency(
+        origin_sharding_spec, target_sharding_spec)
+
+    meta_info = MetaInfo()
+    # NOTE: the cost in shape_consistency_manager.mem_cost is the count in number of numel
+    # get mem cost for MetaInfo
+    mem_cost = shape_consistency_manager.mem_cost(comm_action_sequence)
+    # extract user that has _meta_data and extract element length
+    input_node = next(n for n in node._input_nodes if hasattr(n, '_meta_data'))
+    element_length = input_node._meta_data.element_size()
+
+    mem_cost.fwd.activation *= element_length
+    mem_cost.fwd.temp *= element_length
+    mem_cost.bwd.activation *= element_length
+    mem_cost.bwd.temp *= element_length
+    mem_cost.total.activation *= element_length
+
+    meta_info.memory_cost = mem_cost
+
+    # get computation cost for MetaInfo
+    meta_info.compute_cost = TrainCycleItem(total_cost['forward'] * element_length,
+                                            total_cost['backward'] * element_length,
+                                            total_cost['total'] * element_length)
+
+    # get tensor shape for MetaInfo
+    origin_sharding_spec: ShardingSpec
+    target_sharding_spec: ShardingSpec
+    input_shape = origin_sharding_spec.get_sharded_shape_per_device()
+    output_shape = target_sharding_spec.get_sharded_shape_per_device()
+
+    meta_info.fwd_in = [torch.rand(input_shape, device='meta')]
+    meta_info.fwd_buffer = []
+    meta_info.fwd_out = [torch.rand(output_shape, device='meta')]
+
+    return meta_info
+
+
+def _runtime_apply_meta_info(node: Node, origin_spec_dict, sharding_spec_dict) -> MetaInfo:
+    """
+    This method is used to construct `MetaInto` for shape consistency node
+    """
+
+    # extract node index and user node index
+    args = node.args
+    node_index, user_node_index = args[3], args[4]
+    origin_sharding_spec, target_sharding_spec = origin_spec_dict[node_index], sharding_spec_dict[node_index][
+        user_node_index]
+
+    return _construct_meta_info(node, origin_sharding_spec, target_sharding_spec)
+
+
+def _runtime_comm_spec_apply_meta_info(node: Node, comm_actions_dict: Dict) -> MetaInfo:
+    # extract node_index and op_data_name
+    node_index, op_data_name = node.args[2], node.args[3]
+
+    comm_action = comm_actions_dict[node_index][op_data_name]
+    if isinstance(comm_action.comm_spec, CommSpec):
+        # this case is for all_reduce, there will be no memory cost
+        meta_info = MetaInfo()
+        meta_info.memory_cost = TrainCycleItem(MemoryCost(), MemoryCost(), MemoryCost)
+        output_node = next(n for n in node.users if hasattr(n, '_meta_data'))
+        element_length = output_node._meta_data.element_size()
+
+        total_cost = comm_action.comm_spec.get_comm_cost()
+        meta_info.compute_cost = TrainCycleItem(total_cost['forward'] * element_length,
+                                                total_cost['backward'] * element_length,
+                                                total_cost['total'] * element_length)
+
+        input_shape = output_shape = comm_action.comm_spec.sharding_spec.get_sharded_shape_per_device()
+        meta_info.fwd_in = [torch.rand(input_shape, device='meta')]
+        meta_info.fwd_buffer = []
+        meta_info.fwd_out = [torch.rand(output_shape, device='meta')]
+    else:
+        # this case will be handled by shape consistency manager
+        origin_sharding_spec, target_sharding_spec = comm_action.comm_spec['src_spec'], comm_action.comm_spec[
+            'tgt_spec']
+        meta_info = _construct_meta_info(node, origin_sharding_spec, target_sharding_spec)
+
+    return meta_info
+
+
+def comm_metainfo_pass(gm: GraphModule, sharding_spec_dict: Dict, origin_spec_dict: Dict,
+                       comm_actions_dict: Dict) -> GraphModule:
+    """
+    The method manages all the metainfo of the communication node (run_time_apply, runtime_comm_spec_apply) in the graph.
+    """
+    for node in gm.graph.nodes:
+        if node.target == runtime_apply:
+            setattr(node, 'best_metainfo', _runtime_apply_meta_info(node, origin_spec_dict, sharding_spec_dict))
+        elif node.target == runtime_comm_spec_apply:
+            setattr(node, 'best_metainfo', _runtime_comm_spec_apply_meta_info(node, comm_actions_dict))
+        else:
+            pass
+    return gm

colossalai/auto_parallel/passes/constants.py

+import torch
+
+OUTPUT_SAVED_OPS = [torch.nn.functional.relu, torch.nn.functional.softmax, torch.flatten]
+
+OUTPUT_SAVED_MOD = [
+    torch.nn.ReLU,
+    torch.nn.Softmax,
+]

colossalai/auto_parallel/passes/meta_info_prop.py

+import uuid
+from dataclasses import asdict
+from typing import List
+
+import torch
+import torch.fx
+from torch.fx import GraphModule
+from torch.fx.node import Node
+
+from colossalai.auto_parallel.meta_profiler import MetaInfo
+from colossalai.auto_parallel.passes.constants import OUTPUT_SAVED_MOD, OUTPUT_SAVED_OPS
+from colossalai.fx._compatibility import compatibility
+from colossalai.fx.profiler import GraphInfo
+
+
+def _normalize_tuple(x):
+    if not isinstance(x, tuple):
+        return (x,)
+    return x
+
+
+@compatibility(is_backward_compatible=False)
+class MetaInfoProp:
+
+    def __init__(self, module: GraphModule) -> None:
+        self.module = module
+        self.func_dict = {
+            'placeholder': self.placeholder_handler,
+            'get_attr': self.get_attr_handler,
+            'output': self.output_handler,
+            'call_function': self.node_handler,
+            'call_module': self.node_handler,
+            'call_method': self.node_handler,
+        }
+
+    def _set_data_ptr(self, x):
+        """
+        Set uuid to tensor
+        """
+        if isinstance(x, torch.Tensor):
+            if not x.data_ptr():
+                data_ptr = uuid.uuid4()
+                x.data_ptr = lambda: data_ptr
+
+    def _is_inplace(self, node: Node):
+        """
+        Check if the node is inplace operation.
+        """
+        if node.op == 'call_module':
+            return node.graph.owning_module.get_submodule(node.target).__class__ in OUTPUT_SAVED_MOD
+        elif node.op == "call_function":
+            return node.target in OUTPUT_SAVED_OPS
+        return False
+
+    def run(self) -> GraphModule:
+        """
+        Run the meta information propagation pass on the module.
+        """
+        for node in self.module.graph.nodes:
+            node: Node
+            self.func_dict[node.op](node)
+
+    @compatibility(is_backward_compatible=False)
+    def placeholder_handler(self, node: Node) -> None:
+        """
+        Handle the placeholder node.
+        """
+        graph_info = GraphInfo()
+        out = _normalize_tuple(getattr(node, '_meta_data', None))
+        graph_info.fwd_out = list(out) if out[0] is not None else []
+        node.meta = {**asdict(graph_info)}
+
+    @compatibility(is_backward_compatible=False)
+    def get_attr_handler(self, node: Node) -> None:
+        """
+        Handle the get_attr node.
+        """
+        graph_info = GraphInfo()
+        node.meta = {**asdict(graph_info)}
+
+    @compatibility(is_backward_compatible=False)
+    def output_handler(self, node: Node) -> None:
+        """
+        Handle the output node.
+        """
+        graph_info = GraphInfo()
+        output_tensors = []
+        for par in node._input_nodes:
+            if par.meta:
+                output_tensors += par.meta["fwd_out"]
+        graph_info.fwd_in = output_tensors
+        node.meta = {**asdict(graph_info)}
+
+    @compatibility(is_backward_compatible=False)
+    def node_handler(self, node: Node) -> None:
+        """
+        Handle other kind of nodes
+        """
+        assert hasattr(node, 'best_metainfo'), f"Cannot find best_metainfo in node {node}, {node.op}"
+        graph_info = GraphInfo()
+        meta_info = node.best_metainfo
+        meta_info: MetaInfo
+
+        # set data_ptr for input_tensor in MetaInfo class
+        input_tensors: List[torch.Tensor] = meta_info.fwd_in
+        buffer_tensors: List[torch.Tensor] = meta_info.fwd_buffer
+        output_tensors: List[torch.Tensor] = meta_info.fwd_out
+
+        if self._is_inplace(node):
+            # inplace operation will not create new tensor, and it only has one parent node
+            # TODO: Verify this observation
+            # set data_ptr for input_tensor, buffer_tensor and output_tensor of current node
+            parent_node = list(node._input_nodes.keys())[0]
+            parent_tensor = parent_node.meta.get("fwd_out")[0]
+            parent_tensor: torch.Tensor
+            for tensor in input_tensors:
+                tensor.data_ptr = parent_tensor.data_ptr
+            for tensor in buffer_tensors:
+                tensor.data_ptr = parent_tensor.data_ptr
+            for tensor in output_tensors:
+                tensor.data_ptr = parent_tensor.data_ptr
+
+        else:
+            for par in node._input_nodes:
+                # set data_ptr for the input_tensor of current node from the output_tensor of its parent node
+                for tensor in par.meta.get("fwd_out", []):
+                    tensor: torch.Tensor
+                    target_input_tensor = next(
+                        (x for x in input_tensors if not x.data_ptr() and x.shape == tensor.shape), None)
+                    if target_input_tensor is not None:
+                        target_input_tensor.data_ptr = tensor.data_ptr
+
+            # set data_ptr for tensor in input_tensor that is not set
+            for tensor in input_tensors:
+                if not tensor.data_ptr():
+                    self._set_data_ptr(tensor)
+
+            # set data_ptr for buffer_tensor
+            for tensor in buffer_tensors:
+                self._set_data_ptr(tensor)
+
+            # set data_ptr for output_tensor
+            for tensor in output_tensors:
+                self._set_data_ptr(tensor)
+
+        # attach them to graph_info
+        graph_info.fwd_in = input_tensors
+        graph_info.fwd_tmp = buffer_tensors
+        graph_info.fwd_out = output_tensors
+
+        # fetch other memory informations
+        memory_cost = meta_info.memory_cost
+        graph_info.fwd_mem_tmp = memory_cost.fwd.temp
+        graph_info.fwd_mem_out = memory_cost.fwd.activation
+        graph_info.bwd_mem_tmp = memory_cost.bwd.temp
+        graph_info.bwd_mem_out = memory_cost.bwd.activation
+
+        # fetch flop information
+        # here we use fwd_time and bwd_time to deal with the case that
+        # communication cost is a float
+        compute_cost = meta_info.compute_cost
+        graph_info.fwd_time = compute_cost.fwd
+        graph_info.bwd_time = compute_cost.bwd
+
+        node.meta = {**asdict(graph_info)}

colossalai/fx/passes/experimental/adding_shape_consistency_pass_v2.py → colossalai/auto_parallel/passes/runtime_apply_pass.py

-import builtins
-import copy
-import operator
-from ast import NodeTransformer
 from copy import deepcopy
-from typing import List
+from typing import Dict, List
 
 import torch
-from torch.fx import symbolic_trace
 from torch.fx.node import Node
 
-from colossalai.auto_parallel.tensor_shard.sharding_strategy import CommAction, CommType, OperationDataType
+from colossalai.auto_parallel.meta_profiler import MetaInfo
+from colossalai.auto_parallel.tensor_shard.sharding_strategy import (
+    CommAction,
+    CommType,
+    OperationData,
+    OperationDataType,
+    TrainCycleItem,
+)
 from colossalai.device.device_mesh import DeviceMesh
-from colossalai.fx.passes.split_module import split_module
-from colossalai.tensor.comm_spec import CollectiveCommPattern, CommSpec, _all_reduce, pattern_to_func_dict
+from colossalai.tensor.comm_spec import CommSpec
 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
-from colossalai.tensor.sharding_spec import ShardingSpec, _DimSpec
+from colossalai.tensor.sharding_spec import ShardingSpec
 
 shape_consistency_manager = ShapeConsistencyManager()
 
 
-def runtime_apply(node, origin_dict, input_dict, node_index, user_node_index):
+def runtime_apply(node: Node, origin_dict: Dict, input_dict: Dict, node_index: int, user_node_index: int):
+    """
+    This method will be invoked during runtime to do the shape consistency, which make sure the activations is converted into
+    the user node expected form.
+    """
     origin_sharding_spec = origin_dict[node_index]
     target_sharding_spec = input_dict[node_index][user_node_index]
     return shape_consistency_manager.apply_for_autoparallel_runtime(node, origin_sharding_spec, target_sharding_spec)
 
 
-def runtime_comm_spec_apply(tensor, comm_actions_dict, node_index, op_data):
+def runtime_apply_for_iterable_object(node: Node, origin_dict: Dict, input_dict: Dict, node_index: int,
+                                      user_node_index: int):
+    """
+    This method will be invoked during runtime to do the shape consistency, which makes sure the activations in type of tuple or list
+    is converted into the user node expected form.
+    """
+    rst = []
+    for index, (origin_sharding_spec,
+                target_sharding_spec) in enumerate(zip(origin_dict[node_index],
+                                                       input_dict[node_index][user_node_index])):
+        rst.append(
+            shape_consistency_manager.apply_for_autoparallel_runtime(node[index], origin_sharding_spec,
+                                                                     target_sharding_spec))
+    rst = type(node)(rst)
+    return rst
+
 
-    comm_action = comm_actions_dict[node_index][op_data]
+def runtime_comm_spec_apply(tensor: torch.Tensor, comm_actions_dict: Dict, node_index: int, op_data_name: str):
+    """
+    This method will be invoked during runtime to apply the comm action following the instruction of comm spec.
+    """
+    comm_action = comm_actions_dict[node_index][op_data_name]
     if isinstance(comm_action.comm_spec, CommSpec):
         rst = comm_action.comm_spec.covert_spec_to_action(tensor)
     else:
@@ -37,94 +61,11 @@ def runtime_comm_spec_apply(tensor, comm_actions_dict, node_index, op_data):
     return rst
 
 
-def solution_annotatation_pass(gm: torch.fx.GraphModule, solution: List[int], device_mesh):
-    mod_graph = gm.graph
-    nodes = tuple(mod_graph.nodes)
-
-    # the dict to get origin sharding spec of node
-    origin_node_sharding_spec_dict = {}
-    for node_index, (node, strategy_index) in enumerate(zip(nodes, solution)):
-        strategies_vector = node.strategies_vector
-        setattr(node, 'best_strategy', strategies_vector[strategy_index])
-        setattr(node, 'sharding_spec', strategies_vector[strategy_index].get_sharding_spec_by_name(str(node)))
-        origin_node_sharding_spec_dict[node_index] = strategies_vector[strategy_index].get_sharding_spec_by_name(
-            str(node))
-
-    # apply the sharding spec of parameters
-    for node in nodes:
-        if node.op == 'call_module':
-            target_module = node.graph.owning_module.get_submodule(node.target)
-            for name, param in target_module.named_parameters():
-                target_sharding_spec = node.best_strategy.get_sharding_spec_by_name(name)
-                if target_sharding_spec.dim_partition_dict != {}:
-                    origin_sharding_spec = ShardingSpec(device_mesh, param.shape, {})
-                    setattr(param, 'sharding_spec', origin_sharding_spec)
-                    param_sharded = torch.nn.Parameter(
-                        shape_consistency_manager.apply_for_autoparallel_runtime(param.data, param.sharding_spec,
-                                                                                 target_sharding_spec).detach().clone())
-                else:
-                    param_sharded = param
-                setattr(target_module, name, param_sharded)
-                comm_actions = node.best_strategy.communication_actions
-                for operation_data, comm_action in comm_actions.items():
-                    comm_spec_to_use = comm_action.comm_spec
-                    if operation_data.type == OperationDataType.PARAM and operation_data.name == name and comm_action.comm_type == CommType.HOOK:
-
-                        def wrapper(param, comm_spec):
-
-                            def hook_fn(grad):
-                                _all_reduce(grad, comm_spec)
-
-                            param.register_hook(hook_fn)
-
-                        wrapper(param_sharded, comm_spec_to_use)
-
-            sharded_buffer_dict = {}
-            for name, buffer in target_module.named_buffers():
-                origin_sharding_spec = ShardingSpec(device_mesh, buffer.shape, {})
-                setattr(buffer, 'sharding_spec', origin_sharding_spec)
-                target_sharding_spec = node.best_strategy.get_sharding_spec_by_name(name)
-                buffer_sharded = shape_consistency_manager.apply(buffer, target_sharding_spec)
-                sharded_buffer_dict[name] = buffer_sharded
-
-            for name, buffer_sharded in sharded_buffer_dict.items():
-                setattr(target_module, name, buffer_sharded.detach().clone())
-
-    # the dict to get input sharding specs of user node
-    sharding_spec_convert_dict = {}
-    for index, node in enumerate(nodes):
-        target_sharding_specs = []
-        for user_node in node.strategies_vector.successor_nodes:
-            target_sharding_spec = user_node.best_strategy.get_sharding_spec_by_name(str(node.name))
-            target_sharding_specs.append(target_sharding_spec)
-        sharding_spec_convert_dict[index] = target_sharding_specs
-
-    # the dict to record comm actions of nodes
-    comm_actions_dict = {}
-    for index, node in enumerate(nodes):
-        comm_action_dict = {}
-        for op_data, comm_action in node.best_strategy.communication_actions.items():
-            comm_action_dict[op_data.name] = comm_action
-        comm_actions_dict[index] = comm_action_dict
-
-    # add above dicts into graph
-    for node in nodes:
-        if node.op != 'placeholder':
-            with mod_graph.inserting_before(node):
-                input_specs_node = mod_graph.create_node('placeholder', target='sharding_spec_convert_dict')
-                origin_specs_node = mod_graph.create_node('placeholder', target='origin_node_sharding_spec_dict')
-                comm_actions_dict_node = mod_graph.create_node('placeholder', target='comm_actions_dict')
-            break
-
-    return sharding_spec_convert_dict, origin_node_sharding_spec_dict, comm_actions_dict
-
-
-def shape_consistency_pass(gm: torch.fx.GraphModule):
-    mod_graph = gm.graph
-    nodes = tuple(mod_graph.nodes)
-    input_dict_node = None
-    origin_dict_node = None
-
+def _preprocess_graph(nodes: List[Node]):
+    """
+    This method is used to extract all the placeholders with sharding information,
+    and mapping the nodes into the index of the origin graph.
+    """
     # mapping the node into the origin graph index
     node_to_index_dict = {}
     index = 0
@@ -142,40 +83,110 @@ def shape_consistency_pass(gm: torch.fx.GraphModule):
             continue
         node_to_index_dict[node] = index
         index += 1
-    assert input_dict_node is not None
 
-    # add shape consistency apply function into graph
+    return input_dict_node, origin_dict_node, comm_actions_dict_node, node_to_index_dict
+
+
+def _shape_consistency_apply(gm: torch.fx.GraphModule):
+    """
+    This pass is used to add the shape consistency node to the origin graph.
+    """
+    mod_graph = gm.graph
+    nodes = tuple(mod_graph.nodes)
+
+    input_dict_node, origin_dict_node, _, node_to_index_dict = _preprocess_graph(nodes)
+
     for node in nodes:
         if not hasattr(node, 'best_strategy') or node.op == 'output':
             continue
 
-        for user_node in node.strategies_vector.successor_nodes:
-            user_node_index = user_node.strategies_vector.predecessor_nodes.index(node)
-            with mod_graph.inserting_before(user_node):
-                shape_consistency_node = mod_graph.create_node('call_function',
-                                                               runtime_apply,
-                                                               args=(node, origin_dict_node, input_dict_node,
-                                                                     node_to_index_dict[node], user_node_index))
+        for user_node_index, user_node in enumerate(node.strategies_vector.successor_nodes):
+            if isinstance(node.sharding_spec, (list, tuple)):
+                assert isinstance(
+                    node.target_sharding_specs,
+                    (list,
+                     tuple)), 'target sharding specs should be tuple or list when node.sharding_spec is tuple or list'
+                total_difference = 0
+                for sharding_spec, target_sharding_spec in zip(node.sharding_spec,
+                                                               node.target_sharding_specs[user_node_index]):
+                    total_difference += sharding_spec.sharding_sequence_difference(target_sharding_spec)
+                if total_difference == 0:
+                    continue
+                with mod_graph.inserting_before(user_node):
+                    shape_consistency_node = mod_graph.create_node('call_function',
+                                                                   runtime_apply_for_iterable_object,
+                                                                   args=(node, origin_dict_node, input_dict_node,
+                                                                         node_to_index_dict[node], user_node_index))
+
+            else:
+                assert isinstance(node.sharding_spec,
+                                  ShardingSpec), 'node.sharding_spec should be type of ShardingSpec, tuple or list.'
+                if node.sharding_spec.sharding_sequence_difference(node.target_sharding_specs[user_node_index]) == 0:
+                    continue
+                with mod_graph.inserting_before(user_node):
+                    shape_consistency_node = mod_graph.create_node('call_function',
+                                                                   runtime_apply,
+                                                                   args=(node, origin_dict_node, input_dict_node,
+                                                                         node_to_index_dict[node], user_node_index))
 
-            origin_index_args = user_node.args.index(node)
             new_args = list(user_node.args)
-            new_args[origin_index_args] = shape_consistency_node
-            user_node.args = new_args
+            new_kwargs = dict(user_node.kwargs)
+            # the origin node may be a positional argument or key word argument of user node
+            if node in new_args:
+                # substitute the origin node with shape_consistency_node
+                origin_index_args = new_args.index(node)
+                new_args[origin_index_args] = shape_consistency_node
+                user_node.args = tuple(new_args)
+            elif str(node) in new_kwargs:
+                # substitute the origin node with shape_consistency_node
+                new_kwargs[str(node)] = shape_consistency_node
+                user_node.kwargs = new_kwargs
+
+    return gm
+
+
+def _comm_spec_apply(gm: torch.fx.GraphModule):
+    """
+    This pass is used to add the comm spec apply node to the origin graph.
+    """
+    mod_graph = gm.graph
+    nodes = tuple(mod_graph.nodes)
+
+    _, _, comm_actions_dict_node, node_to_index_dict = _preprocess_graph(nodes)
+
+    for node in nodes:
+        if not hasattr(node, 'best_strategy') or node.op == 'output':
+            continue
 
         comm_actions = node.best_strategy.communication_actions
         for op_data, comm_action in comm_actions.items():
-            comm_object = node.args[comm_action.arg_index]
-            if op_data.type == OperationDataType.PARAM:
+
+            if comm_action.comm_type == CommType.HOOK:
                 continue
             if comm_action.comm_type == CommType.BEFORE:
+                if op_data.type == OperationDataType.OUTPUT:
+                    comm_object = node
+                elif comm_action.key_for_kwarg is not None:
+                    comm_object = node.kwargs[comm_action.key_for_kwarg]
+                else:
+                    comm_object = node.args[comm_action.arg_index]
                 with mod_graph.inserting_before(node):
                     comm_spec_apply_node = mod_graph.create_node('call_function',
                                                                  runtime_comm_spec_apply,
                                                                  args=(comm_object, comm_actions_dict_node,
                                                                        node_to_index_dict[node], op_data.name))
-                new_args = list(node.args)
-                new_args[comm_action.arg_index] = comm_spec_apply_node
-                node.args = new_args
+                # the origin node may be a positional argument or key word argument of user node
+                if comm_action.key_for_kwarg is not None:
+                    # substitute the origin node with comm_spec_apply_node
+                    new_kwargs = dict(node.kwargs)
+                    new_kwargs[comm_action.key_for_kwarg] = comm_spec_apply_node
+                    node.kwargs = new_kwargs
+                else:
+                    # substitute the origin node with comm_spec_apply_node
+                    new_args = list(node.args)
+                    new_args[comm_action.arg_index] = comm_spec_apply_node
+                    node.args = tuple(new_args)
+
             elif comm_action.comm_type == CommType.AFTER:
                 with mod_graph.inserting_after(node):
                     comm_spec_apply_node = mod_graph.create_node('call_function',
@@ -187,7 +198,24 @@ def shape_consistency_pass(gm: torch.fx.GraphModule):
                     if user == comm_spec_apply_node:
                         continue
                     new_args = list(user.args)
-                    new_args[new_args.index(node)] = comm_spec_apply_node
-                    user.args = tuple(new_args)
-            # TODO: consider other OperationDataType, such as OperationDataType.OUTPUT
+                    new_kwargs = dict(user.kwargs)
+                    # the origin node may be a positional argument or key word argument of user node
+                    if node in new_args:
+                        # substitute the origin node with comm_spec_apply_node
+                        new_args[new_args.index(node)] = comm_spec_apply_node
+                        user.args = tuple(new_args)
+                    elif str(node) in new_kwargs:
+                        # substitute the origin node with comm_spec_apply_node
+                        new_kwargs[str(node)] = comm_spec_apply_node
+                        user.kwargs = new_kwargs
+    return gm
+
+
+def runtime_apply_pass(gm: torch.fx.GraphModule):
+    """
+    The method manages all the passes acting on the distributed training runtime.
+    """
+    gm = _shape_consistency_apply(gm)
+    gm = _comm_spec_apply(gm)
+
     return gm

colossalai/auto_parallel/passes/runtime_preparation_pass.py

+import operator
+from copy import deepcopy
+from typing import Dict, List, Union
+
+import torch
+from torch.fx import symbolic_trace
+from torch.fx.node import Node
+
+from colossalai.auto_parallel.tensor_shard.constants import RESHAPE_FUNC_OP
+from colossalai.auto_parallel.tensor_shard.sharding_strategy import (
+    CommAction,
+    CommType,
+    OperationDataType,
+    ShardingStrategy,
+)
+from colossalai.auto_parallel.tensor_shard.solver.strategies_constructor import StrategiesConstructor
+from colossalai.device.device_mesh import DeviceMesh
+from colossalai.tensor.comm_spec import _all_reduce
+from colossalai.tensor.shape_consistency import ShapeConsistencyManager
+from colossalai.tensor.sharding_spec import ShardingSpec
+
+shape_consistency_manager = ShapeConsistencyManager()
+
+
+def size_processing(size: Union[int, torch.Size],
+                    dim_partition_dict: Dict[int, List[int]],
+                    device_mesh_info: Dict[int, int],
+                    target_dim: int = None,
+                    node_name: str = None):
+    """
+    This method will be invoked during runtime to convert size node value depending on distributed information.
+    """
+    if target_dim is not None:
+        assert isinstance(size, int)
+        if target_dim in dim_partition_dict:
+            total_shard_size = 1
+            for shard_dim in dim_partition_dict[target_dim]:
+                total_shard_size *= device_mesh_info[shard_dim]
+            size = size * total_shard_size
+
+    else:
+        size = list(size)
+        for dim, dim_size in enumerate(size):
+            if dim in dim_partition_dict:
+                total_shard_size = 1
+                for shard_dim in dim_partition_dict[dim]:
+                    total_shard_size *= device_mesh_info[shard_dim]
+                size[dim] = dim_size * total_shard_size
+        size = torch.Size(size)
+
+    return size
+
+
+def _solution_annotatation(gm: torch.fx.GraphModule,
+                           solution: List[int],
+                           strategies_constructor: StrategiesConstructor = None):
+    """
+    This method is used to stick the solution strategy to the nodes and add the information
+    required in runtime into graph as placeholder nodes.
+    """
+    mod_graph = gm.graph
+    # TODO: In future PR, strategies_constructor should be a required argument,
+    # instead of optional argument. This is because we don't need to consider nodes with
+    # no strategy in runtime preparation pass.
+    if strategies_constructor is not None:
+        nodes = [strategies_vector.node for strategies_vector in strategies_constructor.leaf_strategies]
+        no_strategy_nodes = strategies_constructor.no_strategy_nodes
+    else:
+        nodes = tuple(mod_graph.nodes)
+        no_strategy_nodes = []
+
+    # the dict to get origin sharding spec of node
+    origin_node_sharding_spec_dict = {}
+    for node_index, (node, strategy_index) in enumerate(zip(nodes, solution)):
+        strategies_vector = node.strategies_vector
+        # stick the solution strategy to the corresponding node
+        setattr(node, 'best_strategy', strategies_vector[strategy_index])
+        setattr(node, 'sharding_spec', strategies_vector[strategy_index].get_sharding_spec_by_name(str(node)))
+        origin_node_sharding_spec_dict[node_index] = strategies_vector[strategy_index].get_sharding_spec_by_name(
+            str(node))
+
+        # attach the corresponding metainfo if node has the attribute `metainfo_vector`
+        if hasattr(node, 'metainfo_vector'):
+            setattr(node, 'best_metainfo', node.metainfo_vector[strategy_index])
+
+    # the dict to get input sharding specs of user node
+    sharding_spec_convert_dict = {}
+    # the dict to record comm actions of nodes
+    comm_actions_dict = {}
+    for index, node in enumerate(nodes):
+        target_sharding_specs = []
+        for user_node in node.strategies_vector.successor_nodes:
+            if user_node in no_strategy_nodes:
+                target_sharding_spec = node.best_strategy.get_sharding_spec_by_name(str(node.name))
+            else:
+                target_sharding_spec = user_node.best_strategy.get_sharding_spec_by_name(str(node.name))
+            target_sharding_specs.append(target_sharding_spec)
+        sharding_spec_convert_dict[index] = target_sharding_specs
+        setattr(node, 'target_sharding_specs', target_sharding_specs)
+        # the get_attr node strategy is kind of pending strategy, which means we will change it
+        # to the same strategy of the user node.
+        if node.op == 'get_attr':
+            assert len(target_sharding_specs) == 1, f'sharing weight is not supported in current version.'
+            target_node = node.strategies_vector.successor_nodes[0]
+            node_name = str(node)
+            if target_node.op == 'call_function' and target_node.target in RESHAPE_FUNC_OP:
+                node_name = str(target_node)
+                target_node = target_node.strategies_vector.successor_nodes[0]
+            user_strategy = target_node.best_strategy
+            op_data_in_user = user_strategy.get_op_data_by_name(node_name)
+            origin_pending_strategy = node.best_strategy
+            origin_op_data = origin_pending_strategy.get_op_data_by_name(str(node))
+
+            new_communication_actions = {}
+            if op_data_in_user in user_strategy.communication_actions:
+                new_communication_action = user_strategy.communication_actions.pop(op_data_in_user)
+                new_communication_action.arg_index = 0
+                new_communication_actions[origin_op_data] = new_communication_action
+            node.best_strategy.communication_actions = new_communication_actions
+
+        comm_action_dict = {}
+        for op_data, comm_action in node.best_strategy.communication_actions.items():
+            comm_action_dict[op_data.name] = comm_action
+        comm_actions_dict[index] = comm_action_dict
+
+    # add above dicts into graph
+    for node in nodes:
+        if node.op != 'placeholder':
+            with mod_graph.inserting_before(node):
+                input_specs_node = mod_graph.create_node('placeholder', target='sharding_spec_convert_dict')
+                origin_specs_node = mod_graph.create_node('placeholder', target='origin_node_sharding_spec_dict')
+                comm_actions_dict_node = mod_graph.create_node('placeholder', target='comm_actions_dict')
+            break
+    return gm, sharding_spec_convert_dict, origin_node_sharding_spec_dict, comm_actions_dict
+
+
+def _size_value_converting(gm: torch.fx.GraphModule, device_mesh: DeviceMesh):
+    """
+    In the auto parallel system, tensors may get shard on different devices, so the size of tensors
+    need to be converted to the size of original tensor and managed by the users, such as torch.view,
+    torch.reshape, etc. These nodes have enough information like input sharding_spec and
+    output sharding_spec to decide how to convert the size value.
+    """
+    mod_graph = gm.graph
+    nodes = tuple(mod_graph.nodes)
+    node_pairs = {}
+
+    for node in nodes:
+
+        if node.op == 'call_method' and node.target == 'size':
+            # extract useful information from size node
+            # dim_partition_dict will instruct the size value on which
+            # dimension should be enlarged.
+            sharding_spec = node.args[0].sharding_spec
+            dim_partition_dict = sharding_spec.dim_partition_dict
+
+            # there are two usages of torch.Tensor.size:
+            #   tensor.size()
+            #   tensor.size(dim)
+            # if a target_dim is assigned, then the output will be
+            # in type of int, instead of torch.Size
+            target_dim = None
+            if len(node.args) > 1:
+                target_dim = node.args[1]
+                if target_dim < 0:
+                    target_dim += node.args[0]._meta_data.dim()
+
+            # DeviceMesh information instructs the scaling of the size value
+            device_mesh_info = {}
+            for dim, dim_size in enumerate(device_mesh.mesh_shape):
+                device_mesh_info[dim] = dim_size
+
+            with mod_graph.inserting_after(node):
+                size_processing_node = mod_graph.create_node('call_function',
+                                                             size_processing,
+                                                             args=(node, dim_partition_dict, device_mesh_info,
+                                                                   target_dim, node.name))
+                # store original node and processing node pair in node_pairs dictioanry
+                # It will be used to replace the original node with processing node in slice object
+                node_pairs[node] = size_processing_node
+                size_processing_node._meta_data = node._meta_data
+
+            user_list = list(node.users.keys())
+            for user in user_list:
+                if user == size_processing_node:
+                    continue
+                new_args = list(user.args)
+                new_kwargs = dict(user.kwargs)
+                # the origin node may be a positional argument or key word argument of user node
+                if node in new_args:
+                    # substitute the origin node with size_processing_node
+                    new_args[new_args.index(node)] = size_processing_node
+                    user.args = tuple(new_args)
+                elif str(node) in new_kwargs:
+                    # substitute the origin node with size_processing_node
+                    new_kwargs[str(node)] = size_processing_node
+                    user.kwargs = new_kwargs
+
+        if node.op == 'call_function' and node.target == operator.getitem:
+
+            getitem_index = node.args[1]
+            # slice object is quite special in torch.fx graph,
+            # On one side, we treat slice object same as type of int,
+            # so we do not create a node for slice object. On the other side,
+            # slice object could take fx.Node as its argument. And the user
+            # relationship cannot be tracked in fx graph.
+            # Therefore, I record the node_pairs in this pass, and use the it
+            # to replace the original node argument inside the slice object if
+            # it has been processed in above pass.
+
+            # There are three main usages of operator.getitem:
+            #   getitem(input, int)
+            #   getitem(input, slice)
+            #   getitem(input, Tuple[slice])
+            # In this pass, we need process the last two cases because
+            # node arguments may potentially appear in these cases.
+            if isinstance(getitem_index, slice):
+                new_start, new_stop, new_step = getitem_index.start, getitem_index.stop, getitem_index.step
+                if getitem_index.start in node_pairs:
+                    new_start = node_pairs[getitem_index.start]
+                elif getitem_index.stop in node_pairs:
+                    new_stop = node_pairs[getitem_index.stop]
+                elif getitem_index.step in node_pairs:
+                    new_step = node_pairs[getitem_index.step]
+                new_slice_item = slice(new_start, new_stop, new_step)
+                new_args = (node.args[0], new_slice_item)
+                node.args = new_args
+
+            elif isinstance(getitem_index, (tuple, list)):
+                if not isinstance(getitem_index[0], slice):
+                    continue
+                new_slice_items = []
+
+                for slice_item in getitem_index:
+                    if slice_item is None:
+                        new_slice_items.append(None)
+                        continue
+
+                    new_start, new_stop, new_step = slice_item.start, slice_item.stop, slice_item.step
+
+                    if slice_item.start in node_pairs:
+                        new_start = node_pairs[slice_item.start]
+                    elif slice_item.stop in node_pairs:
+                        new_stop = node_pairs[slice_item.stop]
+                    elif slice_item.step in node_pairs:
+                        new_step = node_pairs[slice_item.step]
+                    new_slice_item = slice(new_start, new_stop, new_step)
+                    new_slice_items.append(new_slice_item)
+
+                new_args = (node.args[0], tuple(new_slice_items))
+                node.args = new_args
+
+    return gm
+
+
+def _node_args_converting(gm: torch.fx.GraphModule, device_mesh: DeviceMesh):
+    """
+    This pass will process node args to adapt the distributed tensor layout.
+    """
+    mod_graph = gm.graph
+    nodes = tuple(mod_graph.nodes)
+
+    for node in nodes:
+        # skip the placeholder node added in _solution_annotation pass
+        if not hasattr(node, 'sharding_spec'):
+            continue
+
+        def _process_sharding_spec(sharding_spec):
+            if isinstance(sharding_spec, ShardingSpec):
+                dim_partition_dict = sharding_spec.dim_partition_dict
+                device_mesh = sharding_spec.device_mesh
+                return dim_partition_dict, device_mesh
+            if sharding_spec is None:
+                return None, None
+            assert isinstance(sharding_spec,
+                              (tuple, list)), 'sharding_spec should be type of ShardingSpec, tuple, list or None'
+
+            device_mesh = sharding_spec[0].device_mesh
+            dim_partition_dict = []
+            for element in sharding_spec:
+                dim_partition_dict.append(_process_sharding_spec(element))
+            return dim_partition_dict, sharding_spec
+
+        output_dim_partition_dict, device_mesh = _process_sharding_spec(node.sharding_spec)
+        new_args = []
+
+        if node.op == 'call_method':
+            method = getattr(node.args[0]._meta_data.__class__, node.target)
+            # process the node with (input, *shape) style args
+            if method in (torch.Tensor.view, torch.Tensor.reshape):
+
+                for arg in node.args:
+                    if isinstance(arg, Node):
+                        if isinstance(arg._meta_data, (int, tuple, list)):
+                            new_args.append(arg._meta_data)
+                        else:
+                            new_args.append(arg)
+                    else:
+                        assert isinstance(
+                            arg, (int, tuple, list)), 'The argument in view node should be either type of Node or int.'
+                        new_args.append(arg)
+
+                for dim, shard_dims in output_dim_partition_dict.items():
+                    total_shard_size = 1
+                    for shard_dim in shard_dims:
+                        total_shard_size *= device_mesh.shape[shard_dim]
+                    # There are two ways to use torch.view:
+                    # 1. torch.view(input, *shape)
+                    # 2. torch.view(input, shape)
+                    if isinstance(new_args[1], int):
+                        # we will skip the dim with -1 value
+                        if new_args[dim + 1] == -1:
+                            continue
+                        else:
+                            new_args[dim + 1] //= total_shard_size
+                    else:
+                        new_args[1] = list(new_args[1])
+                        # we will skip the dim with -1 value
+                        if new_args[1][dim] == -1:
+                            continue
+                        else:
+                            new_args[1][dim] //= total_shard_size
+                node.args = tuple(new_args)
+
+        elif node.op == 'call_function':
+            target = node.target
+            # process the node with (input, torch.Size) style args
+            if target in (torch.reshape,):
+                for arg in node.args:
+                    if isinstance(arg, Node):
+                        if isinstance(arg._meta_data, (tuple, list)):
+                            new_args.append(list(arg._meta_data))
+                        else:
+                            new_args.append(arg)
+                    else:
+                        assert isinstance(
+                            arg, (tuple, list)), 'The argument in reshape node should be either type of Node or tuple.'
+                        new_args.append(list(arg))
+
+                for dim, shard_dims in output_dim_partition_dict.items():
+                    # we will skip the dim with -1 value
+                    if new_args[1][dim] == -1:
+                        continue
+                    total_shard_size = 1
+                    for shard_dim in shard_dims:
+                        total_shard_size *= device_mesh.shape[shard_dim]
+                    new_args[1][dim] //= total_shard_size
+                node.args = tuple(new_args)
+
+    return gm
+
+
+def _module_params_sharding(gm: torch.fx.GraphModule, device_mesh: DeviceMesh):
+    """
+    Apply the sharding action to the module parameters and buffers following the
+    instructions of solver solution.
+    """
+    mod_graph = gm.graph
+    nodes = tuple(mod_graph.nodes)
+    # This stream is created for overlaping the communication and computation.
+    reduction_stream = torch.cuda.Stream()
+    for node in nodes:
+        if node.op == 'call_module':
+            target_module = node.graph.owning_module.get_submodule(node.target)
+            # TODO: we need to do more actions to take care of the shared parameters.
+            if hasattr(target_module, 'processed') and target_module.processed:
+                continue
+            setattr(target_module, 'processed', True)
+            for name, param in target_module.named_parameters():
+                target_sharding_spec = node.best_strategy.get_sharding_spec_by_name(name)
+                # apply the sharding spec of parameters
+                if target_sharding_spec.dim_partition_dict != {}:
+                    origin_sharding_spec = ShardingSpec(device_mesh, param.shape, {})
+                    setattr(param, 'sharding_spec', origin_sharding_spec)
+                    # TODO: build a ColoParamter class to manager the distributed parameters
+                    # we could use .data here, because all the operations just happen before the real training
+                    # loop, so we don't need to track these operations in the autograd graph.
+                    param.data = shape_consistency_manager.apply_for_autoparallel_runtime(
+                        param.data, param.sharding_spec, target_sharding_spec).detach().clone()
+
+                setattr(target_module, name, param)
+                comm_actions = node.best_strategy.communication_actions
+                for operation_data, comm_action in comm_actions.items():
+                    comm_spec_to_use = comm_action.comm_spec
+                    # register hook to the parameters
+                    if operation_data.type == OperationDataType.PARAM and operation_data.name == name and comm_action.comm_type == CommType.HOOK:
+
+                        def wrapper(param, comm_spec):
+
+                            def hook_fn(grad):
+                                _all_reduce(grad, comm_spec, async_op=False)
+
+                            param.register_hook(hook_fn)
+
+                        wrapper(param, comm_spec_to_use)
+
+            sharded_buffer_dict = {}
+            # apply the sharding spec of buffers
+            for name, buffer in target_module.named_buffers():
+                origin_sharding_spec = ShardingSpec(device_mesh, buffer.shape, {})
+                setattr(buffer, 'sharding_spec', origin_sharding_spec)
+                target_sharding_spec = node.best_strategy.get_sharding_spec_by_name(name)
+                buffer_sharded = shape_consistency_manager.apply(buffer, target_sharding_spec)
+                sharded_buffer_dict[name] = buffer_sharded
+
+            for name, buffer_sharded in sharded_buffer_dict.items():
+                setattr(target_module, name, buffer_sharded.detach().clone())
+
+        if node.op == 'get_attr':
+            root = node.graph.owning_module
+            atoms = node.target.split(".")
+            attr_len = len(atoms)
+            if attr_len == 1:
+                target_module = root
+                target = getattr(root, atoms[0])
+            else:
+                target_module = root
+                for atom in atoms[:-1]:
+                    target_module = getattr(target_module, atom)
+                target = getattr(target_module, atoms[-1])
+
+            target_sharding_spec = node.sharding_spec
+            if target_sharding_spec.dim_partition_dict != {}:
+                origin_sharding_spec = ShardingSpec(device_mesh, target.shape, {})
+                setattr(target, 'sharding_spec', origin_sharding_spec)
+                # TODO: build a ColoParamter class to manager the distributed parameters
+                # we could use .data here, because all the operations just happen before the real training
+                # loop, so we don't need to track these operations in the autograd graph.
+                target.data = shape_consistency_manager.apply_for_autoparallel_runtime(
+                    target.data, target.sharding_spec, target_sharding_spec).detach().clone()
+
+            assert hasattr(target_module, atoms[-1])
+            setattr(target_module, atoms[-1], target)
+
+            comm_actions = node.best_strategy.communication_actions
+            for operation_data, comm_action in comm_actions.items():
+                comm_spec_to_use = comm_action.comm_spec
+                # register hook to the parameters
+                if isinstance(node._meta_data, torch.nn.parameter.Parameter) and comm_action.comm_type == CommType.HOOK:
+
+                    def wrapper(param, comm_spec):
+
+                        def hook_fn(grad):
+                            _all_reduce(grad, comm_spec, async_op=False)
+
+                        param.register_hook(hook_fn)
+
+                    wrapper(target, comm_spec_to_use)
+    return gm
+
+
+def implicit_comm_action_apply(gm: torch.fx.GraphModule):
+    """
+    replace the origin kernel into kernel with implicit communication inside.
+    """
+    pass
+
+
+def runtime_preparation_pass(gm: torch.fx.GraphModule,
+                             solution: List[int],
+                             device_mesh: DeviceMesh,
+                             strategies_constructor: StrategiesConstructor = None):
+    gm, sharding_spec_convert_dict, origin_node_sharding_spec_dict, comm_actions_dict = _solution_annotatation(
+        gm, solution, strategies_constructor)
+    gm = _size_value_converting(gm, device_mesh)
+    gm = _node_args_converting(gm, device_mesh)
+    # TODO: the pass below should be uncommented after the implementation of implicit_comm_action_apply_pass completed.
+    # gm = implicit_comm_action_apply(gm)
+    gm = _module_params_sharding(gm, device_mesh)
+
+    return gm, sharding_spec_convert_dict, origin_node_sharding_spec_dict, comm_actions_dict

colossalai/auto_parallel/tensor_shard/constants.py

-import torch
 import operator
 
+import torch
+
 __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',
@@ -25,7 +26,14 @@ ELEMENTWISE_METHOD_OP = [
     # TODO: contiguous maybe need some extra processes.
     torch.Tensor.contiguous
 ]
-RESHAPE_FUNC_OP = [torch.flatten, torch.reshape]
+RESHAPE_FUNC_OP = [
+    torch.flatten,
+    torch.reshape,
+    torch.transpose,
+    torch.split,
+    torch.permute,
+    operator.getitem,
+]
 RESHAPE_METHOD_OP = [
     torch.Tensor.view,
     torch.Tensor.unsqueeze,
@@ -35,7 +43,7 @@ RESHAPE_METHOD_OP = [
 ]
 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, torch.where, operator.pow, torch.pow, torch.tanh
+    operator.mul, operator.floordiv, operator.truediv, torch.matmul, operator.pow, torch.pow
 ]
 CONV_MODULE_OP = [
     torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d, torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d,

colossalai/auto_parallel/tensor_shard/deprecated/__init__.py

+from .cost_graph import CostGraph
+from .graph_analysis import GraphAnalyser
 from .options import SolverOptions
-from .strategies_constructor import StrategiesConstructor
 from .sharding_strategy import ShardingStrategy, StrategiesVector
-from .cost_graph import CostGraph
 from .solver import Solver
-from .graph_analysis import GraphAnalyser
\ No newline at end of file
+from .strategies_constructor import StrategiesConstructor

colossalai/auto_parallel/tensor_shard/deprecated/_utils.py

@@ -5,10 +5,11 @@ from functools import reduce
 from typing import Dict, List, Optional, Union
 
 import torch
+from torch.fx.node import Node
+
 from colossalai.device.device_mesh import DeviceMesh
 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
 from colossalai.tensor.sharding_spec import ShardingSpec
-from torch.fx.node import Node
 
 from .constants import INFINITY_COST
 
@@ -17,7 +18,7 @@ def generate_sharding_spec(input_: Union[Node, torch.Tensor], device_mesh: Devic
                            dim_partition_dict: Dict[int, List[int]]) -> ShardingSpec:
     """
     Generate the sharding spec of the tensor based on the given dim_partition_dict.
-    
+
 
     Args:
         input_ (Union[Node, torch.Tensor]): the input can be a Node object or a PyTorch tensor. If a node is used, it will look for its meta data associated with this node.
@@ -58,7 +59,7 @@ def generate_resharding_costs(nodes: List[Node],
         nodes (List[Node]): a list of nodes
         sharding_spec_for_input(ShardingSpec): a list of ShardingSpec for the nodes.
         count_backward (Optional[bool]): whether to include the cost of resharding in the backward pass, default is True. False can be used for inference.
-        dtype (Optional[torch.dtype]): the data type for cost calculation, default is None. 
+        dtype (Optional[torch.dtype]): the data type for cost calculation, default is None.
     '''
     # The resharding_cost of weight is counted due to sharing weight cases.
     resharding_costs = {}

colossalai/auto_parallel/tensor_shard/deprecated/op_handler/conv_handler.py

@@ -3,9 +3,9 @@ import warnings
 from functools import reduce
 
 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 .operator_handler import OperatorHandler
 
@@ -71,19 +71,19 @@ class ConvHandler(OperatorHandler):
         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
@@ -541,14 +541,14 @@ class ConvHandler(OperatorHandler):
             # strategies_for_input = [[R, R, R, R], [R, S0, R, R], [R, S1, R, R], [S0, R, R, R], [S0, S1, R, R], [S1, R, R, R], [S1, S0, R, R]]
             strategies_vector_for_input = StrategiesVector(node=nodes[0], in_nodes=[nodes[1], 2], strategies=strategies_for_input)
             setattr(nodes[1], 'strategies_vector', strategies_vector_for_input)
-            
+
             strategies_vector = StrategiesVector(node=nodes[2], in_nodes=[nodes[1], ])
             conv_handler = ConvHandler(input_node=nodes[1], input_index=0, weight=dict(gm.named_modules())[nodes[2].name].weight, output_node=nodes[2],
                                     device_mesh=device_mesh, strategies_vector=strategies_vector, shape_consistency_manager=shape_consistency_manager)
             conv_handler.register_strategy_into_strategies_vector()
             for strategy in conv_handler.strategies_vector:
                 print(f'{strategy.name}: compute_cost is {strategy.compute_cost}, communication_cost is {strategy.communication_cost}, memory_cost is {strategy.memory_cost}, resharding_costs is {strategy.resharding_costs}')
-        
+
         Output:
             S0S1 = S0R x RS1: compute_cost is 8856576, communication_cost is 0, memory_cost is 492032.0, resharding_costs is {mul: [0, 32769.001, 131074.2, 0, 32769.1, 131074.2, 98307.201]}
             S1S0 = S1R x RS0: compute_cost is 8856576, communication_cost is 0, memory_cost is 492032.0, resharding_costs is {mul: [0, 131074.2, 32769.001, 131074.2, 98307.201, 0, 32769.1]}

colossalai/auto_parallel/tensor_shard/deprecated/op_handler/dot_handler.py

@@ -6,9 +6,9 @@ from typing import List
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-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 ..constants import LINEAR_FUNC_OP, LINEAR_MODULE_OP
 from .operator_handler import OperatorHandler
@@ -82,13 +82,13 @@ class MatVecStrategyGenerator(StrategyGenerator):
 
 class MatMulStrategyGenerator(StrategyGenerator):
     """
-    MatMulStrategyGenerator is used to generate the sharding strategies when the second tensor is 
+    MatMulStrategyGenerator is used to generate the sharding strategies when the second tensor is
     a 2D tensor. This is used for nn.Linear, F.linear, torch.matmul and torch.addmm.
 
     A matmul can be formulated as [n, p] x [p, q] = [n, q]
 
     Args:
-        is_linear (bool): whether this generator is used for nn.Linear and F.linear. 
+        is_linear (bool): whether this generator is used for nn.Linear and F.linear.
             This will incur extra transformation of the dim partitioning as the weight is transposed.
     """
 
@@ -255,7 +255,7 @@ class BatchedMatMulStrategyGenerator(StrategyGenerator):
     """
     Generate sharding strategies for the batched matrix multiplication.
 
-    A batched matrix multiplication can be viewed as 
+    A batched matrix multiplication can be viewed as
     [b, i, k] x [b, k, j] -> [b, i, j]
     """
 
@@ -431,7 +431,7 @@ class DotHandler(OperatorHandler):
         sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)
 
         dim_partition_dict_for_output = {0: [mesh_dim_0], 1: [mesh_dim_1]}
-        sharding_spec_for_ouput = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_input)
+        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_input)
 
         # generate resharding cost for this strategy
         resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])
@@ -451,7 +451,7 @@ class DotHandler(OperatorHandler):
 
         # create and register strategy
         sharding_strategies = ShardingStrategy(name,
-                                               output_sharding_spec=sharding_spec_for_ouput,
+                                               output_sharding_spec=sharding_spec_for_output,
                                                compute_cost=compute_cost,
                                                communication_cost=communication_cost,
                                                memory_cost=toatl_memory_cost,
@@ -473,7 +473,7 @@ class DotHandler(OperatorHandler):
         sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)
 
         dim_partition_dict_for_output = {0: [mesh_dim_0]}
-        sharding_spec_for_ouput = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
+        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
 
         # generate resharding cost for this strategy
         resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])
@@ -491,7 +491,7 @@ class DotHandler(OperatorHandler):
         communication_cost_grad_backward = self.device_mesh.all_reduce_cost(weight_memory_cost, mesh_dim_0)
         communication_cost = communication_cost_activation_forward + communication_cost_grad_backward
         sharding_strategies = ShardingStrategy(name,
-                                               output_sharding_spec=sharding_spec_for_ouput,
+                                               output_sharding_spec=sharding_spec_for_output,
                                                compute_cost=compute_cost,
                                                communication_cost=communication_cost,
                                                memory_cost=toatl_memory_cost,
@@ -510,7 +510,7 @@ class DotHandler(OperatorHandler):
         sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)
 
         dim_partition_dict_for_output = {1: [mesh_dim_1]}
-        sharding_spec_for_ouput = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_input)
+        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_input)
 
         # generate resharding cost for this strategy
         resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])
@@ -529,7 +529,7 @@ class DotHandler(OperatorHandler):
         communication_cost = communication_cost_activation_backward + communication_cost_activation_forward
 
         sharding_strategies = ShardingStrategy(name,
-                                               output_sharding_spec=sharding_spec_for_ouput,
+                                               output_sharding_spec=sharding_spec_for_output,
                                                compute_cost=compute_cost,
                                                communication_cost=communication_cost,
                                                memory_cost=toatl_memory_cost,
@@ -548,7 +548,7 @@ class DotHandler(OperatorHandler):
         sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)
 
         dim_partition_dict_for_output = {}
-        sharding_spec_for_ouput = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
+        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
 
         # generate resharding cost for this strategy
         resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])
@@ -564,7 +564,7 @@ class DotHandler(OperatorHandler):
         # compute the communication cost of this strategy
         communication_cost = self.device_mesh.all_reduce_cost(activation_memory_cost, mesh_dim)
         sharding_strategies = ShardingStrategy(name,
-                                               output_sharding_spec=sharding_spec_for_ouput,
+                                               output_sharding_spec=sharding_spec_for_output,
                                                compute_cost=compute_cost,
                                                communication_cost=communication_cost,
                                                memory_cost=toatl_memory_cost,
@@ -583,7 +583,7 @@ class DotHandler(OperatorHandler):
         sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)
 
         dim_partition_dict_for_output = {1: [mesh_dim]}
-        sharding_spec_for_ouput = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
+        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
 
         # generate resharding cost for this strategy
         resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])
@@ -600,7 +600,7 @@ class DotHandler(OperatorHandler):
         communication_cost_activation_backward = self.device_mesh.all_reduce_cost(input_grad_memory_cost, mesh_dim)
         communication_cost = communication_cost_activation_backward
         sharding_strategies = ShardingStrategy(name,
-                                               output_sharding_spec=sharding_spec_for_ouput,
+                                               output_sharding_spec=sharding_spec_for_output,
                                                compute_cost=compute_cost,
                                                communication_cost=communication_cost,
                                                memory_cost=toatl_memory_cost,
@@ -619,7 +619,7 @@ class DotHandler(OperatorHandler):
         sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)
 
         dim_partition_dict_for_output = {0: [mesh_dim_0, mesh_dim_1]}
-        sharding_spec_for_ouput = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
+        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
 
         # generate resharding cost for this strategy
         resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])
@@ -636,7 +636,7 @@ class DotHandler(OperatorHandler):
         communication_cost_weight_backward = self.device_mesh.flatten_device_mesh.all_reduce_cost(weight_memory_cost, 0)
         communication_cost = communication_cost_weight_backward
         sharding_strategies = ShardingStrategy(name,
-                                               output_sharding_spec=sharding_spec_for_ouput,
+                                               output_sharding_spec=sharding_spec_for_output,
                                                compute_cost=compute_cost,
                                                communication_cost=communication_cost,
                                                memory_cost=toatl_memory_cost,
@@ -655,7 +655,7 @@ class DotHandler(OperatorHandler):
         sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)
 
         dim_partition_dict_for_output = {}
-        sharding_spec_for_ouput = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
+        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
 
         # generate resharding cost for this strategy
         resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])
@@ -673,7 +673,7 @@ class DotHandler(OperatorHandler):
             activation_memory_cost, 0)
         communication_cost = communication_cost_forward_activation
         sharding_strategies = ShardingStrategy(name,
-                                               output_sharding_spec=sharding_spec_for_ouput,
+                                               output_sharding_spec=sharding_spec_for_output,
                                                compute_cost=compute_cost,
                                                communication_cost=communication_cost,
                                                memory_cost=toatl_memory_cost,
@@ -692,7 +692,7 @@ class DotHandler(OperatorHandler):
         sharding_spec_for_weight = self._generate_sharding_spec(self.weight, dim_partition_dict_for_weight)
 
         dim_partition_dict_for_output = {1: [mesh_dim_0, mesh_dim_1]}
-        sharding_spec_for_ouput = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
+        sharding_spec_for_output = self._generate_sharding_spec(self.output_data, dim_partition_dict_for_output)
 
         # generate resharding cost for this strategy
         resharding_costs = self._generate_resharding_costs([sharding_spec_for_input, sharding_spec_for_weight])
@@ -709,7 +709,7 @@ class DotHandler(OperatorHandler):
             input_grad_memory_cost, 0)
         communication_cost = communication_cost_activation_backward
         sharding_strategies = ShardingStrategy(name,
-                                               output_sharding_spec=sharding_spec_for_ouput,
+                                               output_sharding_spec=sharding_spec_for_output,
                                                compute_cost=compute_cost,
                                                communication_cost=communication_cost,
                                                memory_cost=toatl_memory_cost,

colossalai/auto_parallel/tensor_shard/deprecated/op_handler/layer_norm_handler.py

@@ -2,10 +2,14 @@ import operator
 from functools import reduce
 
 import torch
-from colossalai.auto_parallel.tensor_shard.deprecated._utils import (enumerate_all_possible_1d_sharding,
-                                                                     enumerate_all_possible_2d_sharding,
-                                                                     generate_sharding_size, ignore_sharding_exception)
-from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import (ShardingStrategy, StrategiesVector)
+
+from colossalai.auto_parallel.tensor_shard.deprecated._utils import (
+    enumerate_all_possible_1d_sharding,
+    enumerate_all_possible_2d_sharding,
+    generate_sharding_size,
+    ignore_sharding_exception,
+)
+from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import ShardingStrategy, StrategiesVector
 
 from .operator_handler import OperatorHandler
 
@@ -63,19 +67,19 @@ class LayerNormHandler(OperatorHandler):
         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
@@ -216,7 +220,7 @@ class LayerNormHandler(OperatorHandler):
             norm_handler.register_strategy()
             for strategy in norm_handler.strategies_vector:
                 print(f'{strategy.name}, computation_cost: {strategy.compute_cost}, memory_cost: {strategy.memory_cost}')
-        
+
         Output:
             RS0 = RS0 x S0, computation_cost: 131072, memory_cost: 524288.0
             RS1 = RS1 x S1, computation_cost: 131072, memory_cost: 524288.0

colossalai/auto_parallel/tensor_shard/deprecated/op_handler/operator_handler.py

+from abc import ABC, abstractmethod
+from typing import Dict, List
 from webbrowser import Opera
+
 import torch
 import torch.nn as nn
-from abc import ABC, abstractmethod
 from torch.fx.node import Node
-from typing import Dict, List
+
+from colossalai.auto_parallel.tensor_shard.deprecated.constants import *
 from colossalai.device.device_mesh import DeviceMesh
 from colossalai.tensor.sharding_spec import ShardingSpec
-from .._utils import generate_resharding_costs, generate_sharding_spec
-from colossalai.auto_parallel.tensor_shard.deprecated.constants import *
 
+from .._utils import generate_resharding_costs, generate_sharding_spec
 from ..sharding_strategy import StrategiesVector
 
 __all__ = ['OperatorHandler']
@@ -60,7 +62,7 @@ class OperatorHandler(ABC):
     @abstractmethod
     def register_strategy(self) -> StrategiesVector:
         """
-        Register 
+        Register
         """
         pass
 

colossalai/auto_parallel/tensor_shard/deprecated/op_handler/reshape_handler.py

@@ -4,9 +4,9 @@ import warnings
 from copy import deepcopy
 
 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
 

colossalai/auto_parallel/tensor_shard/deprecated/strategies_constructor.py

+import builtins
+import math
+import operator
+from copy import deepcopy
+from typing import Dict, List
+
+import torch
 from torch.fx import Graph, Node
-from colossalai.tensor.sharding_spec import ShardingSpec
+
 from colossalai.device.device_mesh import DeviceMesh
 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
+from colossalai.tensor.sharding_spec import ShardingSpec
+
+from ._utils import generate_resharding_costs, generate_sharding_spec
+from .constants import *
+from .op_handler import *
 from .options import SolverOptions
 from .sharding_strategy import ShardingStrategy, StrategiesVector
-from .op_handler import *
-from .constants import *
-from copy import deepcopy
-import math
-import torch
-import operator
-from typing import Dict, List
-from ._utils import generate_sharding_spec, generate_resharding_costs
-import builtins
 
 __all__ = ['StrategiesConstructor']
 

colossalai/auto_parallel/tensor_shard/initialize.py

+from typing import Dict, List, Tuple
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.fx import GraphModule
+from torch.fx.graph import Graph
+
+from colossalai.auto_parallel.passes.runtime_apply_pass import runtime_apply_pass
+from colossalai.auto_parallel.passes.runtime_preparation_pass import runtime_preparation_pass
+from colossalai.auto_parallel.tensor_shard.sharding_strategy import CommAction
+from colossalai.auto_parallel.tensor_shard.solver import (
+    CostGraph,
+    GraphAnalyser,
+    Solver,
+    SolverOptions,
+    StrategiesConstructor,
+)
+from colossalai.device.alpha_beta_profiler import AlphaBetaProfiler
+from colossalai.device.device_mesh import DeviceMesh
+from colossalai.fx.tracer import ColoTracer
+from colossalai.tensor.sharding_spec import ShardingSpec
+
+
+class ModuleWrapper(nn.Module):
+    '''
+    This class is used to wrap the original module, and add the sharding_spec_dict, origin_spec_dict, comm_actions_dict
+    into the forward function.
+    '''
+
+    def __init__(self, module: GraphModule, sharding_spec_dict: Dict[int, List[ShardingSpec]],
+                 origin_spec_dict: Dict[int, ShardingSpec], comm_actions_dict: Dict[int, Dict[str, CommAction]]):
+        '''
+        Args:
+            module: the original module
+            sharding_spec_dict: The sharding_spec_dict is used to record the target sharding specs of each tensor required in user node.
+            origin_spec_dict: The origin_spec_dict is used to record the original sharding spec of each tensor.
+            comm_actions_dict: The comm_actions_dict is used to record the communication actions of each tensor.
+        '''
+        super(ModuleWrapper, self).__init__()
+        self.module = module
+        self.sharding_spec_dict = sharding_spec_dict
+        self.origin_spec_dict = origin_spec_dict
+        self.comm_actions_dict = comm_actions_dict
+
+    def forward(self, *args, **kwargs):
+        return self.module(*args,
+                           sharding_spec_convert_dict=self.sharding_spec_dict,
+                           origin_node_sharding_spec_dict=self.origin_spec_dict,
+                           comm_actions_dict=self.comm_actions_dict,
+                           **kwargs)
+
+
+def extract_meta_args_from_dataloader(data_loader: torch.utils.data.DataLoader, data_process_func: callable):
+    '''
+    This method is used to extract the meta_args from the dataloader under the instruction of the data_process_func.
+    '''
+    # TODO: implement this function
+    pass
+
+
+def search_best_logical_mesh_shape(world_size: int, alpha_beta_dict: Dict[Tuple[int], Tuple[float]]):
+    '''
+    This method is used to search the best logical mesh shape for the given world size
+    based on the alpha_beta_dict.
+
+    For example:
+        if the world_size is 8, and the possible logical shape will be (1, 8), (2, 4), (4, 2), (8, 1).
+    '''
+    # TODO: implement this function
+    return (world_size, 1)
+
+
+def extract_alpha_beta_for_device_mesh(alpha_beta_dict: Dict[Tuple[int], Tuple[float]], logical_mesh_shape: Tuple[int]):
+    '''
+    This method is used to extract the mesh_alpha and mesh_beta for the given logical_mesh_shape
+    from the alpha_beta_dict. These two values will be used to estimate the communication cost.
+    '''
+    # TODO: implement this function
+    pass
+
+
+def build_strategy_constructor(graph: Graph, device_mesh: DeviceMesh):
+    '''
+    This method is used to build the strategy_constructor for the given graph.
+    After this method, each node in the graph will have a strategies_vector which
+    is constructed by the related node handler.
+    '''
+    solver_options = SolverOptions()
+    strategies_constructor = StrategiesConstructor(graph, device_mesh, solver_options)
+    strategies_constructor.build_strategies_and_cost()
+
+    return strategies_constructor
+
+
+def solve_solution(gm: GraphModule, strategy_constructor: StrategiesConstructor, memory_budget: float = -1.0):
+    '''
+    This method is used to solve the best solution for the given graph.
+    The solution is a list of integers, each integer represents the best strategy index of the corresponding node.
+    '''
+    graph_analyser = GraphAnalyser(gm)
+    liveness_list = graph_analyser.liveness_analysis()
+    cost_graph = CostGraph(strategy_constructor.leaf_strategies)
+    cost_graph.simplify_graph()
+    solver = Solver(gm.graph, strategy_constructor, cost_graph, graph_analyser, memory_budget=memory_budget)
+    ret = solver.call_solver_serialized_args()
+    solution = list(ret[0])
+
+    return solution
+
+
+def transform_to_sharded_model(gm: GraphModule, solution: List[int], device_mesh: DeviceMesh,
+                               strategies_constructor: StrategiesConstructor):
+    '''
+    This method is used to transform the original graph to the sharded graph.
+    The model parameters will be sharded according to the solution and the grad hooks
+    will be added to the sharded graph using the runtime_preparation_pass.
+    The communication node will be added into the graph using the runtime_apply_pass.
+    '''
+    gm, sharding_spec_dict, origin_spec_dict, comm_actions_dict = runtime_preparation_pass(
+        gm, solution, device_mesh, strategies_constructor)
+    gm = runtime_apply_pass(gm)
+    gm.recompile()
+    sharding_spec_dicts = (sharding_spec_dict, origin_spec_dict, comm_actions_dict)
+
+    return gm, sharding_spec_dicts
+
+
+def initialize_device_mesh(world_size: int = -1,
+                           alpha_beta_dict: Dict[Tuple[int], Tuple[float]] = None,
+                           logical_mesh_shape: Tuple[int] = None):
+    '''
+    This method is used to initialize the device mesh.
+
+    Args:
+        world_size(optional): the size of device mesh. If the world_size is -1,
+            the world size will be set to the number of GPUs in the current machine.
+        alpha_beta_dict(optional): the alpha_beta_dict contains the alpha and beta values
+            for each devices. if the alpha_beta_dict is None, the alpha_beta_dict will be
+            generated by profile_alpha_beta function.
+        logical_mesh_shape(optional): the logical_mesh_shape is used to specify the logical
+            mesh shape. If the logical_mesh_shape is None, the logical_mesh_shape will be
+            generated by search_best_logical_mesh_shape function.
+    '''
+    # if world_size is not set, use the world size from torch.distributed
+    if world_size == -1:
+        world_size = dist.get_world_size()
+    device1d = [i for i in range(world_size)]
+
+    if alpha_beta_dict is None:
+        # if alpha_beta_dict is not given, use a series of executions to profile alpha and beta values for each device
+        alpha_beta_dict = profile_alpha_beta(device1d)
+
+    if logical_mesh_shape is None:
+        # search for the best logical mesh shape
+        logical_mesh_shape = search_best_logical_mesh_shape(world_size, alpha_beta_dict)
+
+    # extract alpha and beta values for the chosen logical mesh shape
+    mesh_alpha, mesh_beta = extract_alpha_beta_for_device_mesh(alpha_beta_dict, logical_mesh_shape)
+    physical_mesh = torch.tensor(device1d)
+    device_mesh = DeviceMesh(physical_mesh_id=physical_mesh,
+                             mesh_shape=logical_mesh_shape,
+                             mesh_alpha=mesh_alpha,
+                             mesh_beta=mesh_beta,
+                             init_process_group=True)
+    return device_mesh
+
+
+def initialize_model(model: nn.Module,
+                     meta_args: Dict[str, torch.Tensor],
+                     device_mesh: DeviceMesh,
+                     memory_budget: float = -1.0,
+                     save_solver_solution: bool = False,
+                     load_solver_solution: bool = False,
+                     solution_path: str = None,
+                     return_solution: bool = False):
+    '''
+    This method is used to initialize the sharded model which could be used as normal pytorch model.
+
+    Args:
+        model: the model to be sharded.
+        meta_args: the meta_args is used to specify the input shapes of the model.
+        device_mesh: the device mesh to execute the model.
+        memory_budget(optional): the max cuda memory could be used. If the memory budget is -1.0,
+            the memory budget will be infinity.
+        save_solver_solution(optional): if the save_solver_solution is True, the solution will be saved
+            to the solution_path.
+        load_solver_solution(optional): if the load_solver_solution is True, the solution will be loaded
+            from the solution_path.
+        solution_path(optional): the path to save or load the solution.
+        return_solution(optional): if the return_solution is True, the solution will be returned. The returned
+            solution will be used to debug or help to analyze the sharding result. Therefore, we will not just
+            return a series of integers, but return the best strategies.
+    '''
+    tracer = ColoTracer()
+
+    graph = tracer.trace(root=model, meta_args=meta_args)
+    gm = GraphModule(model, graph, model.__class__.__name__)
+    gm.recompile()
+    strategies_constructor = build_strategy_constructor(graph, device_mesh)
+    if load_solver_solution:
+        solution = torch.load(solution_path)
+    else:
+        solution = solve_solution(gm, strategies_constructor, memory_budget)
+        if save_solver_solution:
+            torch.save(solution, solution_path)
+
+    gm, sharding_spec_dicts = transform_to_sharded_model(gm, solution, device_mesh, strategies_constructor)
+    model_to_return = ModuleWrapper(gm, *sharding_spec_dicts)
+
+    if return_solution:
+        solution_to_return = []
+        nodes = [strategies_vector.node for strategies_vector in strategies_constructor.leaf_strategies]
+        for index, node in enumerate(nodes):
+            solution_to_return.append(f'{node.name} {node.strategies_vector[solution[index]].name}')
+        return model_to_return, solution_to_return
+    else:
+        return model_to_return
+
+
+def autoparallelize(model: nn.Module,
+                    meta_args: Dict[str, torch.Tensor] = None,
+                    data_loader: torch.utils.data.DataLoader = None,
+                    data_process_func: callable = None,
+                    alpha_beta_dict: Dict[Tuple[int], Tuple[float]] = None,
+                    logical_mesh_shape: Tuple[int] = None,
+                    save_solver_solution: bool = False,
+                    load_solver_solution: bool = False,
+                    solver_solution_path: str = None,
+                    return_solution: bool = False,
+                    memory_budget: float = -1.0):
+    '''
+    This method is used to initialize the device mesh, extract the meta_args, and
+    use them to create a sharded model.
+
+    Args:
+        model: the model to be sharded.
+        meta_args(optional): the meta_args is used to specify the input shapes of the model.
+            If the meta_args is None, the meta_args will be extracted from the data_loader.
+        data_loader(optional): the data_loader to be used in normal training loop.
+        data_process_func(optional): the data_process_func is used to process the data from the data_loader.
+        alpha_beta_dict(optional): the alpha_beta_dict contains the alpha and beta values
+            for each devices. if the alpha_beta_dict is None, the alpha_beta_dict will be
+            generated by profile_alpha_beta function.
+        logical_mesh_shape(optional): the logical_mesh_shape is used to specify the logical
+            mesh shape. If the logical_mesh_shape is None, the logical_mesh_shape will be
+            generated by search_best_logical_mesh_shape function.
+        save_solver_solution(optional): if the save_solver_solution is True, the solution will be saved
+            to the solution_path.
+        load_solver_solution(optional): if the load_solver_solution is True, the solution will be loaded
+            from the solution_path.
+        solver_solution_path(optional): the path to save or load the solution.
+        return_solution(optional): if the return_solution is True, the solution will be returned.
+        memory_budget(optional): the max cuda memory could be used. If the memory budget is -1.0,
+            the memory budget will be infinity.
+    '''
+    device_mesh = initialize_device_mesh(alpha_beta_dict=alpha_beta_dict, logical_mesh_shape=logical_mesh_shape)
+    if meta_args is None:
+        meta_args = extract_meta_args_from_dataloader(data_loader, data_process_func)
+
+    rst_to_unpack = initialize_model(model,
+                                     meta_args,
+                                     device_mesh,
+                                     save_solver_solution=save_solver_solution,
+                                     load_solver_solution=load_solver_solution,
+                                     solver_solution_path=solver_solution_path,
+                                     return_solution=return_solution,
+                                     memory_budget=memory_budget)
+
+    if return_solution:
+        model, solution = rst_to_unpack
+        return model, solution
+    else:
+        model = rst_to_unpack
+        return model

colossalai/auto_parallel/tensor_shard/node_handler/__init__.py

+from .addmm_handler import ADDMMFunctionHandler
 from .batch_norm_handler import BatchNormModuleHandler
+from .binary_elementwise_handler import BinaryElementwiseHandler
 from .bmm_handler import AddBMMFunctionHandler, BMMFunctionHandler
 from .conv_handler import ConvFunctionHandler, ConvModuleHandler
+from .embedding_handler import EmbeddingFunctionHandler, EmbeddingModuleHandler
+from .experimental import PermuteHandler, ViewHandler
+from .getattr_handler import GetattrHandler
+from .getitem_handler import GetItemHandler
 from .layer_norm_handler import LayerNormModuleHandler
 from .linear_handler import LinearFunctionHandler, LinearModuleHandler
+from .matmul_handler import MatMulHandler
 from .normal_pooling_handler import NormPoolingHandler
-from .output_handler import OuputHandler
-from .placeholder_handler import PlacehodlerHandler
+from .output_handler import OutputHandler
+from .placeholder_handler import PlaceholderHandler
 from .registry import operator_registry
 from .reshape_handler import ReshapeHandler
+from .softmax_handler import SoftmaxHandler
+from .sum_handler import SumHandler
+from .tensor_constructor_handler import TensorConstructorHandler
 from .unary_elementwise_handler import UnaryElementwiseHandler
 from .where_handler import WhereHandler
 
 __all__ = [
     'LinearFunctionHandler', 'LinearModuleHandler', 'BMMFunctionHandler', 'AddBMMFunctionHandler',
     'LayerNormModuleHandler', 'BatchNormModuleHandler', 'ConvModuleHandler', 'ConvFunctionHandler',
-    'UnaryElementwiseHandler', 'ReshapeHandler', 'PlacehodlerHandler', 'OuputHandler', 'WhereHandler',
-    'NormPoolingHandler', 'operator_registry'
+    'UnaryElementwiseHandler', 'ReshapeHandler', 'PlaceholderHandler', 'OutputHandler', 'WhereHandler',
+    'NormPoolingHandler', 'BinaryElementwiseHandler', 'MatMulHandler', 'operator_registry', 'ADDMMFunctionHandler',
+    'GetItemHandler', 'GetattrHandler', 'ViewHandler', 'PermuteHandler', 'TensorConstructorHandler',
+    'EmbeddingModuleHandler', 'EmbeddingFunctionHandler', 'SumHandler', 'SoftmaxHandler'
 ]

colossalai/auto_parallel/tensor_shard/node_handler/addmm_handler.py

+from typing import Dict, List, Union
+
+import torch
+
+from colossalai.tensor.shape_consistency import CollectiveCommPattern, CommSpec, ShapeConsistencyManager
+
+from ..sharding_strategy import CommAction, CommType, OperationData, OperationDataType, ShardingStrategy
+from ..utils import comm_actions_for_oprands, recover_sharding_spec_for_broadcast_shape
+from .node_handler import NodeHandler
+from .registry import operator_registry
+from .strategy import LinearProjectionStrategyGenerator, StrategyGenerator
+
+__all__ = ['ADDMMFunctionHandler']
+
+
+@operator_registry.register(torch.addmm)
+@operator_registry.register(torch.Tensor.addmm)
+class ADDMMFunctionHandler(NodeHandler):
+    """
+    This is a NodeHandler class which deals with the batched matrix multiplication operation in PyTorch.
+    Such operations including `torch.bmm` and `torch.Tensor.bmm` require the tensor to be 3D, thus, there is
+    no logical-physical shape conversion in this handler.
+    """
+
+    def _infer_op_data_type(self, tensor: torch.Tensor) -> OperationDataType:
+        if isinstance(tensor, torch.nn.parameter.Parameter):
+            data_type = OperationDataType.PARAM
+        else:
+            data_type = OperationDataType.ARG
+        return data_type
+
+    def get_operation_data_mapping(self) -> Dict[str, OperationData]:
+
+        # input operand
+        input_data = self.node.args[1]._meta_data
+        physical_input_operand = OperationData(name=str(self.node.args[1]),
+                                               type=self._infer_op_data_type(input_data),
+                                               data=input_data)
+
+        # other operand
+        other_data = self.node.args[2]._meta_data
+        physical_other_operand = OperationData(name=str(self.node.args[2]),
+                                               type=self._infer_op_data_type(other_data),
+                                               data=other_data)
+        # bias physical shape
+        bias_logical_shape = self.node._meta_data.shape
+        bias_data = self.node.args[0]._meta_data
+        physical_bias_operand = OperationData(name=str(self.node.args[0]),
+                                              type=self._infer_op_data_type(bias_data),
+                                              data=bias_data,
+                                              logical_shape=bias_logical_shape)
+
+        # output
+        physical_output = OperationData(name=str(self.node), type=OperationDataType.OUTPUT, data=self.node._meta_data)
+
+        mapping = {
+            "input": physical_input_operand,
+            "other": physical_other_operand,
+            "output": physical_output,
+            'bias': physical_bias_operand
+        }
+
+        return mapping
+
+    def get_strategy_generator(self) -> List[StrategyGenerator]:
+        op_data_mapping = self.get_operation_data_mapping()
+        generators = []
+        generators.append(
+            LinearProjectionStrategyGenerator(op_data_mapping, self.device_mesh, linear_projection_type='addmm'))
+        return generators
+
+    def post_process(self, strategy: ShardingStrategy) -> Union[ShardingStrategy, List[ShardingStrategy]]:
+        # convert bias from its logical sharding spec to its physical sharding spec
+        op_data_mapping = self.get_operation_data_mapping()
+
+        bias_op_data = op_data_mapping['bias']
+        bias_physical_shape = bias_op_data.data.shape
+        bias_logical_shape = bias_op_data.logical_shape
+        bias_sharding_spec = strategy.get_sharding_spec_by_name(bias_op_data.name)
+        bias_sharding_spec, removed_dims = recover_sharding_spec_for_broadcast_shape(
+            bias_sharding_spec, bias_logical_shape, bias_physical_shape)
+        strategy.sharding_specs[bias_op_data] = bias_sharding_spec
+
+        if len(removed_dims) > 0:
+            comm_action = comm_actions_for_oprands(node=self.node,
+                                                   removed_dims=removed_dims,
+                                                   op_data=bias_op_data,
+                                                   sharding_spec=bias_sharding_spec)
+            strategy.communication_actions[bias_op_data] = comm_action
+
+        return strategy

colossalai/auto_parallel/tensor_shard/node_handler/batch_norm_handler.py

@@ -2,8 +2,10 @@ from typing import Dict, List
 
 import torch
 
-from ..sharding_strategy import OperationData, OperationDataType
-from .node_handler import ModuleHandler
+from colossalai.auto_parallel.meta_profiler.metainfo import MetaInfo
+
+from ..sharding_strategy import OperationData, OperationDataType, StrategiesVector
+from .node_handler import MetaInfoModuleHandler, ModuleHandler
 from .registry import operator_registry
 from .strategy import BatchNormStrategyGenerator, StrategyGenerator
 
@@ -13,7 +15,7 @@ __all__ = ['BatchNormModuleHandler']
 @operator_registry.register(torch.nn.BatchNorm1d)
 @operator_registry.register(torch.nn.BatchNorm2d)
 @operator_registry.register(torch.nn.BatchNorm3d)
-class BatchNormModuleHandler(ModuleHandler):
+class BatchNormModuleHandler(MetaInfoModuleHandler):
     """
     A BatchNormModuleHandler which deals with the sharding strategies for nn.BatchNormXd module.
     """

colossalai/auto_parallel/tensor_shard/node_handler/binary_elementwise_handler.py

+from typing import Dict, List, Union
+
+import torch
+from torch.fx.node import Node
+
+from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, ShardingStrategy
+from colossalai.tensor.shape_consistency import CollectiveCommPattern, CommSpec, ShapeConsistencyManager
+
+from ..constants import BCAST_FUNC_OP
+from ..utils import comm_actions_for_oprands, recover_sharding_spec_for_broadcast_shape
+from .node_handler import MetaInfoNodeHandler, NodeHandler
+from .registry import operator_registry
+from .strategy import BinaryElementwiseStrategyGenerator, StrategyGenerator
+
+__all__ = ['BinaryElementwiseHandler']
+
+
+@operator_registry.register(BCAST_FUNC_OP)
+class BinaryElementwiseHandler(MetaInfoNodeHandler):
+    """
+    An BinaryBcastOpHandler is a node handler which deals with operations which have two
+    operands and broadcasting occurs such as torch.add.
+    """
+
+    def get_operation_data_mapping(self) -> Dict[str, OperationData]:
+        bcast_shape = self.node._meta_data.shape
+
+        def _get_op_data_type(tensor):
+            if isinstance(tensor, torch.nn.parameter.Parameter):
+                return OperationDataType.PARAM
+            else:
+                return OperationDataType.ARG
+
+        def _get_arg_value(idx):
+            if isinstance(self.node.args[idx], Node):
+                meta_data = self.node.args[idx]._meta_data
+            else:
+                # this is in fact a real data like int 1
+                # but we can deem it as meta data
+                # as it won't affect the strategy generation
+                assert isinstance(self.node.args[idx], (int, float))
+                meta_data = torch.Tensor([self.node.args[idx]]).to('meta')
+            return meta_data
+
+        input_meta_data = _get_arg_value(0)
+        other_meta_data = _get_arg_value(1)
+        output_meta_data = self.node._meta_data
+
+        input_op_data = OperationData(name=str(self.node.args[0]),
+                                      type=_get_op_data_type(input_meta_data),
+                                      data=input_meta_data,
+                                      logical_shape=bcast_shape)
+        other_op_data = OperationData(name=str(self.node.args[1]),
+                                      type=_get_op_data_type(other_meta_data),
+                                      data=other_meta_data,
+                                      logical_shape=bcast_shape)
+        output_op_data = OperationData(name=str(self.node),
+                                       type=OperationDataType.OUTPUT,
+                                       data=output_meta_data,
+                                       logical_shape=bcast_shape)
+
+        mapping = {'input': input_op_data, 'other': other_op_data, 'output': output_op_data}
+        return mapping
+
+    def get_strategy_generator(self) -> List[StrategyGenerator]:
+        op_data_mapping = self.get_operation_data_mapping()
+        generators = []
+        generators.append(BinaryElementwiseStrategyGenerator(op_data_mapping, self.device_mesh))
+        return generators
+
+    def post_process(self, strategy: ShardingStrategy) -> Union[ShardingStrategy, List[ShardingStrategy]]:
+        # convert bias from its logical sharding spec to its physical sharding spec
+        op_data_mapping = self.get_operation_data_mapping()
+
+        for op_name, op_data in op_data_mapping.items():
+            if not isinstance(op_data.data, torch.Tensor):
+                # remove the sharding spec if the op_data is not a tensor, e.g. torch.pow(tensor, 2)
+                strategy.sharding_specs.pop(op_data)
+            else:
+                # convert the logical sharding spec to physical sharding spec if broadcast
+                # e.g. torch.rand(4, 4) + torch.rand(4)
+                physical_shape = op_data.data.shape
+                logical_shape = op_data.logical_shape
+                sharding_spec = strategy.get_sharding_spec_by_name(op_data.name)
+                sharding_spec, removed_dims = recover_sharding_spec_for_broadcast_shape(
+                    sharding_spec, logical_shape, physical_shape)
+
+                strategy.sharding_specs[op_data] = sharding_spec
+                if len(removed_dims) > 0:
+                    comm_action = comm_actions_for_oprands(node=self.node,
+                                                           removed_dims=removed_dims,
+                                                           op_data=op_data,
+                                                           sharding_spec=sharding_spec)
+                    strategy.communication_actions[op_data] = comm_action
+
+        return strategy