From 44f49803c17e1826262a8b52518380fea5de29f0 Mon Sep 17 00:00:00 2001 From: Cypher30 <1529318642@qq.com> Date: Fri, 17 Feb 2023 10:14:13 +0800 Subject: [PATCH 1/5] [autoparallel] tensor related meta information prototype --- .../meta_profiler/meta_registry/tensor.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py new file mode 100644 index 000000000000..86f083a7470c --- /dev/null +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py @@ -0,0 +1,30 @@ +from typing import List, Tuple + +import torch + +from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem +from colossalai.fx.profiler.memory_utils import activation_size +from colossalai.fx.profiler.opcount import flop_mapping + +from ..registry import meta_register + +__all__ = [] + + +def tensor_related_metainfo(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: + """torch.Tensor related operations metainfo generator + torch.tensor: all zero, fwd_out + torch.Tensor.size: all zero, fwd_out + torch.Tensor.to: all zero, fwd_out + torch.Tensor.type: bwd_mem_out, bwd_mem_tmp, fwd_out + torch.Tensor.contiguous: bwd_mem_out, bwd_mem_tmp, fwd_out + torch.Tensor.transpose: bwd_mem_out, fwd_out + torch.Tensor.permute: bwd_mem_out, fwd_out + torch.Tensor.split: bwd_mem_out, fwd_out + torch.Tensor.view: bwd_mem_out, fwd_out + + + Returns: + Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs + """ + pass From f16f21f7ed6a2cc9b3661f1ec21863e98987b2c5 Mon Sep 17 00:00:00 2001 From: Cypher30 <1529318642@qq.com> Date: Fri, 17 Feb 2023 11:38:33 +0800 Subject: [PATCH 2/5] [autoparallel] tensor related meta information --- .../meta_profiler/meta_registry/__init__.py | 1 + .../meta_profiler/meta_registry/pooling.py | 1 - .../meta_profiler/meta_registry/tensor.py | 78 ++++++++++--- .../test_metainfo/test_tensor_metainfo.py | 103 ++++++++++++++++++ 4 files changed, 167 insertions(+), 16 deletions(-) create mode 100644 tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_tensor_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 359590c1fc04..df9eb6498377 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/__init__.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/__init__.py @@ -5,3 +5,4 @@ from .linear import * from .norm import * from .pooling import * +from .tensor import * diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/pooling.py b/colossalai/auto_parallel/meta_profiler/meta_registry/pooling.py index 79780c92eed4..21272ea09ac1 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/pooling.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/pooling.py @@ -14,7 +14,6 @@ @meta_register.register(torch.nn.AdaptiveAvgPool1d) @meta_register.register(torch.nn.AdaptiveAvgPool2d) @meta_register.register(torch.nn.AdaptiveAvgPool3d) -@meta_register.register(torch.flatten) def avgpool_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: """Meta info for AdaptiveAvgPool The aten graph of AdaptiveAvgPool is diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py index 86f083a7470c..5e67494517b3 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py @@ -1,4 +1,4 @@ -from typing import List, Tuple +from typing import Callable, List, Tuple import torch @@ -8,23 +8,71 @@ from ..registry import meta_register -__all__ = [] +__all__ = ["tensor_related_metainfo"] -def tensor_related_metainfo(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: - """torch.Tensor related operations metainfo generator - torch.tensor: all zero, fwd_out - torch.Tensor.size: all zero, fwd_out - torch.Tensor.to: all zero, fwd_out - torch.Tensor.type: bwd_mem_out, bwd_mem_tmp, fwd_out - torch.Tensor.contiguous: bwd_mem_out, bwd_mem_tmp, fwd_out - torch.Tensor.transpose: bwd_mem_out, fwd_out - torch.Tensor.permute: bwd_mem_out, fwd_out - torch.Tensor.split: bwd_mem_out, fwd_out - torch.Tensor.view: bwd_mem_out, fwd_out +def tensor_related_metainfo(bwd_mem_out_factor: float = 1, bwd_mem_tmp_factor: float = 0) -> Callable: + """torch.Tensor related metainfo generator template + Args: + bwd_mem_out_factor (float, optional): backward activation memory cost factor. Defaults to 1. + bwd_mem_tmp_factor (float, optional): backward temp memory cost factor. Defaults to 0. Returns: - Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs + Callable: torch.Tensor related metainfo generator """ - pass + + def meta_func(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: + """torch.Tensor related metainfo generator + + Returns: + Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs + """ + engaged_tensors = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data + + # compute costs are all zero + compute_cost = TrainCycleItem(fwd=0, bwd=0, total=0) + + # memory costs + # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward + fwd_mem_cost = MemoryCost(activation=activation_size(engaged_tensors) * 2, parameter=0, temp=0, buffer=0) + + bwd_mem_cost = MemoryCost(activation=activation_size(engaged_tensors) * bwd_mem_out_factor, + parameter=0, + temp=activation_size(engaged_tensors) * bwd_mem_tmp_factor, + buffer=0) + + total_mem_cost = MemoryCost(activation=fwd_mem_cost.activation + bwd_mem_cost.activation, + parameter=fwd_mem_cost.parameter + bwd_mem_cost.parameter, + temp=fwd_mem_cost.temp + bwd_mem_cost.temp, + buffer=fwd_mem_cost.buffer + bwd_mem_cost.buffer) + + memory_cost = TrainCycleItem(fwd=fwd_mem_cost, bwd=bwd_mem_cost, total=total_mem_cost) + + # store fwd_in, fwd_buffer, fwd_out + fwd_in = [] + fwd_buffer = [] + if isinstance(engaged_tensors, tuple) or isinstance(engaged_tensors, list) or isinstance(engaged_tensors, dict): + fwd_out = [torch.zeros_like(tensor) for tensor in engaged_tensors] + else: + # enaged_tensors is a single tensor + fwd_out = [torch.zeros_like(engaged_tensors)] + + return compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out + + return meta_func + + +# register torch.Tensor related metainfo +# (0, 0) +meta_register.register([torch.tensor, torch.Tensor.size, torch.Tensor.to, torch.Tensor.unsqueeze, + torch.unsqueeze])(tensor_related_metainfo(0, 0)) + +# (1, 0) +meta_register.register([ + torch.Tensor.flatten, torch.flatten, torch.Tensor.transpose, torch.transpose, torch.Tensor.permute, torch.permute, + torch.Tensor.split, torch.split, torch.Tensor.view +])(tensor_related_metainfo(1, 0)) + +# (1, 1) +meta_register.register([torch.Tensor.type, torch.Tensor.contiguous])(tensor_related_metainfo(1, 1)) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_tensor_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_tensor_metainfo.py new file mode 100644 index 000000000000..a0ab66fdc060 --- /dev/null +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_tensor_metainfo.py @@ -0,0 +1,103 @@ +from functools import partial + +import pytest +import torch +import torch.multiprocessing as mp +import torch.nn as nn + +from colossalai.auto_parallel.tensor_shard.node_handler import LinearModuleHandler +from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( + MemoryCost, + OperationData, + OperationDataType, + ShardingStrategy, + StrategiesVector, + TrainCycleItem, +) +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 print_results + +if torch.__version__ >= '1.12.0': + from colossalai.auto_parallel.meta_profiler import MetaInfo, meta_register + + +class SplitModule(nn.Module): + + def __init__(self) -> None: + super().__init__() + + def forward(self, x): + return x.split(512, dim=0) + + +@pytest.mark.skipif(torch.__version__ < '1.12.0', reason="need pytorch 1.12.0 or higher for aten level operations") +def test_tensor_meta_info(): + """test tensor related meta information + We will just use torch.Tensor.split for the test + """ + meta_func = meta_register.get(torch.Tensor.split) + + # construct meta tensors + input_tensor = torch.rand(1024, 1024, device="meta") + output_tensor = input_tensor.split(512, dim=0) + + # construct operation data + input_data = OperationData( + name="input", + data=input_tensor, + type=OperationDataType.ARG, + logical_shape=input_tensor.shape, + ) + output_data = OperationData( + name="output", + data=output_tensor, + type=OperationDataType.OUTPUT, + logical_shape=input_tensor.shape, + ) + split_info_data = OperationData( + name='split_info', + type=OperationDataType.ARG, + data=0, + logical_shape=None, + ) + + # construct args + args = [input_data, output_data, split_info_data] + kwargs = {'inplace': False} + + # estimated results + compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out = meta_func(*args, **kwargs) + + # actual results + model = SplitModule() + input_real_tensor = torch.rand(1024, 1024).cuda() + + input_real_tensor.requires_grad = True + + # fwd + torch.cuda.reset_peak_memory_stats() + mem_stamp0 = torch.cuda.memory_allocated() + output_real_tensor = model(input_real_tensor) + fwd_allocated = torch.cuda.memory_allocated() - mem_stamp0 + fwd_peak = torch.cuda.max_memory_allocated() - mem_stamp0 + + # bwd + upstream_grad = [torch.rand_like(tensor) for tensor in output_real_tensor] + torch.cuda.reset_peak_memory_stats() + mem_stamp0 = torch.cuda.memory_allocated() + torch.autograd.backward(output_real_tensor, upstream_grad) + bwd_allocated = torch.cuda.memory_allocated() - mem_stamp0 + bwd_peak = torch.cuda.max_memory_allocated() - mem_stamp0 + + print_results([input_real_tensor], output_real_tensor, compute_cost, memory_cost, fwd_allocated, fwd_peak, + bwd_allocated, bwd_peak) + + +if __name__ == "__main__": + test_tensor_meta_info() From 08d306d99f52796df5b48e274d45233bbb9d0c51 Mon Sep 17 00:00:00 2001 From: Cypher30 <1529318642@qq.com> Date: Fri, 17 Feb 2023 15:38:48 +0800 Subject: [PATCH 3/5] [autoparallel] tensor related meta information --- .../meta_profiler/meta_registry/tensor.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py index 5e67494517b3..618e17ab0b0f 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py @@ -28,18 +28,25 @@ def meta_func(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[tor Returns: Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs """ - engaged_tensors = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data + inputs = next( + filter( + lambda x: x.type == OperationDataType.ARG or x.type == OperationDataType.PARAM and isinstance( + x.data, torch.Tensor), args)).data + outputs = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data # compute costs are all zero compute_cost = TrainCycleItem(fwd=0, bwd=0, total=0) # memory costs # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward - fwd_mem_cost = MemoryCost(activation=activation_size(engaged_tensors) * 2, parameter=0, temp=0, buffer=0) + fwd_mem_cost = MemoryCost(activation=activation_size(inputs) + activation_size(outputs), + parameter=0, + temp=0, + buffer=0) - bwd_mem_cost = MemoryCost(activation=activation_size(engaged_tensors) * bwd_mem_out_factor, + bwd_mem_cost = MemoryCost(activation=activation_size(outputs) * bwd_mem_out_factor, parameter=0, - temp=activation_size(engaged_tensors) * bwd_mem_tmp_factor, + temp=activation_size(outputs) * bwd_mem_tmp_factor, buffer=0) total_mem_cost = MemoryCost(activation=fwd_mem_cost.activation + bwd_mem_cost.activation, @@ -52,11 +59,14 @@ def meta_func(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[tor # store fwd_in, fwd_buffer, fwd_out fwd_in = [] fwd_buffer = [] - if isinstance(engaged_tensors, tuple) or isinstance(engaged_tensors, list) or isinstance(engaged_tensors, dict): - fwd_out = [torch.zeros_like(tensor) for tensor in engaged_tensors] + if isinstance(outputs, tuple) or isinstance(outputs, list) or isinstance(outputs, dict): + if isinstance(outputs[0], torch.Tensor): + fwd_out = [torch.zeros_like(tensor) for tensor in outputs] + else: + fwd_out = outputs else: # enaged_tensors is a single tensor - fwd_out = [torch.zeros_like(engaged_tensors)] + fwd_out = [torch.zeros_like(outputs)] return compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out From 526ca50ec4c6ef7b9860b7ec8cbb07785e2ac67d Mon Sep 17 00:00:00 2001 From: Cypher30 <1529318642@qq.com> Date: Sat, 18 Feb 2023 10:21:49 +0800 Subject: [PATCH 4/5] [autoparallel] tensor related meta information --- .../meta_profiler/meta_registry/tensor.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py index 618e17ab0b0f..154a94217a68 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py @@ -28,10 +28,6 @@ def meta_func(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[tor Returns: Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs """ - inputs = next( - filter( - lambda x: x.type == OperationDataType.ARG or x.type == OperationDataType.PARAM and isinstance( - x.data, torch.Tensor), args)).data outputs = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data # compute costs are all zero @@ -39,10 +35,7 @@ def meta_func(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[tor # memory costs # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward - fwd_mem_cost = MemoryCost(activation=activation_size(inputs) + activation_size(outputs), - parameter=0, - temp=0, - buffer=0) + fwd_mem_cost = MemoryCost(activation=activation_size(outputs) * 2, parameter=0, temp=0, buffer=0) bwd_mem_cost = MemoryCost(activation=activation_size(outputs) * bwd_mem_out_factor, parameter=0, @@ -75,8 +68,8 @@ def meta_func(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[tor # register torch.Tensor related metainfo # (0, 0) -meta_register.register([torch.tensor, torch.Tensor.size, torch.Tensor.to, torch.Tensor.unsqueeze, - torch.unsqueeze])(tensor_related_metainfo(0, 0)) +meta_register.register([torch.Tensor.to, torch.Tensor.unsqueeze, torch.unsqueeze, + torch.arange])(tensor_related_metainfo(0, 0)) # (1, 0) meta_register.register([ From 0bffb722f8f611dfd2c278c7c85f1f9b63bb6c8d Mon Sep 17 00:00:00 2001 From: Cypher30 <1529318642@qq.com> Date: Sat, 18 Feb 2023 10:46:08 +0800 Subject: [PATCH 5/5] [autoparallel] tensor related meta information --- .../auto_parallel/meta_profiler/meta_registry/tensor.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py index 154a94217a68..332e649d2d7e 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py @@ -53,10 +53,8 @@ def meta_func(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[tor fwd_in = [] fwd_buffer = [] if isinstance(outputs, tuple) or isinstance(outputs, list) or isinstance(outputs, dict): - if isinstance(outputs[0], torch.Tensor): - fwd_out = [torch.zeros_like(tensor) for tensor in outputs] - else: - fwd_out = outputs + # tuple of tensors + fwd_out = [torch.zeros_like(tensor) for tensor in outputs] else: # enaged_tensors is a single tensor fwd_out = [torch.zeros_like(outputs)] @@ -68,7 +66,7 @@ def meta_func(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[tor # register torch.Tensor related metainfo # (0, 0) -meta_register.register([torch.Tensor.to, torch.Tensor.unsqueeze, torch.unsqueeze, +meta_register.register([torch.tensor, torch.Tensor.to, torch.Tensor.unsqueeze, torch.unsqueeze, torch.arange])(tensor_related_metainfo(0, 0)) # (1, 0)