megatron.p2p_communication.send_forward_backward_recv_forward_backward

megatron.p2p_communication.send_forward_backward_recv_forward_backward(output_tensor, input_tensor_grad, recv_prev, recv_next, tensor_shape=None, timers=None)

Batched send and recv with previous and next ranks in pipeline.