Fix deadlock in PipeEngine._exec_recv_grads

[i4never]

I'm using Megtron-DeepSpeed with TP/PP/DP. In my case there are three tensors need to communicate between pipelines:

• hidden_state (floating, need grad)
• attention_mask (int32, no grad)
• cached_rotray_embedding (floating, no grad)

Only first tensor has grad which meets the restriction of PipelineEngine here:

DeepSpeed/deepspeed/runtime/pipe/engine.py

Lines 734 to 736 in 3dd7ccf

	# TODO: Improve pipe partitioning to pass multiple tensors that require grads
	assert all([torch.is_tensor(elt) and elt.requires_grad is False for elt in outputs[1:]])
	outputs_tail = outputs[1:]

Only grads of first tensor sended in first stage:

DeepSpeed/deepspeed/runtime/pipe/engine.py

Lines 1106 to 1109 in 3dd7ccf

	if self.is_grad_partitioned:
	# First two sends are partitioned gradient
	p2p.send(inputs[0], self.prev_stage)
	p2p.send(inputs[1], self.prev_stage)

But the next stage try to recv more than one grad because tensor.is_floating_point() is used to filter outputs. In my case cached_rotray_embedding is floating tensor with no grad which caught by filter. Next stage expecting more data than sended makes training hangs.

DeepSpeed/deepspeed/runtime/pipe/engine.py

Lines 1206 to 1209 in 3dd7ccf

	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()]

Since only one grad is send anyway, we don't need is_floating_point filter here.

[i4never]

rebase master

[i4never]

Hi @tjruwase, Could you review this pls?

[tjruwase]

Hi @tjruwase, Could you review this pls?

@i4never, thanks for this PR. This is very old code, which is documented as hacky, and unfortunately the author is no longer available. Although, I think your changes is correct, I think it is best to be more cautious. So, can you please add some unit tests?