[fx]add autoparallel passes (#1121)

* [CLI] add CLI launcher

* Revert "[CLI] add CLI launcher"

This reverts commit df7e6506d4.

* feature/add autoparallel passes

colossalai/fx/passes/shard_1d_pass.py

@@ -0,0 +1,120 @@
+import torch
+from torch.fx.node import map_arg
+from torch.fx.node import Node
+from torch.fx.passes.split_module import split_module
+
+import colossalai
+from colossalai.context import ParallelMode
+from colossalai.core import global_context as gpc
+
+
+def all_gather_function(input_):
+    world_size = gpc.get_world_size(ParallelMode.PARALLEL_1D)
+    rank = gpc.get_local_rank(ParallelMode.PARALLEL_1D)
+    tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
+    tensor_list[rank] = input_
+    group = gpc.get_group(ParallelMode.PARALLEL_1D)
+    torch.distributed.all_gather(tensor_list, input_, group=group)
+    output = torch.cat(tensor_list, dim=-1).contiguous()
+    return output
+
+
+def all_reduce_function(input_):
+    if gpc.get_world_size(ParallelMode.PARALLEL_1D) == 1:
+        return input_
+    torch.distributed.all_reduce(input_, group=gpc.get_group(ParallelMode.PARALLEL_1D))
+    return input_
+
+
+def weight_split(weight, dim):
+    #TODO: this function will be refactored by using ColoTensor dist_spec when a stable reshaper feature is ready to use.
+    num_partition = gpc.get_world_size(ParallelMode.TENSOR)
+    shape = weight.shape
+    length = shape[dim] // num_partition
+    sharded_weight_list = []
+    for i in range(num_partition):
+        sharded_weight_list.append(weight.narrow(dim, i * length, length))
+    return sharded_weight_list[gpc.get_local_rank(ParallelMode.PARALLEL_1D)]
+
+
+def replace_all_uses_except_replaced(node, replace_node):
+    """
+        Replace all uses of ``node`` in the Graph with the Node ``replace_node``,
+        except the user of ``node`` is ``replace_node``.
+
+        Args:
+
+            replace_node (Node): The node to replace all uses of ``node`` with.
+
+        Returns:
+
+            The list of Nodes on which this change was made.
+    """
+    to_process = list(node.users)
+    for use_node in to_process:
+        if use_node == replace_node:
+            continue
+
+        def may_replace_node(n):
+            if n == node:
+                return replace_node
+            else:
+                return n
+
+        new_args = map_arg(use_node.args, may_replace_node)
+        new_kwargs = map_arg(use_node.kwargs, may_replace_node)
+        use_node._args = new_args
+        use_node._kwargs = new_kwargs
+        for old_use in use_node._input_nodes.keys():
+            old_use.users.pop(use_node)
+        use_node._input_nodes = {}
+        map_arg(use_node._args, lambda n: use_node._input_nodes.setdefault(n))
+        map_arg(use_node._kwargs, lambda n: use_node._input_nodes.setdefault(n))
+        for new_use in use_node._input_nodes.keys():
+            new_use.users.setdefault(use_node)
+    return to_process
+
+
+def column_shard_linear_pass(gm: torch.fx.GraphModule):
+    mod_graph = gm.graph
+    for node in mod_graph.nodes:
+        if node.op == "call_module":
+            target_module = node.graph.owning_module.get_submodule(node.target)
+            if isinstance(target_module, torch.nn.Linear):
+                target_module.weight.data = weight_split(target_module.weight.data, dim=0)
+                if target_module.bias is not None:
+                    target_module.bias.data = weight_split(target_module.bias.data, dim=0)
+
+                # inserting communication node after the sharded linear node
+                with mod_graph.inserting_after(node):
+                    new_node = mod_graph.create_node('call_function', all_gather_function, args=(node,))
+                    replace_all_uses_except_replaced(node, new_node)
+    gm.recompile()
+    return gm
+
+
+def row_shard_linear_pass(gm: torch.fx.GraphModule):
+    mod_graph = gm.graph
+    for node in mod_graph.nodes:
+        if node.op == "call_module":
+            target_module = node.graph.owning_module.get_submodule(node.target)
+            if isinstance(target_module, torch.nn.Linear):
+                target_module.weight.data = weight_split(target_module.weight.data, dim=-1)
+
+                # insert input sharding node before the sharded linear node
+                with mod_graph.inserting_before(node):
+                    input_node_list = list(node._input_nodes.keys())
+                    assert len(input_node_list) == 1, 'linear forward must have and only have one input tensor.'
+                    input_node = input_node_list[0]
+                    new_input_node = mod_graph.create_node('call_function', weight_split, args=(input_node, -1))
+                    replace_all_uses_except_replaced(input_node, new_input_node)
+
+                # inserting communication node after the sharded linear node
+                with mod_graph.inserting_after(node):
+                    new_node = mod_graph.create_node('call_function', all_reduce_function, args=(node,))
+                    replace_all_uses_except_replaced(node, new_node)
+    gm.recompile()
+    return gm
+
+
+#TODO: add elementwise op process pass, then we can try to use column and row mixed strategy.