From 3589cad83a41e89447c171a8cd8f67edbb9d66b3 Mon Sep 17 00:00:00 2001 From: Quentin Anthony Date: Thu, 26 Oct 2023 12:35:20 -0400 Subject: [PATCH] Remove PP Grad Tail Check (#2538) * Only communicate grad tail if it exists Co-authored-by: Dashiell Stander * Revert previous patch and just always send the grad tail * Formatting --------- Co-authored-by: Dashiell Stander Co-authored-by: Olatunji Ruwase Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> --- deepspeed/runtime/pipe/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py index 2542dd602b1f..b89c1aca7f2a 100644 --- 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 = []