Remove PP Grad Tail Check (#2538)

* Only communicate grad tail if it exists Co-authored-by: Dashiell Stander <[email protected]> * Revert previous patch and just always send the grad tail * Formatting --------- Co-authored-by: Dashiell Stander <[email protected]> Co-authored-by: Olatunji Ruwase <[email protected]> Co-authored-by: Logan Adams <[email protected]>
microsoft · Oct 26, 2023 · 3589cad · 3589cad
1 parent b891ec2
commit 3589cad
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
@@ -988,7 +988,7 @@ def _exec_send_grads(self, buffer_id):
             if isinstance(inputs, tuple):
                 first_input = inputs[0]
                 assert all([torch.is_tensor(elt) for elt in inputs[1:]])
-                inputs_grad_tail = [elt.grad for elt in inputs[1:] if elt.grad is not None]
+                inputs_grad_tail = [elt.grad for elt in inputs[1:]]
             elif torch.is_tensor(inputs):
                 first_input = inputs
                 inputs_grad_tail = []