From 4a8d15d0eff1e0ba12682e83f267459986611915 Mon Sep 17 00:00:00 2001 From: Max Kovalenko <75629718+deepcharm@users.noreply.github.com> Date: Mon, 9 Oct 2023 04:32:25 +0300 Subject: [PATCH] Move tensors to device if mp is not enabled (#4461) Tensors are not moved to device in gather_partitioned_activations in case model parallel is not enabled. Co-authored-by: Olatunji Ruwase --- deepspeed/runtime/activation_checkpointing/checkpointing.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py index 77407a52026a..108cb37b57fb 100644 --- a/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ b/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -270,6 +270,8 @@ def gather_partitioned_activations(tensors, device=None): # 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())) + if device is not None: + item = item.to(device) inputs.append(item) continue