Get PyTorch batched communication API working for interleaved schedule

from megatron import mpu

def _communicate(tensor_send_next, tensor_send_prev, recv_prev, recv_next,

                 use_ring_exchange=False):

def _communicate(tensor_send_next, tensor_send_prev, recv_prev, recv_next):

    """Communicate tensors between stages. Used as helper method in other

    communication methods that are used in megatron/schedules.py.

                   previous rank.

        recv_next: boolean for whether tensor should be received from

                   next rank.

        use_ring_exchange: boolean for whether torch.distributed.ring_exchange()

                           API should be used.

    Returns:

        (tensor_recv_prev, tensor_recv_next)

            tensor_send_prev = mpu.split_tensor_into_1d_equal_chunks(tensor_send_prev)

    # Send tensors in both the forward and backward directions as appropriate.

    if use_ring_exchange:

        torch.distributed.ring_exchange(tensor_send_prev=tensor_send_prev,

                                        tensor_recv_prev=tensor_recv_prev,

                                        tensor_send_next=tensor_send_next,

                                        tensor_recv_next=tensor_recv_next,

                                        group=mpu.get_pipeline_model_parallel_group())

    else:

    ops = []

    if tensor_send_prev is not None:

        send_prev_op = torch.distributed.P2POp(

            torch.distributed.irecv, tensor_recv_next,

            mpu.get_pipeline_model_parallel_next_rank())

        ops.append(recv_next_op)

    if len(ops) > 0:

        reqs = torch.distributed.batch_isend_irecv(ops)

        for req in reqs:

            req.wait()

    return tensor_recv_prev, tensor_recv_next

def recv_forward(timers=None, use_ring_exchange=False):

def recv_forward(timers=None):

    """Receive tensor from previous rank in pipeline (forward receive)."""

    if mpu.is_pipeline_first_stage():

        input_tensor = None

            tensor_send_next=None,

            tensor_send_prev=None,

            recv_prev=True,

            recv_next=False,

            use_ring_exchange=use_ring_exchange)

            recv_next=False)

        if timers is not None:

            timers('forward-recv').stop()

    return input_tensor

def recv_backward(timers=None, use_ring_exchange=False):

def recv_backward(timers=None):

    """Receive tensor from next rank in pipeline (backward receive)."""

    if mpu.is_pipeline_last_stage():

        output_tensor_grad = None

            tensor_send_next=None,

            tensor_send_prev=None,

            recv_prev=False,

            recv_next=True,

            use_ring_exchange=use_ring_exchange)

            recv_next=True)

        if timers is not None:

            timers('backward-recv').stop()

    return output_tensor_grad

def send_forward(output_tensor, timers=None, use_ring_exchange=False):

def send_forward(output_tensor, timers=None):

    """Send tensor to next rank in pipeline (forward send)."""

    if not mpu.is_pipeline_last_stage():

        if timers is not None:

            tensor_send_next=output_tensor,

            tensor_send_prev=None,

            recv_prev=False,

            recv_next=False,

            use_ring_exchange=use_ring_exchange)

            recv_next=False)

        if timers is not None:

            timers('forward-send').stop()

def send_backward(input_tensor_grad, timers=None, use_ring_exchange=False):

def send_backward(input_tensor_grad, timers=None):

    """Send tensor to previous rank in pipeline (backward send)."""

    if not mpu.is_pipeline_first_stage():

        if timers is not None:

            tensor_send_next=None,

            tensor_send_prev=input_tensor_grad,

            recv_prev=False,

            recv_next=False,

            use_ring_exchange=use_ring_exchange)

            recv_next=False)

        if timers is not None:

            timers('backward-send').stop()

def send_forward_recv_backward(output_tensor, timers=None, use_ring_exchange=False):

def send_forward_recv_backward(output_tensor, timers=None):

    """Batched send and recv with next rank in pipeline."""

    if mpu.is_pipeline_last_stage():

        output_tensor_grad = None

            tensor_send_next=output_tensor,

            tensor_send_prev=None,

            recv_prev=False,

            recv_next=True,

            use_ring_exchange=use_ring_exchange)

            recv_next=True)

        if timers is not None:

            timers('forward-send-backward-recv').stop()

    return output_tensor_grad

def send_backward_recv_forward(input_tensor_grad, timers=None, use_ring_exchange=False):

def send_backward_recv_forward(input_tensor_grad, timers=None):

    """Batched send and recv with previous rank in pipeline."""

    if mpu.is_pipeline_first_stage():

        input_tensor = None

            tensor_send_next=None,

            tensor_send_prev=input_tensor_grad,

            recv_prev=True,

            recv_next=False,

            use_ring_exchange=use_ring_exchange)

            recv_next=False)

        if timers is not None:

            timers('backward-send-forward-recv').stop()

    return input_tensor

        tensor_send_next=output_tensor,

        tensor_send_prev=None,

        recv_prev=recv_prev,

        recv_next=False,

        use_ring_exchange=True)

        recv_next=False)

    if timers is not None:

        timers('forward-send-forward-recv').stop()

    return input_tensor

        tensor_send_next=None,

        tensor_send_prev=input_tensor_grad,

        recv_prev=False,

        recv_next=recv_next,

        use_ring_exchange=True)

        recv_next=recv_next)

    if timers is not None:

        timers('backward-send-backward-recv').stop()

    return output_tensor_grad

        tensor_send_next=output_tensor,

        tensor_send_prev=input_tensor_grad,

        recv_prev=recv_prev,

        recv_next=recv_next,

        use_ring_exchange=True)

        recv_next=recv_next)

    if timers is not None:

        timers('forward-backward-send-forward-backward-recv').stop()

    return input_tensor, output_tensor_grad

    # Run warmup forward passes.

    mpu.set_virtual_pipeline_model_parallel_rank(0)

    input_tensors[0].append(

        p2p_communication.recv_forward(timers, use_ring_exchange=True))

        p2p_communication.recv_forward(timers))

    for k in range(num_warmup_microbatches):

        output_tensor = forward_step_helper(k)

    if not forward_only:

        if all_warmup_microbatches:

            output_tensor_grads[num_model_chunks-1].append(

                p2p_communication.recv_backward(timers, use_ring_exchange=True))

                p2p_communication.recv_backward(timers))

        for k in range(num_microbatches_remaining, num_microbatches):

            input_tensor_grad = backward_step_helper(k)

            next_backward_model_chunk_id = get_model_chunk_id(k+1, forward=False)