megatron.p2p_communication

Description

Functions

recv_backward([tensor_shape, timers])	Receive tensor from next rank in pipeline (backward receive).
recv_forward([tensor_shape, dtype_, timers])	Receive tensor from previous rank in pipeline (forward receive).
send_backward(input_tensor_grad[, ...])	Send tensor to previous rank in pipeline (backward send).
send_backward_recv_backward(...[, ...])	Batched recv from next rank and send to previous rank in pipeline.
send_backward_recv_forward(input_tensor_grad)	Batched send and recv with previous rank in pipeline.
send_forward(output_tensor[, tensor_shape, ...])	Send tensor to next rank in pipeline (forward send).
send_forward_backward_recv_forward_backward(...)	Batched send and recv with previous and next ranks in pipeline.
send_forward_recv_backward(output_tensor[, ...])	Batched send and recv with next rank in pipeline.
send_forward_recv_forward(output_tensor, ...)	Batched recv from previous rank and send to next rank in pipeline.