Update engine.py

Remove PP Grad Tail Check (until #2538 is merged to upstream)
microsoft · Mar 9, 2023 · 64d6c5a · 64d6c5a
1 parent 457850d
commit 64d6c5a
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
@@ -990,9 +990,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 = []