Avoid partitioning small activations (#1154)

microsoft · Jun 18, 2021 · b1669c0 · b1669c0
1 parent 1ba3e8e
commit b1669c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/runtime/activation_checkpointing/checkpointing.py b/deepspeed/runtime/activation_checkpointing/checkpointing.py
@@ -416,7 +416,7 @@ def save_args_for_backward(*all_args):
 
             inputs = []
             for i, item in enumerate(args[:-1]):
-                if not torch.is_tensor(item):
+                if not torch.is_tensor(item) or mp_size > item.numel():
                     inputs.append(item)
                     continue