Skip to content

Commit

Permalink
[FIX] fix deadlock in PipeEngine._exec_recv_grads
Browse files Browse the repository at this point in the history
  • Loading branch information
i4never committed May 21, 2024
1 parent 1d81967 commit ad4fe8a
Showing 1 changed file with 1 addition and 3 deletions.
4 changes: 1 addition & 3 deletions deepspeed/runtime/pipe/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -1204,9 +1204,7 @@ def _exec_recv_grads(self, buffer_id):
# branches on is_grad_partitioned so we don't filter out the
# metadata tensor.
if self.is_grad_partitioned:
sizes_and_dtypes = [(list(t.size()), t.dtype)
for t in outputs[:2]] + [(list(t.size()), t.dtype)
for t in outputs[2:] if t.is_floating_point()]
sizes_and_dtypes = [(list(t.size()), t.dtype) for t in outputs[:2]]
else:
sizes_and_dtypes = [(list(t.size()), t.dtype) for t in outputs if t.is_floating_point()]
self.grad_layer = self._allocate_buffers(sizes_and_dtypes, num_buffers=1)[0]
Expand Down

0 comments on commit ad4fe8a

Please sign in to comment.