From f16f718bf8f517e7418ed95feaa762992e493d0c Mon Sep 17 00:00:00 2001 From: hyunwoongko Date: Sun, 26 Sep 2021 04:22:58 +0900 Subject: [PATCH 1/3] Add flexibility of pipeline module and engine --- .../runtime/activation_checkpointing/checkpointing.py | 8 +++++++- deepspeed/runtime/pipe/engine.py | 11 ++++++----- deepspeed/runtime/pipe/module.py | 9 ++++++++- 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index e1bbf1913041..ff07dffb5cc8 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -365,11 +365,14 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): global contiguous_data_buffers, data_offsets inputs = [] + num_non_fp_tensors = 0 for i, item in enumerate(args): if not is_activation_to_checkpoint(item): inputs.append(item) + num_non_fp_tensors += 1 continue + i -= num_non_fp_tensors partition_size = get_partition_size(item) partition = item.detach().contiguous().view(-1).narrow( 0, @@ -427,15 +430,18 @@ def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint global contiguous_size_buffers, size_offsets new_args = [] + num_non_fp_tensors = 0 for i, (arg, inp) in enumerate(zip(args, inputs)): size = torch.tensor(arg.size()) if torch.is_tensor(arg) else None if not is_activation_to_checkpoint(arg): new_args.append(arg) new_args.append(size) + num_non_fp_tensors += 1 continue arg.data = inp.data new_args.append(arg) + i -= num_non_fp_tensors if contiguous_checkpoint: numel = size.numel() @@ -893,4 +899,4 @@ def is_configured(): Return: True of configured, else False """ - return deepspeed_checkpointing_enabled + return deepspeed_checkpointing_enabled \ No newline at end of file diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py index 196cbe8c6217..74aea0fb5858 100644 --- a/deepspeed/runtime/pipe/engine.py +++ b/deepspeed/runtime/pipe/engine.py @@ -49,7 +49,7 @@ class PipelineEngine(DeepSpeedEngine): This engine is created by ``deepspeed.initialize()`` when a :class:`PipelineModule` is provided. """ - def __init__(self, *super_args, **super_kwargs): + def __init__(self, has_bool_tensors=False, *super_args, **super_kwargs): super().__init__(*super_args, **super_kwargs) assert isinstance(self.module, PipelineModule), "model must base PipelineModule" @@ -57,6 +57,7 @@ def __init__(self, *super_args, **super_kwargs): # We schedule the all-reduces, so disable it in super().backward() self.enable_backward_allreduce = False + self.has_bool_tensors = has_bool_tensors # used to disable the pipeline all-reduce when used with 1-bit Adam/1-bit LAMB self.pipeline_enable_backward_allreduce = True @@ -837,7 +838,7 @@ def _exec_send_activations(self, buffer_id): # NCCL does not like to send torch.BoolTensor types, so cast the mask to half(). # We could do char, but with half() we can eventually flatten with other fp16 # messages (TODO) - if self.module.__class__.__name__ == 'GPT2ModelPipe': + if self.module.__class__.__name__ == 'GPT2ModelPipe' or self.has_bool_tensors: outputs = list(outputs) outputs[-1] = outputs[-1].half() outputs = tuple(outputs) @@ -856,7 +857,7 @@ def _exec_send_activations(self, buffer_id): f'{type(outputs)}') # Restore the boolean tensor - if self.module.__class__.__name__ == 'GPT2ModelPipe': + if self.module.__class__.__name__ == 'GPT2ModelPipe' or self.has_bool_tensors: outputs = list(outputs) outputs[-1] = outputs[-1].bool() outputs = tuple(outputs) @@ -885,7 +886,7 @@ def _exec_send_grads(self, buffer_id): # a grad that needs to be communicated. We free the buffer immediately # after, so no need to restore it. The receiver also has a hack that skips # the recv. This is because NCCL does not let us send torch.BoolTensor :-(. - if self.module.__class__.__name__ == 'GPT2ModelPipe': + if self.module.__class__.__name__ == 'GPT2ModelPipe' or self.has_bool_tensors: inputs = list(inputs) inputs.pop() inputs = tuple(inputs) @@ -948,7 +949,7 @@ def _exec_recv_activations(self, buffer_id): # NCCL does not like to send torch.BoolTensor types, so un-cast the # attention mask - if self.module.__class__.__name__ == 'GPT2ModelPipe': + if self.module.__class__.__name__ == 'GPT2ModelPipe' or self.has_bool_tensors: recvd[-1] = recvd[-1].bool() recvd = tuple(recvd) diff --git a/deepspeed/runtime/pipe/module.py b/deepspeed/runtime/pipe/module.py index 815c8eb2160b..8126317d93cd 100644 --- a/deepspeed/runtime/pipe/module.py +++ b/deepspeed/runtime/pipe/module.py @@ -95,7 +95,8 @@ def __init__(self, base_seed=1234, partition_method='parameters', activation_checkpoint_interval=0, - activation_checkpoint_func=checkpointing.checkpoint): + activation_checkpoint_func=checkpointing.checkpoint, + checkpointable_layers=None): """Modules to be parallelized with pipeline parallelism. The key constraint that enables pipeline parallelism is the @@ -137,6 +138,10 @@ def forward(self, inputs): self.loss_fn = loss_fn + self.checkpointable_layers = checkpointable_layers + if checkpointable_layers is not None: + assert isinstance(checkpointable_layers, list), "param `checkpointable_layers` must be type of list." + self.seed_layers = seed_layers self.seed_fn = seed_fn self.base_seed = base_seed @@ -599,6 +604,8 @@ def _is_checkpointable(self, funcs): if self.__class__.__name__ == 'GPT2ModelPipe': return all('ParallelTransformerLayerPipe' in f.__class__.__name__ for f in funcs) + if self.checkpointable_layers is not None: + return all(f.__class__.__name__ in self.checkpointable_layers for f in funcs) params = [f.parameters() for f in funcs if isinstance(f, torch.nn.Module)] return any(len(list(p)) > 0 for p in params) From 75ba9741b95bd9b25e232e88e0945bf1d2133bcf Mon Sep 17 00:00:00 2001 From: hyunwoongko Date: Tue, 28 Sep 2021 03:53:10 +0900 Subject: [PATCH 2/3] Separate PRs --- deepspeed/runtime/activation_checkpointing/checkpointing.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index ff07dffb5cc8..6acefeb249e4 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -365,14 +365,11 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): global contiguous_data_buffers, data_offsets inputs = [] - num_non_fp_tensors = 0 for i, item in enumerate(args): if not is_activation_to_checkpoint(item): inputs.append(item) - num_non_fp_tensors += 1 continue - i -= num_non_fp_tensors partition_size = get_partition_size(item) partition = item.detach().contiguous().view(-1).narrow( 0, @@ -430,18 +427,15 @@ def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint global contiguous_size_buffers, size_offsets new_args = [] - num_non_fp_tensors = 0 for i, (arg, inp) in enumerate(zip(args, inputs)): size = torch.tensor(arg.size()) if torch.is_tensor(arg) else None if not is_activation_to_checkpoint(arg): new_args.append(arg) new_args.append(size) - num_non_fp_tensors += 1 continue arg.data = inp.data new_args.append(arg) - i -= num_non_fp_tensors if contiguous_checkpoint: numel = size.numel() From 2c86eaa1ce5dd178f7e632fc36700310f66120fb Mon Sep 17 00:00:00 2001 From: hyunwoongko Date: Tue, 28 Sep 2021 03:54:26 +0900 Subject: [PATCH 3/3] Separate PRs --- deepspeed/runtime/activation_checkpointing/checkpointing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index 6acefeb249e4..e1bbf1913041 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -893,4 +893,4 @@ def is_configured(): Return: True of configured, else False """ - return deepspeed_checkpointing_enabled \ No newline at end of file + return deepspeed_checkpointing_enabled