From 5a8c34146dd13a39b3e767ce26e8c9f8cac7e84d Mon Sep 17 00:00:00 2001 From: Hyunwoong Ko Date: Thu, 16 Sep 2021 18:42:04 +0900 Subject: [PATCH 1/9] Update checkpointing.py --- .../runtime/activation_checkpointing/checkpointing.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index e1bbf1913041..4be9c2d3ebba 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,16 +430,19 @@ 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() if i >= len(contiguous_size_buffers): From 01cc99c0522c49a4d0eb898c7a84fa0a011137de Mon Sep 17 00:00:00 2001 From: hyunwoongko Date: Sun, 19 Sep 2021 14:02:26 +0900 Subject: [PATCH 2/9] Fix formatting --- deepspeed/runtime/activation_checkpointing/checkpointing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index 4be9c2d3ebba..09d157a6e1d0 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -371,7 +371,7 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): 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( @@ -442,7 +442,7 @@ def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint arg.data = inp.data new_args.append(arg) i -= num_non_fp_tensors - + if contiguous_checkpoint: numel = size.numel() if i >= len(contiguous_size_buffers): From f16f718bf8f517e7418ed95feaa762992e493d0c Mon Sep 17 00:00:00 2001 From: hyunwoongko Date: Sun, 26 Sep 2021 04:22:58 +0900 Subject: [PATCH 3/9] 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 4/9] 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 5/9] 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 From b8009a972950cf08d5c8be9c1bc74c0957a7f529 Mon Sep 17 00:00:00 2001 From: Hyunwoong Ko Date: Tue, 28 Sep 2021 04:00:54 +0900 Subject: [PATCH 6/9] Update checkpointing.py --- deepspeed/runtime/activation_checkpointing/checkpointing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index 09d157a6e1d0..236ff25dd229 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -387,7 +387,7 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device) - for i in range(num_layers) + for _ in range(num_layers) ] contiguous_data_buffers.append(tensor_list) data_offsets.append(0) @@ -396,7 +396,7 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device) - for i in range(num_layers) + for _ in range(num_layers) ] contiguous_data_buffers[i] = tensor_list data_offsets[i] = 0 From 96f943ef105e514f5e4b21f51e0cb94ce0f98187 Mon Sep 17 00:00:00 2001 From: Hyunwoong Ko Date: Tue, 28 Sep 2021 05:01:14 +0900 Subject: [PATCH 7/9] Update checkpointing.py --- .../activation_checkpointing/checkpointing.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index 236ff25dd229..ccf4c52f5bff 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -366,13 +366,17 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): inputs = [] num_non_fp_tensors = 0 - for i, item in enumerate(args): + i = 0 + + for item in args: if not is_activation_to_checkpoint(item): inputs.append(item) num_non_fp_tensors += 1 continue + + if i >= num_non_fp_tensors: + i -= num_non_fp_tensors - i -= num_non_fp_tensors partition_size = get_partition_size(item) partition = item.detach().contiguous().view(-1).narrow( 0, @@ -423,6 +427,8 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): partition = partition.cpu() if CPU_CHECKPOINT else partition inputs.append(partition) + i += 1 + return inputs @@ -431,7 +437,9 @@ def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint new_args = [] num_non_fp_tensors = 0 - for i, (arg, inp) in enumerate(zip(args, inputs)): + i = 0 + + for arg, inp in 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) @@ -441,7 +449,9 @@ def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint arg.data = inp.data new_args.append(arg) - i -= num_non_fp_tensors + + if i >= num_non_fp_tensors: + i -= num_non_fp_tensors if contiguous_checkpoint: numel = size.numel() @@ -468,7 +478,9 @@ def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint new_args.append(contiguous_size) else: new_args.append(size) - + + i += 1 + return new_args From ac933c6a41e539453c93064bf6a30b54cb66e2e8 Mon Sep 17 00:00:00 2001 From: hyunwoongko Date: Sat, 2 Oct 2021 11:08:32 +0900 Subject: [PATCH 8/9] Reflect code review for contiguous activation checkpointing --- .../activation_checkpointing/checkpointing.py | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index 9a20c87017d7..2f98416f3236 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -368,11 +368,20 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): global contiguous_data_buffers, data_offsets inputs = [] - for i, item in enumerate(args): + num_non_fp_tensors = 0 + + for arg_index, item in enumerate(args): if not is_activation_to_checkpoint(item): inputs.append(item) + num_non_fp_tensors += 1 continue + if arg_index >= num_non_fp_tensors: + i = arg_index - num_non_fp_tensors + else: + # to prevent arg_index=0, num_non_fp_tensors=1 + i = arg_index + partition_size = get_partition_size(item) partition = item.detach().contiguous().view(-1).narrow( 0, @@ -387,7 +396,7 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device) - for i in range(num_layers) + for _ in range(num_layers) ] contiguous_data_buffers.append(tensor_list) data_offsets.append(0) @@ -396,7 +405,7 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device) - for i in range(num_layers) + for _ in range(num_layers) ] contiguous_data_buffers[i] = tensor_list data_offsets[i] = 0 @@ -430,16 +439,25 @@ def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint global contiguous_size_buffers, size_offsets new_args = [] - for i, (arg, inp) in enumerate(zip(args, inputs)): + num_non_fp_tensors = 0 + + for arg_index, (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) + if arg_index >= num_non_fp_tensors: + i = arg_index - num_non_fp_tensors + else: + # to prevent arg_index=0, num_non_fp_tensors=1 + i = arg_index + if contiguous_checkpoint: numel = size.numel() if i >= len(contiguous_size_buffers): From 41418019a9ea02e75490589a8f009f352999ba00 Mon Sep 17 00:00:00 2001 From: hyunwoongko Date: Sat, 2 Oct 2021 11:15:59 +0900 Subject: [PATCH 9/9] remove useless condition --- .../activation_checkpointing/checkpointing.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index 2f98416f3236..5cc45c49e8a1 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -376,12 +376,7 @@ def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): num_non_fp_tensors += 1 continue - if arg_index >= num_non_fp_tensors: - i = arg_index - num_non_fp_tensors - else: - # to prevent arg_index=0, num_non_fp_tensors=1 - i = arg_index - + i = arg_index - num_non_fp_tensors partition_size = get_partition_size(item) partition = item.detach().contiguous().view(-1).narrow( 0, @@ -451,12 +446,7 @@ def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint arg.data = inp.data new_args.append(arg) - - if arg_index >= num_non_fp_tensors: - i = arg_index - num_non_fp_tensors - else: - # to prevent arg_index=0, num_non_fp_tensors=1 - i = arg_index + i = arg_index - num_non_fp_tensors if contiguous_checkpoint: numel = size.numel()