From d7ecaf362b6c20412838ade71446dcc2d9d05f6c Mon Sep 17 00:00:00 2001 From: HELSON Date: Thu, 7 Apr 2022 17:38:45 +0800 Subject: [PATCH] [zero] fix init bugs in zero context (#686) * adapt model weight initialization for methods in Pytorch nn.init --- colossalai/zero/init_ctx/init_context.py | 125 ++++++++++++------ tests/test_moe/test_moe_zero_init.py | 49 ++++--- tests/test_moe/test_moe_zero_model.py | 7 +- tests/test_moe/test_moe_zero_optim.py | 3 +- .../test_init_context.py | 4 +- .../test_shard_model_v2.py | 7 +- .../test_sharded_optim_v2.py | 3 +- .../test_state_dict.py | 5 +- 8 files changed, 117 insertions(+), 86 deletions(-) diff --git a/colossalai/zero/init_ctx/init_context.py b/colossalai/zero/init_ctx/init_context.py index 52a166d89..f5efe3d11 100644 --- a/colossalai/zero/init_ctx/init_context.py +++ b/colossalai/zero/init_ctx/init_context.py @@ -3,6 +3,8 @@ import functools from typing import Optional import torch +import torch.nn as nn +import torch.distributed as dist from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc from colossalai.context.singleton_meta import SingletonMeta @@ -10,7 +12,6 @@ from colossalai.logging import get_dist_logger from colossalai.zero.shard_utils import BaseShardStrategy from colossalai.zero.sharded_model._utils import cast_tensor_to_fp16 from colossalai.zero.sharded_param import ShardedParamV2 -from torch.distributed import ProcessGroup from contextlib import AbstractContextManager @@ -93,24 +94,21 @@ class ZeroContextConfig(object): replicated (bool, optional): Whether the param is replicated across data parallel group. Some parameters are not replicated, e.g. parameters in MOE experts. shard_param (bool, optional): Is param sharded after exiting the context. Defaults to False. - rm_torch_payload_on_the_fly (bool, optional): If set to `True`, remove tensor payload on `param.data` after module init finished. - This will reduce memory usage when initializing model. - But it's not suitable for all models, especially when there are `weight init` operations in `__init__`. - If set to `False`, remove tensor payload on param.data afther the context exist. - This is used when you add some logic to operate tensors in __init__ of module. - See torchvision resnet18. Defaults to False. """ - def __init__(self, - target_device: torch.device, - replicated: bool = True, - shard_param: bool = False, - rm_torch_payload_on_the_fly: bool = False): + def __init__(self, target_device: torch.device, replicated: bool = True, shard_param: bool = False): super().__init__() + + if shard_param: + assert replicated, "Non-replicated parameters can't be sharded." + + # replicated no-shard parameters should locate in cuda, since we will broadcast them soon + if replicated and not shard_param: + assert target_device.type == 'cuda', "Replicated no-shard paramters should locate in cuda." + self.target_device = target_device self.is_replicated: bool = replicated self.shard_param: bool = shard_param - self.rm_torch_payload_on_the_fly: bool = rm_torch_payload_on_the_fly class ZeroInitContext(InsertPostInitMethodToModuleSubClasses): @@ -123,35 +121,27 @@ class ZeroInitContext(InsertPostInitMethodToModuleSubClasses): Args: target_device (torch.device): The device where param data are after exiting the context. shard_strategy (BaseShardStrategy): Shard strategy instance. + seed (int, optional): Random seed for weight initialization shard_param (bool, optional): Is param sharded after exiting the context. Defaults to False. - rm_torch_payload_on_the_fly (bool, optional): If set to `True`, remove tensor payload on `param.data` after module init finished. - This will reduce memory usage when initializing model. - But it's not suitable for all models, especially when there are `weight init` operations in `__init__`. - If set to `False`, remove tensor payload on param.data afther the context exist. - This is used when you add some logic to operate tensors in __init__ of module. - See torchvision resnet18. Defaults to False. model_numel_tensor (torch.Tensor, optional): A tensor which will store the number of elements of model. Defaults to torch.zeros(1, dtype=torch.int). - dp_process_group (Optional[ProcessGroup], optional): Data parallel process group. Defaults to None. """ def __init__(self, target_device: torch.device, shard_strategy: BaseShardStrategy, + seed: int = 2**10 - 1, shard_param: bool = False, - rm_torch_payload_on_the_fly: bool = False, - model_numel_tensor: torch.Tensor = torch.zeros(1, dtype=torch.long), - dp_process_group: Optional[ProcessGroup] = None): + model_numel_tensor: torch.Tensor = torch.zeros(1, dtype=torch.long)): super().__init__() self.shard_strategy = shard_strategy - self.initialized_param_list = [] + self.sharded_param_list = [] + self.unshard_param_list = [] self.model_numel_tensor = model_numel_tensor - self.dp_process_group = dp_process_group or gpc.get_group(ParallelMode.DATA) + self.seed = seed + self.dp_process_group = gpc.get_group(ParallelMode.DATA) - self.config = ZeroContextConfig(target_device=target_device, - replicated=True, - shard_param=shard_param, - rm_torch_payload_on_the_fly=rm_torch_payload_on_the_fly) + self.config = ZeroContextConfig(target_device=target_device, replicated=True, shard_param=shard_param) ZeroContextMgr().current_context = self @@ -167,9 +157,35 @@ class ZeroInitContext(InsertPostInitMethodToModuleSubClasses): def shard_param(self): return self.config.shard_param - @property - def rm_torch_payload_on_the_fly(self): - return self.config.rm_torch_payload_on_the_fly + @staticmethod + def calc_fanin_fanout(tensor: torch.Tensor): + """We use this function to substitute fan-in and fan-out calculation in torch.nn.init. + This can help us get correct fan-in and fan-out for sharded tensor. + """ + assert isinstance(tensor, nn.Parameter), "Sharded tensor initilization is only allowed for paramters" + + # get correct shape of input tensor + if not hasattr(tensor, 'colo_attr') or not tensor.colo_attr.param_is_sharded: + tensor_shape = tensor.shape + else: + tensor_shape = tensor.colo_attr.sharded_data_tensor.origin_shape + + dimensions = len(tensor_shape) + if dimensions < 2: + raise ValueError("Fan in and fan out can not be computed for tensor with fewer than 2 dimensions") + + num_input_fmaps = tensor_shape[1] + num_output_fmaps = tensor_shape[0] + receptive_field_size = 1 + if dimensions > 2: + # math.prod is not always available, accumulate the product manually + # we could use functools.reduce but that is not supported by TorchScript + for s in tensor_shape[2:]: + receptive_field_size *= s + fan_in = num_input_fmaps * receptive_field_size + fan_out = num_output_fmaps * receptive_field_size + + return fan_in, fan_out def _pre_context_exec(self): """ @@ -177,15 +193,40 @@ class ZeroInitContext(InsertPostInitMethodToModuleSubClasses): """ self.logger = get_dist_logger("ZeroInitContext") + # substitute fan-in and fan-out calculation + self.nn_fanin_fanout = nn.init._calculate_fan_in_and_fan_out + nn.init._calculate_fan_in_and_fan_out = self.calc_fanin_fanout + + # reserve rng states + self.cpu_rng_state = torch.get_rng_state() + self.cuda_rng_state = torch.cuda.get_rng_state() + + # set new seed for initialization, since we initialize sharded tensor separately + # we don't want all processes have the same seed + # otherwise all sharded tensors are same after init + offset = self.seed + 1 # we want to have more 1 in binary format seed + torch.manual_seed(self.seed + offset * dist.get_rank()) + def _post_context_exec(self): """The callback function when exiting context. """ - if not self.rm_torch_payload_on_the_fly: - for param in self.initialized_param_list: - assert hasattr(param, 'colo_attr') - param.colo_attr.remove_torch_payload() + for param in self.sharded_param_list: + assert hasattr(param, 'colo_attr') + param.colo_attr.remove_torch_payload() - del self.initialized_param_list + del self.sharded_param_list + + src_rank = gpc.get_ranks_in_group(ParallelMode.DATA)[0] + for param in self.unshard_param_list: + assert hasattr(param, 'colo_attr') + if param.is_replicated: + dist.broadcast(tensor=param.data, src=src_rank, group=self.dp_process_group) + + del self.unshard_param_list + + nn.init._calculate_fan_in_and_fan_out = self.nn_fanin_fanout + torch.set_rng_state(self.cpu_rng_state) + torch.cuda.set_rng_state(self.cuda_rng_state) def _post_init_method(self, module: torch.nn.Module): """ @@ -219,11 +260,14 @@ class ZeroInitContext(InsertPostInitMethodToModuleSubClasses): if param.grad is not None: param.grad = param.grad.to(target_device) - param.colo_attr = ShardedParamV2(param, rm_torch_payload=self.rm_torch_payload_on_the_fly) + param.colo_attr = ShardedParamV2(param, rm_torch_payload=False) if self.shard_param: self.shard_strategy.shard([param.colo_attr.sharded_data_tensor], self.dp_process_group) - self.initialized_param_list.append(param) + param.data = param.colo_attr.sharded_data_tensor.payload + self.sharded_param_list.append(param) + else: + self.unshard_param_list.append(param) # We must cast buffers # If we use BN, buffers may be on CPU and Float @@ -250,8 +294,7 @@ class ZeroContextMgr(metaclass=SingletonMeta): def no_shard_zero_context(is_replicated: bool = True) -> AbstractContextManager: return ZeroContextMgr().hijack_context_config(target_device=torch.device('cuda', torch.cuda.current_device()), replicated=is_replicated, - shard_param=False, - rm_torch_payload_on_the_fly=False) + shard_param=False) def no_shard_zero_decrator(is_replicated: bool = True): diff --git a/tests/test_moe/test_moe_zero_init.py b/tests/test_moe/test_moe_zero_init.py index fca39b83a..3c308a421 100644 --- a/tests/test_moe/test_moe_zero_init.py +++ b/tests/test_moe/test_moe_zero_init.py @@ -51,36 +51,36 @@ def run_moe_zero_init(init_device_type, shard_strategy_class): with ZeroInitContext(target_device=init_device, shard_strategy=shard_strategy_class(), shard_param=True, - model_numel_tensor=model_numel_tensor, - rm_torch_payload_on_the_fly=False): + model_numel_tensor=model_numel_tensor): model = MoeModel() - for name, param in model.named_parameters(): - assert hasattr(param, 'colo_attr') + for name, param in model.named_parameters(): + assert hasattr(param, 'colo_attr') - # the weights in the gate should be fp32 - if 'gate' in name: - assert param.colo_attr.sharded_data_tensor.dtype == torch.float32 - else: - assert param.colo_attr.sharded_data_tensor.dtype == torch.half + # the weights in the gate should be fp32 + if 'gate' in name: + assert param.colo_attr.sharded_data_tensor.dtype == torch.float32 + else: + assert param.colo_attr.sharded_data_tensor.dtype == torch.half - # the parameters in moe experts and its gate should not be sharded - if ('experts' in name) or ('gate' in name) or ('residual_combine' in name): - assert not param.colo_attr.sharded_data_tensor.is_sharded - else: - assert param.colo_attr.sharded_data_tensor.is_sharded + # the parameters in moe experts and its gate should not be sharded + if ('experts' in name) or ('gate' in name) or ('residual_combine' in name): + assert not param.colo_attr.sharded_data_tensor.is_sharded + assert param.colo_attr.sharded_data_tensor.data_ptr() == param.data.data_ptr() + else: + assert param.colo_attr.sharded_data_tensor.is_sharded - # the parameters in moe experts is not replicated - if 'experts' in name: - assert not param.is_replicated - else: - assert param.is_replicated + # the parameters in moe experts is not replicated + if 'experts' in name: + assert not param.is_replicated + else: + assert param.is_replicated - if param.colo_attr.param_is_sharded: - assert param.colo_attr.sharded_data_tensor.payload.device.type == init_device.type, \ - f'{param.colo_attr.sharded_data_tensor.payload.device.type} vs. {init_device.type}' - else: - assert param.colo_attr.sharded_data_tensor.payload.device.type == 'cuda' + if param.colo_attr.param_is_sharded: + assert param.colo_attr.sharded_data_tensor.payload.device.type == init_device.type, \ + f'{param.colo_attr.sharded_data_tensor.payload.device.type} vs. {init_device.type}' + else: + assert param.colo_attr.sharded_data_tensor.payload.device.type == 'cuda' def _run_dist(rank, world_size, port): @@ -91,7 +91,6 @@ def _run_dist(rank, world_size, port): @pytest.mark.dist @pytest.mark.parametrize("world_size", [2, 4]) -@pytest.mark.skip("Under development") @rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*") def test_moe_zero_init(world_size): run_func = partial(_run_dist, world_size=world_size, port=free_port()) diff --git a/tests/test_moe/test_moe_zero_model.py b/tests/test_moe/test_moe_zero_model.py index 87a72a8e1..34c43bc06 100644 --- a/tests/test_moe/test_moe_zero_model.py +++ b/tests/test_moe/test_moe_zero_model.py @@ -28,12 +28,9 @@ def run_model_test(enable_autocast, shard_strategy_class): get_components_func = non_distributed_component_funcs.get_callable('no_leaf_module') _, train_dataloader, _, _, criterion = get_components_func() - rm_torch_payload_on_the_fly = False - - with ZeroInitContext(target_device=torch.cuda.current_device(), + with ZeroInitContext(target_device=torch.device('cuda', torch.cuda.current_device()), shard_strategy=shard_strategy, - shard_param=True, - rm_torch_payload_on_the_fly=rm_torch_payload_on_the_fly): + shard_param=True): zero_model = MoeModel() zero_model = ShardedModelV2(zero_model, shard_strategy, use_memory_tracer=True) diff --git a/tests/test_moe/test_moe_zero_optim.py b/tests/test_moe/test_moe_zero_optim.py index 5281f92f1..5a5063d6c 100644 --- a/tests/test_moe/test_moe_zero_optim.py +++ b/tests/test_moe/test_moe_zero_optim.py @@ -60,8 +60,7 @@ def _run_test_sharded_optim_v2(cpu_offload, shard_strategy_class, use_cpuadam, g with ZeroInitContext( target_device=torch.device('cpu') if cpu_offload else torch.device(f'cuda:{get_current_device()}'), shard_strategy=shard_strategy, - shard_param=True, - rm_torch_payload_on_the_fly=False): + shard_param=True): zero_model = MoeModel() zero_model = ShardedModelV2( diff --git a/tests/test_zero_data_parallel/test_init_context.py b/tests/test_zero_data_parallel/test_init_context.py index ac1eaa7af..094227b6d 100644 --- a/tests/test_zero_data_parallel/test_init_context.py +++ b/tests/test_zero_data_parallel/test_init_context.py @@ -28,7 +28,6 @@ def run_model_test(init_device_type, shard_strategy_class): for get_components_func in non_distributed_component_funcs: model_builder, _, _, _, _ = get_components_func() - model_numel_tensor = torch.zeros(1, dtype=torch.int) if init_device_type == 'cuda': init_device = torch.device(f"cuda:{get_current_device()}") elif init_device_type == 'cpu': @@ -40,8 +39,7 @@ def run_model_test(init_device_type, shard_strategy_class): with ZeroInitContext(target_device=init_device, shard_strategy=shard_strategy_class(), shard_param=True, - model_numel_tensor=model_numel_tensor, - rm_torch_payload_on_the_fly=False): + model_numel_tensor=model_numel_tensor): model = model_builder(checkpoint=True) for param in model.parameters(): diff --git a/tests/test_zero_data_parallel/test_shard_model_v2.py b/tests/test_zero_data_parallel/test_shard_model_v2.py index fa889bcde..449ceedc0 100644 --- a/tests/test_zero_data_parallel/test_shard_model_v2.py +++ b/tests/test_zero_data_parallel/test_shard_model_v2.py @@ -29,12 +29,9 @@ def run_model_test(enable_autocast, shard_strategy_class): get_components_func = non_distributed_component_funcs.get_callable(model_name) model_builder, train_dataloader, _, _, criterion = get_components_func() - rm_torch_payload_on_the_fly = False - - with ZeroInitContext(target_device=torch.cuda.current_device(), + with ZeroInitContext(target_device=torch.device('cuda', torch.cuda.current_device()), shard_strategy=shard_strategy, - shard_param=True, - rm_torch_payload_on_the_fly=rm_torch_payload_on_the_fly): + shard_param=True): zero_model = model_builder(checkpoint=True) zero_model = ShardedModelV2(zero_model, shard_strategy, use_memory_tracer=True) diff --git a/tests/test_zero_data_parallel/test_sharded_optim_v2.py b/tests/test_zero_data_parallel/test_sharded_optim_v2.py index 40bad3075..01339cc34 100644 --- a/tests/test_zero_data_parallel/test_sharded_optim_v2.py +++ b/tests/test_zero_data_parallel/test_sharded_optim_v2.py @@ -60,8 +60,7 @@ def _run_test_sharded_optim_v2(cpu_offload, shard_strategy_class, use_cpuadam, g with ZeroInitContext( target_device=torch.device(f'cpu:0') if cpu_offload else torch.device(f'cuda:{get_current_device()}'), shard_strategy=shard_strategy, - shard_param=True, - rm_torch_payload_on_the_fly=False): + shard_param=True): zero_model = model_builder(checkpoint=True) zero_model = ShardedModelV2( zero_model, diff --git a/tests/test_zero_data_parallel/test_state_dict.py b/tests/test_zero_data_parallel/test_state_dict.py index fd1cc2f07..41cae05a5 100644 --- a/tests/test_zero_data_parallel/test_state_dict.py +++ b/tests/test_zero_data_parallel/test_state_dict.py @@ -27,10 +27,9 @@ def run_zero_state_dict(shard_strategy_class): get_components_func = non_distributed_component_funcs.get_callable(model_name) model_builder, train_dataloader, test_dataloader, optimizer, criterion = get_components_func() - with ZeroInitContext(target_device=torch.cuda.current_device(), + with ZeroInitContext(target_device=torch.device('cuda', torch.cuda.current_device()), shard_strategy=shard_strategy, - shard_param=True, - rm_torch_payload_on_the_fly=False): + shard_param=True): zero_model = model_builder(checkpoint=True) zero_model = ShardedModelV2(zero_model, shard_strategy)