[autoparallel] adapt solver with resnet (#1583)

* [autoparallel]adapt solver with resnet

* polish code

* polish code

tests/test_auto_parallel/test_batch_norm_handler.py

@@ -0,0 +1,119 @@
+import torch
+from torch.fx import GraphModule
+import torch.nn as nn
+import pytest
+
+from colossalai.fx.proxy import ColoProxy
+from colossalai.fx.tracer.tracer import ColoTracer
+from colossalai.tensor.sharding_spec import ShardingSpec, _DimSpec
+from colossalai.auto_parallel.solver.op_handler.batch_norm_handler import BatchNormHandler
+from colossalai.auto_parallel.solver.sharding_strategy import ShardingStrategy, StrategiesVector
+from colossalai.tensor.shape_consistency import ShapeConsistencyManager
+from colossalai.device.device_mesh import DeviceMesh
+
+
+class BNModel(nn.Module):
+
+    def __init__(self, c):
+        super().__init__()
+        self.bn = nn.BatchNorm2d(c)
+
+    def forward(self, x):
+        x = x * 2
+        x = self.bn(x)
+        return x
+
+
+def test_bn_handler():
+    physical_mesh_id = torch.arange(0, 4)
+    mesh_shape = (2, 2)
+    # [[0, 1]
+    #  [2, 3]]
+    device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
+    entire_shape = torch.Size((4, 16, 64, 64))
+    shape_consistency_manager = ShapeConsistencyManager()
+
+    tracer = ColoTracer()
+    model = BNModel(16)
+    input_sample = {'x': torch.rand(4, 16, 64, 64).to('meta')}
+    # graph():
+    #     %x : torch.Tensor [#users=1] = placeholder[target=x]
+    #     %mul : [#users=1] = call_function[target=operator.mul](args = (%x, 2), kwargs = {})
+    #     %bn : [#users=1] = call_module[target=bn](args = (%mul,), kwargs = {})
+    #     return bn
+    graph = tracer.trace(root=model, meta_args=input_sample)
+    gm = GraphModule(model, graph, model.__class__.__name__)
+    gm.recompile()
+    # [x, mul, bn, output]
+    nodes = [node for node in gm.graph.nodes]
+
+    # find the sharding strategies for the input node of the bn node
+    # 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(nodes[1])
+    sharding_option = (None, 0, 1)
+    for first_sharding_index in sharding_option:
+        for second_sharding_index in sharding_option:
+            if first_sharding_index is not None and second_sharding_index == first_sharding_index:
+                continue
+            if first_sharding_index is None:
+                first_dim_spec = _DimSpec([])
+            else:
+                first_dim_spec = _DimSpec([first_sharding_index])
+
+            if second_sharding_index is None:
+                second_dim_spec = _DimSpec([])
+            else:
+                second_dim_spec = _DimSpec([second_sharding_index])
+
+            replica_dim_spec = _DimSpec([])
+            sharding_sequence = [first_dim_spec, second_dim_spec, replica_dim_spec, replica_dim_spec]
+            sharding_spec = ShardingSpec(device_mesh=device_mesh,
+                                         entire_shape=entire_shape,
+                                         sharding_sequence=sharding_sequence)
+            strategy_name = str(sharding_spec.sharding_sequence)
+            sharding_strategy = ShardingStrategy(name=strategy_name, output_sharding_spec=sharding_spec)
+            strategies_vector_for_input.append(sharding_strategy)
+    setattr(nodes[1], 'strategies_vector', strategies_vector_for_input)
+
+    # generate bn strategy
+    strategies_vector = StrategiesVector(node=nodes[2])
+    bn_handler = BatchNormHandler(node=nodes[2],
+                                  device_mesh=device_mesh,
+                                  strategies_vector=strategies_vector,
+                                  shape_consistency_manager=shape_consistency_manager)
+    bn_handler.register_strategy()
+    # ['RS0 = RS0 x S0', 'S1S0 = RS0 x S0', 'RS1 = RS1 x S1', 'S0S1 = RS1 x S1', 'RR = RR x R', 'S0R = RR x R', 'S1R = RR x R', 'S01R = RR x R', 'RS01 = RS01 x S01',
+    # 'S0R = S0R x R WITH SYNC_BN', 'S1R = S1R x R WITH SYNC_BN', 'S0S1 = S0S1 x S1 WITH SYNC_BN', 'S1S0 = S1S0 x S0 WITH SYNC_BN', 'S01R = S01R x R WITH SYNC_BN']
+    strategy_name_list = [strategy.name for strategy in bn_handler.strategies_vector]
+
+    # RS = RS x S and strategies based on it, such as
+    # SS = RS x S
+    assert 'RS0 = RS0 x S0' in strategy_name_list
+    assert 'S1S0 = RS0 x S0' in strategy_name_list
+    assert 'RS1 = RS1 x S1' in strategy_name_list
+    assert 'S0S1 = RS1 x S1' in strategy_name_list
+
+    # RR = RR x R and strategies based on it, such as
+    # SR = SR x R
+    assert 'RR = RR x R' in strategy_name_list
+    assert 'S0R = RR x R' in strategy_name_list
+    assert 'S1R = RR x R' in strategy_name_list
+    assert 'S01R = RR x R' in strategy_name_list
+
+    # RS01 = RS01 x S01
+    assert 'RS01 = RS01 x S01' in strategy_name_list
+
+    # SR = SR x R WITH SYNC_BN
+    assert 'S0R = S0R x R WITH SYNC_BN' in strategy_name_list
+    assert 'S1R = S1R x R WITH SYNC_BN' in strategy_name_list
+
+    # SS = SS x S WITH SYNC_BN
+    assert 'S0S1 = S0S1 x S1 WITH SYNC_BN' in strategy_name_list
+    assert 'S1S0 = S1S0 x S0 WITH SYNC_BN' in strategy_name_list
+
+    # S01R = S01R x R WITH SYNC_BN
+    assert 'S01R = S01R x R WITH SYNC_BN' in strategy_name_list
+
+
+if __name__ == '__main__':
+    test_bn_handler()

tests/test_auto_parallel/test_conv_handler.py

@@ -6,7 +6,7 @@ import pytest
 from colossalai.fx.proxy import ColoProxy
 from colossalai.fx.tracer.tracer import ColoTracer
 from colossalai.tensor.sharding_spec import ShardingSpec, _DimSpec
-from colossalai.auto_parallel.solver.conv_handler import ConvHandler
+from colossalai.auto_parallel.solver.op_handler.conv_handler import ConvHandler
 from colossalai.auto_parallel.solver.sharding_strategy import ShardingStrategy, StrategiesVector
 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
 from colossalai.device.device_mesh import DeviceMesh

tests/test_auto_parallel/test_cost_graph.py

@@ -6,8 +6,6 @@ import pytest
 from colossalai.fx.proxy import ColoProxy
 from colossalai.fx.tracer.tracer import ColoTracer
 from colossalai.tensor.sharding_spec import ShardingSpec, _DimSpec
-from colossalai.auto_parallel.solver.conv_handler import ConvHandler, CONV_STRATEGIES_LIST
-from colossalai.auto_parallel.solver.sharding_strategy import ShardingStrategy, StrategiesVector
 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
 from colossalai.device.device_mesh import DeviceMesh
 from colossalai.auto_parallel.solver.strategies_constructor import StrategiesConstructor

tests/test_auto_parallel/test_dot_handler.py

@@ -6,7 +6,7 @@ import pytest
 from colossalai.fx.proxy import ColoProxy
 from colossalai.fx.tracer.tracer import ColoTracer
 from colossalai.tensor.sharding_spec import ShardingSpec, _DimSpec
-from colossalai.auto_parallel.solver.dot_handler import DotHandler
+from colossalai.auto_parallel.solver.op_handler.dot_handler import DotHandler
 from colossalai.auto_parallel.solver.sharding_strategy import ShardingStrategy, StrategiesVector
 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
 from colossalai.device.device_mesh import DeviceMesh

tests/test_auto_parallel/test_liveness_analysis.py

@@ -32,21 +32,21 @@ def test_liveness_analysis():
     gm = ColoGraphModule(root=model, graph=graph, class_name=model.__class__.__name__)
 
     graph_analyser = GraphAnalyser(gm)
-    liveness_dict = graph_analyser.liveness_analysis()
-    stage_count = len(liveness_dict)
+    liveness_list = graph_analyser.liveness_analysis()
+    stage_count = len(liveness_list)
 
-    # 8 stages including input and output
-    assert stage_count == 8
+    # if a LiveStage is covered by another LiveStage, we just keep the larger one.
+    assert stage_count == 1
 
     # a variable named `relu` must exist
     # and this live var must have inplace = True
-    assert liveness_dict[5].all_live_vars.exists('relu')
-    relu_var = liveness_dict[5].all_live_vars.get('relu')
+    assert liveness_list[0].all_live_vars.exists('relu')
+    relu_var = liveness_list[0].all_live_vars.get('relu')
     assert relu_var.is_inplace
 
     # the unique vars must be fewer than the all vars since in-place ops exist
-    all_live_vars = liveness_dict[7].all_live_vars
-    unique_live_vars = liveness_dict[7].unique_live_vars
+    all_live_vars = liveness_list[0].all_live_vars
+    unique_live_vars = liveness_list[0].unique_live_vars
     assert len(unique_live_vars) + 1 == len(all_live_vars)
 
 

tests/test_auto_parallel/test_solver.py

@@ -0,0 +1,78 @@
+import torch
+from torch.fx import GraphModule
+import torch.nn as nn
+import pytest
+
+from colossalai.fx.tracer.tracer import ColoTracer
+from colossalai.tensor.shape_consistency import ShapeConsistencyManager
+from colossalai.device.device_mesh import DeviceMesh
+from colossalai.auto_parallel.solver.strategies_constructor import StrategiesConstructor
+from colossalai.auto_parallel.solver.cost_graph import CostGraph
+from colossalai.auto_parallel.solver.graph_analysis import GraphAnalyser
+from copy import deepcopy
+from colossalai.auto_parallel.solver import Solver
+
+
+class ConvModel(nn.Module):
+
+    def __init__(self, c_in, c_out):
+        super().__init__()
+        self.conv1 = nn.Conv2d(c_in, c_out, kernel_size=3)
+        self.conv2 = nn.Conv2d(c_out, c_out, kernel_size=3)
+        self.conv3 = nn.Conv2d(c_out, c_out, kernel_size=3)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = x * 2
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = x / 2
+        x = self.conv3(x)
+        x = self.relu(x)
+        return x
+
+
+@pytest.mark.skip("for higher testing speed")
+def test_solver():
+    physical_mesh_id = torch.arange(0, 4)
+    mesh_shape = (2, 2)
+    # [[0, 1]
+    #  [2, 3]]
+    device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
+    entire_shape = torch.Size((4, 16, 64, 64))
+    shape_consistency_manager = ShapeConsistencyManager()
+
+    tracer = ColoTracer()
+    model = ConvModel(16, 32)
+    input_sample = {'x': torch.rand(4, 16, 64, 64).to('meta')}
+
+    # graph():
+    #     %x : torch.Tensor [#users=1] = placeholder[target=x]
+    #     %mul : [#users=1] = call_function[target=operator.mul](args = (%x, 2), kwargs = {})
+    #     %conv1 : [#users=1] = call_module[target=conv1](args = (%mul,), kwargs = {})
+    #     %conv2 : [#users=1] = call_module[target=conv2](args = (%conv1,), kwargs = {})
+    #     %truediv : [#users=1] = call_function[target=operator.truediv](args = (%conv2, 2), kwargs = {})
+    #     %conv3 : [#users=1] = call_module[target=conv3](args = (%truediv,), kwargs = {})
+    #     %relu : [#users=1] = call_module[target=relu](args = (%conv3,), kwargs = {})
+    #     return relu
+    graph = tracer.trace(root=model, meta_args=input_sample)
+    gm = GraphModule(model, graph, model.__class__.__name__)
+    gm.recompile()
+
+    solver_options = {'fast_mode': True}
+    strategies_constructor = StrategiesConstructor(graph, device_mesh, shape_consistency_manager, 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()
+
+    # [ 0 0 13 13 13 13 13 0]
+    strategies_combination_list = ret[0]
+    assert solver.leaf_strategies[2][13].name == 'S01R = S01R x RR'
+
+
+if __name__ == '__main__':
+    test_solver()

tests/test_auto_parallel/test_strategies_constructor.py

@@ -6,7 +6,7 @@ import pytest
 from colossalai.fx.proxy import ColoProxy
 from colossalai.fx.tracer.tracer import ColoTracer
 from colossalai.tensor.sharding_spec import ShardingSpec, _DimSpec
-from colossalai.auto_parallel.solver.conv_handler import ConvHandler, CONV_STRATEGIES_LIST
+from colossalai.auto_parallel.solver.op_handler.conv_handler import CONV_STRATEGIES_LIST
 from colossalai.auto_parallel.solver.sharding_strategy import ShardingStrategy, StrategiesVector
 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
 from colossalai.device.device_mesh import DeviceMesh