[autoparallel] adapt autoparallel tests with latest api

colossalai/auto_parallel/tensor_shard/node_handler/strategy/matmul_strategy_generator.py

@@ -247,12 +247,12 @@ def collate_strategies(self) -> List[ShardingStrategy]:
         strategies.append(self.split_rhs_space_both_contract(1, 0))
 
         # RR= RS x SR
-        # strategies.append(self.recompute_split_both_contract(0))
-        # strategies.append(self.recompute_split_both_contract(1))
+        strategies.append(self.recompute_split_both_contract(0))
+        strategies.append(self.recompute_split_both_contract(1))
 
-        # # RS = RR x RS
-        # strategies.append(self.split_rhs_space_only(0))
-        # strategies.append(self.split_rhs_space_only(1))
+        # RS = RR x RS
+        strategies.append(self.split_rhs_space_only(0))
+        strategies.append(self.split_rhs_space_only(1))
 
         # S01R = S01R x RR
         strategies.append(self.split_lhs_1st_dim_1d(0, 1))
@@ -263,8 +263,8 @@ def collate_strategies(self) -> List[ShardingStrategy]:
         # RS01 = RR x RS01
         strategies.append(self.split_rhs_2nd_dim_1d(0, 1))
 
-        # # RR = RR x RR
-        # strategies.append(self.non_split())
+        # RR = RR x RR
+        strategies.append(self.non_split())
 
         return strategies
 

colossalai/auto_parallel/tensor_shard/solver/cost_graph.py

@@ -62,9 +62,6 @@ def _build_cost_graph(self):
                         else:
                             edge_cost[(j, i)] = resharding_cost_item.total
                 self.edge_costs[node_pair] = edge_cost
-            # add parents and children attribute to node
-            # parent_nodes = [node for node in strategies_vector.predecessor_nodes]
-            # children_nodes = [node for node in strategies_vector.successor_nodes]
             parent_nodes = []
             children_nodes = []
 

tests/test_auto_parallel/test_tensor_shard/test_bias_addition_forward.py

@@ -4,21 +4,11 @@
 import torch
 import torch.multiprocessing as mp
 
-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 OperationDataType
-from colossalai.auto_parallel.tensor_shard.solver import (
-    CostGraph,
-    GraphAnalyser,
-    Solver,
-    SolverOptions,
-    StrategiesConstructor,
-)
+from colossalai.auto_parallel.tensor_shard.initialize import initialize_model
 from colossalai.device.device_mesh import DeviceMesh
-from colossalai.fx import ColoGraphModule, ColoTracer
 from colossalai.initialize import launch
 from colossalai.logging import disable_existing_loggers
-from colossalai.testing import assert_close, assert_close_loose, rerun_if_address_is_in_use
+from colossalai.testing import assert_close, rerun_if_address_is_in_use
 from colossalai.testing.pytest_wrapper import run_on_environment_flag
 from colossalai.utils import free_port
 
@@ -63,42 +53,9 @@ def check_linear_module(rank, world_size, port):
     # [[0, 1]
     #  [2, 3]]
     device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True)
-    tracer = ColoTracer()
-    # graph():
-    #     %x : torch.Tensor [#users=1] = placeholder[target=x]
-    #     %linear_weight : [#users=1] = get_attr[target=linear.weight]
-    #     %linear_bias : [#users=1] = get_attr[target=linear.bias]
-    #     %linear : [#users=1] = call_function[target=torch._C._nn.linear](args = (%x, %linear_weight), kwargs = {})
-    #     %add : [#users=1] = call_function[target=operator.add](args = (%linear, %linear_bias), kwargs = {})
-    #     %mul : [#users=1] = call_function[target=operator.mul](args = (%add, 2), kwargs = {})
-    #     return mul
-    graph = tracer.trace(root=model, meta_args={'x': torch.rand(4, 4).to('meta')})
-    # def forward(self, x : torch.Tensor):
-    #     linear_weight = self.linear.weight
-    #     linear_bias = self.linear.bias
-    #     linear = torch._C._nn.linear(x, linear_weight);  x = linear_weight = None
-    #     add = linear + linear_bias;  linear = linear_bias = None
-    #     mul = add * 2;  add = None
-    #     return mul
-    gm = ColoGraphModule(model, graph)
-    gm.recompile()
-    node_list = list(graph.nodes)
-
-    solver_options = SolverOptions()
-    strategies_constructor = StrategiesConstructor(graph, device_mesh, solver_options)
-    strategies_constructor.build_strategies_and_cost()
-    linear_node = node_list[3]
-    cost_graph = CostGraph(strategies_constructor.leaf_strategies)
-    cost_graph.simplify_graph()
-    graph_analyser = GraphAnalyser(gm)
-    solver = Solver(gm.graph, strategies_constructor, cost_graph, graph_analyser)
-    ret = solver.call_solver_serialized_args()
-    solution = list(ret[0])
-    gm, sharding_spec_dict, origin_spec_dict, comm_actions_dict = runtime_preparation_pass(gm, solution, device_mesh)
-
-    gm = runtime_apply_pass(gm)
-    gm.recompile()
-    output = gm(input, sharding_spec_dict, origin_spec_dict, comm_actions_dict)
+    meta_args = {'x': torch.rand(4, 4).to('meta')}
+    gm = initialize_model(model, meta_args=meta_args, device_mesh=device_mesh)
+    output = gm(input)
     assert_close(output, output_compare)
 
 
@@ -113,47 +70,9 @@ def check_conv_module(rank, world_size, port):
     # [[0, 1]
     #  [2, 3]]
     device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True)
-    tracer = ColoTracer()
-    # graph():
-    #     %x : torch.Tensor [#users=1] = placeholder[target=x]
-    #     %conv_weight : [#users=1] = get_attr[target=conv.weight]
-    #     %conv_bias : [#users=1] = get_attr[target=conv.bias]
-    #     %conv2d : [#users=1] = call_function[target=torch.conv2d](args = (%x, %conv_weight), kwargs = {})
-    #     %view : [#users=1] = call_method[target=view](args = (%conv_bias, [1, -1, 1, 1]), kwargs = {})
-    #     %add : [#users=1] = call_function[target=operator.add](args = (%conv2d, %view), kwargs = {})
-    #     %mul : [#users=1] = call_function[target=operator.mul](args = (%add, 2), kwargs = {})
-    #     return mul
-    graph = tracer.trace(root=model, meta_args={'x': torch.rand(4, 3, 64, 64).to('meta')})
-    # def forward(self, x : torch.Tensor):
-    #     conv_weight = self.conv.weight
-    #     conv_bias = self.conv.bias
-    #     conv2d = torch.conv2d(x, conv_weight);  x = conv_weight = None
-    #     view = conv_bias.view([1, -1, 1, 1]);  conv_bias = None
-    #     add = conv2d + view;  conv2d = view = None
-    #     mul = add * 2;  add = None
-    #     return mul
-    gm = ColoGraphModule(model, graph)
-
-    gm.recompile()
-
-    node_list = list(graph.nodes)
-    conv_node = node_list[3]
-    solver_options = SolverOptions()
-    strategies_constructor = StrategiesConstructor(graph, device_mesh, solver_options)
-    strategies_constructor.build_strategies_and_cost()
-
-    cost_graph = CostGraph(strategies_constructor.leaf_strategies)
-    cost_graph.simplify_graph()
-    graph_analyser = GraphAnalyser(gm)
-    solver = Solver(gm.graph, strategies_constructor, cost_graph, graph_analyser)
-    ret = solver.call_solver_serialized_args()
-    solution = list(ret[0])
-
-    gm, sharding_spec_dict, origin_spec_dict, comm_actions_dict = runtime_preparation_pass(gm, solution, device_mesh)
-
-    gm = runtime_apply_pass(gm)
-    gm.recompile()
-    output = gm(input, sharding_spec_dict, origin_spec_dict, comm_actions_dict)
+    meta_args = {'x': torch.rand(4, 3, 64, 64).to('meta')}
+    gm = initialize_model(model, meta_args=meta_args, device_mesh=device_mesh)
+    output = gm(input)
     assert_close(output, output_compare)