github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 01:28:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

[autoparallel] add output handler and placeholder handler (#1694)

Browse Source 
* [autoparallel] add output handler and placeholder handler

* Delete test_solver_with_resnet.py

* fix test bugs
This commit is contained in:
YuliangLiu0306
2022-10-13 13:42:36 +08:00
committed by GitHub
parent 56088e6d98
commit 42b882ef06
19 changed files with 343 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_auto_parallel/test_node_handler/test_batch_norm_handler_v2.py
View File

@@ -58,7 +58,7 @@ def test_bn_module_handler():
assert mapping['output'].data.shape == torch.Size([4, 16, 64, 64])
assert mapping['output'].type == OperationDataType.OUTPUT
strategies_vector = handler.register_strategy()
strategies_vector = handler.register_strategy(compute_resharding_cost=False)
strategy_name_list = [val.name for val in strategies_vector]
# RS = RS x S

4
tests/test_auto_parallel/test_node_handler/test_bmm_handler.py
View File

@@ -68,7 +68,7 @@ def test_2d_device_mesh(module):
assert mapping['output'].data.shape == torch.Size([4, 8, 8])
assert mapping['output'].type == OperationDataType.OUTPUT
strategies_vector = handler.register_strategy()
strategies_vector = handler.register_strategy(compute_resharding_cost=False)
strategy_name_list = [val.name for val in strategies_vector]
# one batch dim
@@ -138,7 +138,7 @@ def test_1d_device_mesh(module):
assert mapping['output'].data.shape == torch.Size([4, 8, 8])
assert mapping['output'].type == OperationDataType.OUTPUT
strategies_vector = handler.register_strategy()
strategies_vector = handler.register_strategy(compute_resharding_cost=False)
strategy_name_list = [val.name for val in strategies_vector]
assert len(strategy_name_list) == 1
# one batch dim

4
tests/test_auto_parallel/test_node_handler/test_conv_handler_v2.py
View File

@@ -58,7 +58,7 @@ def test_conv_module_handler():
assert mapping['output'].data.shape == torch.Size([4, 16, 64, 64])
assert mapping['output'].type == OperationDataType.OUTPUT
strategies_vector = handler.register_strategy()
strategies_vector = handler.register_strategy(compute_resharding_cost=False)
strategy_name_list = [val.name for val in strategies_vector]
# SS = SR x RS
@@ -165,7 +165,7 @@ def test_conv_function_handler():
assert mapping['output'].data.shape == torch.Size([4, 16, 64, 64])
assert mapping['output'].type == OperationDataType.OUTPUT
handler.register_strategy()
handler.register_strategy(compute_resharding_cost=False)
strategy_name_list = [val.name for val in strategies_vector]
# SS = SR x RS

4
tests/test_auto_parallel/test_node_handler/test_getitem_handler.py
View File

@@ -47,13 +47,13 @@ def test_getitem_function_handler():
conv_handler = ConvFunctionHandler(node=conv_mod_node,
device_mesh=device_mesh,
strategies_vector=conv_strategies_vector)
conv_handler.register_strategy()
conv_handler.register_strategy(compute_resharding_cost=False)
setattr(conv_mod_node, 'strategies_vector', conv_strategies_vector)
getitem_handler = GetItemHandler(node=getitem_mod_node,
device_mesh=device_mesh,
strategies_vector=getitem_strategies_vector)
getitem_handler.register_strategy()
getitem_handler.register_strategy(compute_resharding_cost=False)
# check operation data mapping
mapping = getitem_handler.get_operation_data_mapping()

2
tests/test_auto_parallel/test_node_handler/test_layer_norm_handler_v2.py
View File

@@ -58,7 +58,7 @@ def test_ln_module_handler():
assert mapping['output'].data.shape == torch.Size([4, 16])
assert mapping['output'].type == OperationDataType.OUTPUT
strategies_vector = handler.register_strategy()
strategies_vector = handler.register_strategy(compute_resharding_cost=False)
strategy_name_list = [val.name for val in strategies_vector]
# SR = SR x R

4
tests/test_auto_parallel/test_node_handler/test_linear_handler_v2.py
View File

@@ -57,7 +57,7 @@ def test_linear_module_handler():
assert mapping['output'].type == OperationDataType.OUTPUT
assert mapping['output'].logical_shape == torch.Size([16, 32])
strategies_vector = handler.register_strategy()
strategies_vector = handler.register_strategy(compute_resharding_cost=False)
strategy_name_list = [val.name for val in strategies_vector]
# one strategy will be converted to different physical sharding spec
assert len(strategy_name_list) > 8
@@ -138,7 +138,7 @@ def test_linear_function_handler():
assert mapping['output'].data.shape == torch.Size([4, 32])
assert mapping['output'].type == OperationDataType.OUTPUT
strategies_vector = handler.register_strategy()
strategies_vector = handler.register_strategy(compute_resharding_cost=False)
strategy_name_list = [val.name for val in strategies_vector]
# one strategy will be converted to different physical sharding spec
assert len(strategy_name_list) > 8

57
tests/test_auto_parallel/test_node_handler/test_output_handler.py Normal file
View File

@@ -0,0 +1,57 @@
import torch
import torch.nn as nn
from colossalai.fx import ColoTracer, ColoGraphModule
from colossalai.auto_parallel.solver.op_handler.output_handler import OuputHandler
from colossalai.auto_parallel.solver.sharding_strategy import OperationData, OperationDataType, StrategiesVector
from colossalai.device.device_mesh import DeviceMesh
class OutputModel(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
y = x * 2
return x, y
def test_output_handler():
model = OutputModel()
tracer = ColoTracer()
# graph():
# %x : torch.Tensor [#users=2] = placeholder[target=x]
# %mul : [#users=1] = call_function[target=operator.mul](args = (%x, 2), kwargs = {})
# return (x, mul)
graph = tracer.trace(model, meta_args={
"x": torch.rand(4, 4, 64, 64).to('meta'),
})
gm = ColoGraphModule(model, graph)
physical_mesh_id = torch.arange(0, 4)
mesh_shape = (2, 2)
device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
output_node = list(graph.nodes)[2]
output_strategies_vector = StrategiesVector(output_node)
# build handler
otuput_handler = OuputHandler(node=output_node, device_mesh=device_mesh, strategies_vector=output_strategies_vector)
otuput_handler.register_strategy(compute_resharding_cost=False)
# check operation data mapping
mapping = otuput_handler.get_operation_data_mapping()
for name, op_data in mapping.items():
op_data: OperationData
# make sure they have valid values
assert op_data.data is not None
assert mapping['output'].name == "output"
assert mapping['output'].data.is_meta
assert mapping['output'].type == OperationDataType.OUTPUT
strategy_name_list = [val.name for val in otuput_handler.strategies_vector]
assert "Replica Output" in strategy_name_list
if __name__ == '__main__':
test_output_handler()

58
tests/test_auto_parallel/test_node_handler/test_placeholder_handler.py Normal file
View File

@@ -0,0 +1,58 @@
import torch
import torch.nn as nn
from colossalai.fx import ColoTracer, ColoGraphModule
from colossalai.auto_parallel.solver.op_handler.placeholder_handler import PlacehodlerHandler
from colossalai.auto_parallel.solver.sharding_strategy import OperationData, OperationDataType, StrategiesVector
from colossalai.device.device_mesh import DeviceMesh
class PlaceholderModel(nn.Module):
def __init__(self):
super().__init__()
def forward(self, input):
return input
def test_placeholder_handler():
model = PlaceholderModel()
tracer = ColoTracer()
# graph():
# %input_1 : torch.Tensor [#users=1] = placeholder[target=input]
# return input_1
graph = tracer.trace(model, meta_args={
"input": torch.rand(4, 4, 64, 64).to('meta'),
})
gm = ColoGraphModule(model, graph)
physical_mesh_id = torch.arange(0, 4)
mesh_shape = (2, 2)
device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
placeholder_node = list(graph.nodes)[0]
placeholder_strategies_vector = StrategiesVector(placeholder_node)
# build handler
placeholder_handler = PlacehodlerHandler(node=placeholder_node,
device_mesh=device_mesh,
strategies_vector=placeholder_strategies_vector)
placeholder_handler.register_strategy(compute_resharding_cost=False)
# check operation data mapping
mapping = placeholder_handler.get_operation_data_mapping()
for name, op_data in mapping.items():
op_data: OperationData
# make sure they have valid values
assert op_data.data is not None
assert mapping['output'].name == "input_1"
assert mapping['output'].data.is_meta
assert mapping['output'].data.shape == torch.Size((4, 4, 64, 64))
assert mapping['output'].type == OperationDataType.OUTPUT
strategy_name_list = [val.name for val in placeholder_handler.strategies_vector]
assert "Replica Placeholder" in strategy_name_list
if __name__ == '__main__':
test_placeholder_handler()

4
tests/test_auto_parallel/test_node_handler/test_reshape_handler_v2.py
View File

@@ -46,13 +46,13 @@ def test_reshape_handler():
conv_handler = ConvFunctionHandler(node=conv_mod_node,
device_mesh=device_mesh,
strategies_vector=conv_strategies_vector)
conv_handler.register_strategy()
conv_handler.register_strategy(compute_resharding_cost=False)
setattr(conv_mod_node, 'strategies_vector', conv_strategies_vector)
reshape_handler = ReshapeHandler(node=reshape_node,
device_mesh=device_mesh,
strategies_vector=reshape_strategies_vector)
reshape_handler.register_strategy()
reshape_handler.register_strategy(compute_resharding_cost=False)
# check operation data mapping
mapping = reshape_handler.get_operation_data_mapping()

4
tests/test_auto_parallel/test_node_handler/test_unary_element_wise_handler_v2.py
View File

@@ -48,13 +48,13 @@ def test_elementwise_handler():
conv_handler = ConvFunctionHandler(node=conv_mod_node,
device_mesh=device_mesh,
strategies_vector=conv_strategies_vector)
conv_handler.register_strategy()
conv_handler.register_strategy(compute_resharding_cost=False)
setattr(conv_mod_node, 'strategies_vector', conv_strategies_vector)
relu_handler = UnaryElementwiseHandler(node=relu_mod_node,
device_mesh=device_mesh,
strategies_vector=relu_strategies_vector)
relu_handler.register_strategy()
relu_handler.register_strategy(compute_resharding_cost=False)
# check operation data mapping
mapping = relu_handler.get_operation_data_mapping()

2
tests/test_auto_parallel/test_node_handler/test_where_handler_v2.py
View File

@@ -75,7 +75,7 @@ def test_where_handler():
assert mapping['output'].data.shape == torch.Size([4, 4, 64, 64])
assert mapping['output'].type == OperationDataType.OUTPUT
handler.register_strategy()
handler.register_strategy(compute_resharding_cost=False)
strategy_name_list = [val.name for val in strategies_vector]
# 4*3 + 4*3/2*2 + 1
assert len(strategy_name_list) == 25

121
tests/test_auto_parallel/test_solver_with_resnet.py
View File

@@ -1,121 +0,0 @@
import torch
from torch.fx import GraphModule
import torch.nn as nn
import pytest
from colossalai.fx.tracer.tracer import ColoTracer
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
from colossalai.auto_parallel.solver.cost_graph import CostGraph
from copy import deepcopy
from colossalai.auto_parallel.solver import Solver
from torchvision.models import resnet34, resnet50
from colossalai.auto_parallel.solver.constants import *
from colossalai.auto_parallel.solver.graph_analysis import GraphAnalyser
from colossalai.auto_parallel.solver.options import SolverOptions
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_cost_graph():
physical_mesh_id = torch.arange(0, 8)
mesh_shape = (2, 4)
# [[0, 1]
# [2, 3]]
device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
shape_consistency_manager = ShapeConsistencyManager()
tracer = ColoTracer()
# model = ConvModel(16, 32)
# input_sample = {'x': torch.rand(4, 16, 64, 64).to('meta')}
model = resnet50(num_classes=100000)
input_sample = {'x': torch.rand(128, 3, 224, 224).to('meta')}
graph = tracer.trace(root=model, meta_args=input_sample)
# graph():
# %x : torch.Tensor [#users=1] = placeholder[target=x]
# %conv1 : [#users=1] = call_module[target=conv1](args = (%x,), kwargs = {})
# %bn1 : [#users=1] = call_module[target=bn1](args = (%conv1,), kwargs = {})
# %relu : [#users=1] = call_module[target=relu](args = (%bn1,), kwargs = {})
# %maxpool : [#users=2] = call_module[target=maxpool](args = (%relu,), kwargs = {})
# %layer1_0_conv1 : [#users=1] = call_module[target=layer1.0.conv1](args = (%maxpool,), kwargs = {})
# %layer1_0_bn1 : [#users=1] = call_module[target=layer1.0.bn1](args = (%layer1_0_conv1,), kwargs = {})
# %layer1_0_relu : [#users=1] = call_module[target=layer1.0.relu](args = (%layer1_0_bn1,), kwargs = {})
# %layer1_0_conv2 : [#users=1] = call_module[target=layer1.0.conv2](args = (%layer1_0_relu,), kwargs = {})
# %layer1_0_bn2 : [#users=1] = call_module[target=layer1.0.bn2](args = (%layer1_0_conv2,), kwargs = {})
# %add : [#users=1] = call_function[target=operator.add](args = (%layer1_0_bn2, %maxpool), kwargs = {})
# %layer1_0_relu_1 : [#users=2] = call_module[target=layer1.0.relu](args = (%add,), kwargs = {})
# %layer1_1_conv1 : [#users=1] = call_module[target=layer1.1.conv1](args = (%layer1_0_relu_1,), kwargs = {})
# %layer1_1_bn1 : [#users=1] = call_module[target=layer1.1.bn1](args = (%layer1_1_conv1,), kwargs = {})
# %layer1_1_relu : [#users=1] = call_module[target=layer1.1.relu](args = (%layer1_1_bn1,), kwargs = {})
# %layer1_1_conv2 : [#users=1] = call_module[target=layer1.1.conv2](args = (%layer1_1_relu,), kwargs = {})
# %layer1_1_bn2 : [#users=1] = call_module[target=layer1.1.bn2](args = (%layer1_1_conv2,), kwargs = {})
# %add_1 : [#users=1] = call_function[target=operator.add](args = (%layer1_1_bn2, %layer1_0_relu_1), kwargs = {})
# ...
# %avgpool : [#users=1] = call_module[target=avgpool](args = (%layer4_2_relu_1,), kwargs = {})
# %flatten : [#users=1] = call_function[target=torch.flatten](args = (%avgpool, 1), kwargs = {})
# %fc : [#users=1] = call_module[target=fc](args = (%flatten,), kwargs = {})
# return fc
gm = GraphModule(model, graph, model.__class__.__name__)
gm.recompile()
graph_analyser = GraphAnalyser(gm)
liveness_list = graph_analyser.liveness_analysis()
solver_options = SolverOptions(fast=True)
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()
solver = Solver(gm.graph, strategies_constructor, cost_graph, graph_analyser)
ret = solver.call_solver_serialized_args()
print(ret[0])
solver._recover_merged_node_strategy()
print(solver.last_s_val)
strategies_list = solver.last_s_val
computation_cost = 0
communication_cost = 0
communication_cost_bn = 0
memory_cost = 0
for index, node in enumerate(graph.nodes):
if node.op == 'call_module':
submod = node.graph.owning_module.get_submodule(node.target)
if type(submod) in BATCHNORM_MODULE_OP:
communication_cost_bn += node.strategies_vector[strategies_list[index]].communication_cost
print(node.name, node.strategies_vector[strategies_list[index]].name)
computation_cost += node.strategies_vector[strategies_list[index]].compute_cost
communication_cost += node.strategies_vector[strategies_list[index]].communication_cost
node_memory_cost = node.strategies_vector[strategies_list[index]].memory_cost
if isinstance(node_memory_cost, tuple):
node_memory_cost = node_memory_cost[0]
memory_cost += node_memory_cost
print(f'computation cost is {computation_cost}')
print(f'communication cost is {communication_cost}')
print(f'memory cost is {memory_cost}')
print(f'bn communication cost is {communication_cost_bn}')
if __name__ == '__main__':
test_cost_graph()