From 629172b31922af9a8bf1373622a64d46006e67ac Mon Sep 17 00:00:00 2001 From: Boyuan Yao <70263930+Cypher30@users.noreply.github.com> Date: Tue, 8 Nov 2022 15:05:26 +0800 Subject: [PATCH] [autoparallel] add batch norm metainfo (#1815) * [fx] metainfo class for auto parallel * [fx] add unit test for linear metainfo * [fx] fix bwd param for linear * [fx] modify unit test * [fx] modify unit test * [fx] modify import * [fx] modify import * [fx] modify import * [fx] move meta profiler to auto parallel * [fx] add conv metainfo class * [fx] restore profiler * [fx] restore meta profiler * [autoparallel] modify unit test * [fx] modify unit test * [autoparallel] add batchnorm metainfo class * [autoparallel] fix batchnorm unit test function declaration * [fx] restore profiler --- .../meta_profiler/meta_registry/__init__.py | 1 + .../meta_profiler/meta_registry/norm.py | 100 ++++++++++++++++++ .../test_metainfo/test_batchnorm_metainfo.py | 61 +++++++++++ 3 files changed, 162 insertions(+) create mode 100644 colossalai/auto_parallel/meta_profiler/meta_registry/norm.py create mode 100644 tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_batchnorm_metainfo.py diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/__init__.py b/colossalai/auto_parallel/meta_profiler/meta_registry/__init__.py index 0763e5167..cbef23da5 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/__init__.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/__init__.py @@ -1,2 +1,3 @@ from .conv import * from .linear import * +from .norm import * diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/norm.py b/colossalai/auto_parallel/meta_profiler/meta_registry/norm.py new file mode 100644 index 000000000..b5818dd87 --- /dev/null +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/norm.py @@ -0,0 +1,100 @@ +from typing import Callable, Dict, List, Tuple, Union + +import torch + +from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( + MemoryCost, + OperationData, + OperationDataType, + ShardingStrategy, + StrategiesVector, + TrainCycleItem, +) +from colossalai.fx.profiler.memory_utils import activation_size +from colossalai.fx.profiler.opcount import flop_mapping +from colossalai.tensor.sharding_spec import ShardingSpec + +from ..registry import meta_register + +__all__ = ['batchnormnd_meta_info'] + + +@meta_register.register(torch.nn.BatchNorm1d) +@meta_register.register(torch.nn.BatchNorm2d) +@meta_register.register(torch.nn.BatchNorm3d) +def batchnormnd_meta_info(*args) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: + """BatchNorm1d, BatchNorm2d, BatchNorm3d, meta info generator + The aten graph of BatchNorm2d is like + + graph(): + %input_2 : [#users=2] = placeholder[target=placeholder](default=) + %cudnn_batch_norm_default : [#users=4] = call_function[target=torch.ops.aten.cudnn_batch_norm.default](args = (%input_2, None, None, None, None, None, None, None), kwargs = {}) + %zeros_like_default : [#users=1] = call_function[target=torch.ops.aten.zeros_like.default](args = (%cudnn_batch_norm_default,), kwargs = {dtype: None, layout: None, device: None, pin_memory: None}) + %detach_default : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%input_2,), kwargs = {}) + %detach_default_1 : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%cudnn_batch_norm_default,), kwargs = {}) + %detach_default_2 : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%cudnn_batch_norm_default,), kwargs = {}) + %detach_default_3 : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%cudnn_batch_norm_default,), kwargs = {}) + %cudnn_batch_norm_backward_default : [#users=3] = call_function[target=torch.ops.aten.cudnn_batch_norm_backward.default](args = (%detach_default, %zeros_like_default, None, None, None, %detach_default_1, %detach_default_2, None, %detach_default_3), kwargs = {}) + %detach_default_4 : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%cudnn_batch_norm_backward_default,), kwargs = {}) + %detach_default_5 : [#users=0] = call_function[target=torch.ops.aten.detach.default](args = (%detach_default_4,), kwargs = {}) + %detach_default_6 : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%cudnn_batch_norm_backward_default,), kwargs = {}) + %detach_default_7 : [#users=0] = call_function[target=torch.ops.aten.detach.default](args = (%detach_default_6,), kwargs = {}) + %detach_default_8 : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%cudnn_batch_norm_backward_default,), kwargs = {}) + %detach_default_9 : [#users=0] = call_function[target=torch.ops.aten.detach.default](args = (%detach_default_8,), kwargs = {}) + Returns: + Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs + """ + + input_tensor = next(filter(lambda x: x.type == OperationDataType.ARG, args)).data + output_tensor = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data + weight_tensor = next(filter(lambda x: x.name == "weight", args)).data + bias_tensor = next(filter(lambda x: x.name == "bias", args)).data + mean_tensor = next(filter(lambda x: x.name == "running_mean", args)).data + var_tensor = next(filter(lambda x: x.name == "running_var", args)).data + num_batch = next(filter(lambda x: x.name == "num_batches_tracked", args)).data + + # construct fwd args + # the fwd inputs are input, weight, bias, running_mean, running_var and some other args + # indicating the status of the module + # the fwd outputs are output, saved mean, saved inv std and num batches tracked + fwd_in_args = [input_tensor, weight_tensor, bias_tensor, mean_tensor, var_tensor, True, 0.1, 1e-5] + fwd_out_args = [output_tensor, mean_tensor, var_tensor, num_batch] + + # construct bwd args + # the bwd inputs are upstream grad, input, weight, running_mean, running_var, saved mean, + # saved inv std and some other args indicating the status of the module + # the bwd outputs are input grad, weight grad and bias grad + bwd_in_args = [ + output_tensor, output_tensor, weight_tensor, mean_tensor, var_tensor, mean_tensor, var_tensor, 1e-5, num_batch + ] + bwd_out_args = [input_tensor, weight_tensor, bias_tensor] + + # calculate cost + fwd_compute_cost = flop_mapping[torch.ops.aten.cudnn_batch_norm.default](fwd_in_args, fwd_out_args) + bwd_compute_cost = flop_mapping[torch.ops.aten.cudnn_batch_norm_backward.default](bwd_in_args, bwd_out_args) + compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost) + + # calculate memory cost + # the fwd activation cost is output plus saved mean and saved inv std + fwd_memory_cost = MemoryCost(activation=activation_size([output_tensor, mean_tensor, var_tensor]), + parameter=activation_size([weight_tensor, bias_tensor]), + temp=0, + buffer=activation_size([mean_tensor, var_tensor])) + + # the bwd memory cost is quite tricky here, BatchNorm will remove saved mean + # and saved inv std during backward phase + bwd_memory_cost = MemoryCost(activation=activation_size([input_tensor]), + parameter=activation_size([weight_tensor, bias_tensor]), + temp=activation_size([mean_tensor, var_tensor]), + buffer=activation_size([mean_tensor, var_tensor])) + + # total cost is the sum of forward and backward cost + total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation, + parameter=fwd_memory_cost.parameter + bwd_memory_cost.parameter) + + memory_cost = TrainCycleItem(fwd=fwd_memory_cost, bwd=bwd_memory_cost, total=total_cost) + + # store fwd_in + fwd_in = [input_tensor] + + return compute_cost, memory_cost, fwd_in diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_batchnorm_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_batchnorm_metainfo.py new file mode 100644 index 000000000..b63d333ba --- /dev/null +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_batchnorm_metainfo.py @@ -0,0 +1,61 @@ +from functools import partial + +import pytest +import torch +import torch.multiprocessing as mp +import torch.nn as nn + +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.pytest_wrapper import run_on_environment_flag +from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use +from colossalai.utils import free_port +from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import mem_test_for_node_strategy + + +def _batchnorm_module_mem_test(rank, world_size, port): + """This function is for conv memory test + Test and print real memory cost and estimated, this test will not be executed except with the tag AUTO_PARALLEL + + Args: + Args: + rank: device rank + bias: indicate whether conv module need bias + world_size: number of devices + port: port for initializing process group + """ + disable_existing_loggers() + launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + model = nn.Sequential(nn.BatchNorm2d(128)).cuda() + input = torch.rand(4, 128, 64, 64).cuda() + input.requires_grad = True + physical_mesh_id = torch.arange(0, 4) + mesh_shape = (2, 2) + device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) + + # index of conv node in computation graph + node_index = 1 + # total number of conv strategies + strategy_number = 4 + mem_test_for_node_strategy(rank=rank, + model=model, + device_mesh=device_mesh, + node_index=node_index, + strategy_number=strategy_number, + input_args=[input], + meta_arg_names=['input']) + + +@run_on_environment_flag(name='AUTO_PARALLEL') +@pytest.mark.dist +@rerun_if_address_is_in_use() +def test_batchnorm_meta_concrete_info_match(): + world_size = 4 + run_func_module = partial(_batchnorm_module_mem_test, world_size=world_size, port=free_port()) + mp.spawn(run_func_module, nprocs=world_size) + + +if __name__ == '__main__': + test_batchnorm_meta_concrete_info_match()