Move tensors to device if mp is not enabled (microsoft#4461)

Tensors are not moved to device in gather_partitioned_activations in case model parallel is not enabled. Co-authored-by: Olatunji Ruwase <[email protected]>
baodii · Nov 7, 2023 · 4a8d15d · 4a8d15d
1 parent 63b9656
commit 4a8d15d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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