From 82831afe8f42de5ed10005a072d67cd91016c7b9 Mon Sep 17 00:00:00 2001 From: Guo Yejun Date: Sun, 6 Nov 2022 17:03:59 -0800 Subject: [PATCH] skip torch.zeros and tensor.copy_ when model parallel is not used --- .../runtime/activation_checkpointing/checkpointing.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index b8d03e142b49..e1235387428f 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -270,8 +270,8 @@ def gather_partitioned_activations(tensors, device=None): inputs.append(item) continue - # don't need to do all_gather if model parallel size is 1 - if mp_size == 1: + # don't need to do all_gather if model parallel is not enabled + if mp_group is None or mp_size == 1: item = item.view(list(size.numpy())) inputs.append(item) continue @@ -290,8 +290,7 @@ def gather_partitioned_activations(tensors, device=None): if i == mp_rank: part_i.copy_(item) partitions.append(part_i) - if mp_group is not None: - dist.all_gather(partitions, partitions[mp_rank], group=mp_group) + dist.all_gather(partitions, partitions[mp_rank], group=mp_group) input_tensor = flat_tensor.view(list(size.numpy())) item.data = input_tensor.data