        input_tensor,
        forward_data_store,
        config,
        collect_non_loss_data,
    )

    if not forward_only:
        backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism and layernorm all-reduce for sequence parallelism).
        config.finalize_model_grads_func([model])

    return forward_data_store

def forward_backward_pipelining_with_interleaving(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,
    micro_batch_size: int,
    decoder_seq_length: int = None,
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
):
    """Run interleaved 1F1B schedule (model split into model chunks), with
    communication between pipeline stages as needed.

    Returns dictionary with losses if the last stage, empty dict otherwise."""
    assert isinstance(model, list), "interleaved pipeline parallelism expected model chunking"
    assert all(isinstance(chunk, torch.nn.Module) for chunk in model), "invalid model chunking"
    assert isinstance(
        data_iterator, list
    ), "interleaved pipeline parallelism expected each model chunk to have a data iterator"

    config = get_model_config(model[0])
    if config.overlap_p2p_comm and config.batch_p2p_comm:
        raise ValueError("Can not use both overlap_p2p_comm and batch_p2p_comm")

    if config.timers is not None:
        config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)

    # Disable async grad reductions
    no_sync_func = config.no_sync_func
    if no_sync_func is None:
        no_sync_func = contextlib.nullcontext
    no_sync_context = None

    def disable_grad_sync():
        """Disable asynchronous grad reductions"""
        nonlocal no_sync_context
        if no_sync_context is None:
            no_sync_context = no_sync_func()
            no_sync_context.__enter__()

    def enable_grad_sync():
        """Enable asynchronous grad reductions"""
        nonlocal no_sync_context
        if no_sync_context is not None:
            no_sync_context.__exit__(None, None, None)
            no_sync_context = None

    disable_grad_sync()

    # Model chunk IDs with synchronized grads
    synchronized_model_chunks = set()

    input_tensors = [[] for _ in range(len(model))]
    output_tensors = [[] for _ in range(len(model))]
    forward_data_store = []
    if not forward_only:
        output_tensor_grads = [[] for _ in range(len(model))]

    pipeline_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
    pipeline_parallel_rank = parallel_state.get_pipeline_model_parallel_rank()

    if num_microbatches % pipeline_parallel_size != 0:
        msg = f'number of microbatches ({num_microbatches}) is not divisible by '
        msg += f'pipeline-model-parallel-size ({pipeline_parallel_size}) '
        msg += 'when using interleaved schedule'
        raise RuntimeError(msg)

    model_type = get_model_type(model[0])
    if model_type == ModelType.encoder_and_decoder:
        raise RuntimeError("Interleaving is not supported with an encoder and decoder model.")

    if decoder_seq_length is not None and decoder_seq_length != seq_length:
        raise RuntimeError(
            "Interleaving is not supported with a different decoder sequence length."
        )

    tensor_shape = [seq_length, micro_batch_size, config.hidden_size]
    if config.sequence_parallel:
        tensor_shape[0] = tensor_shape[0] // parallel_state.get_tensor_model_parallel_world_size()

    # Compute number of warmup and remaining microbatches.
    num_model_chunks = len(model)
    total_num_microbatches = num_microbatches * num_model_chunks
    all_warmup_microbatches = False
    if forward_only:
        num_warmup_microbatches = total_num_microbatches
    else:
        # Run all forward passes and then all backward passes if number of
        # microbatches is just the number of pipeline stages.
        # Otherwise, perform (num_model_chunks-1)*pipeline_parallel_size on
        # all workers, followed by more microbatches after depending on
        # stage ID (more forward passes for earlier stages, later stages can
        # immediately start with 1F1B).
        if num_microbatches == pipeline_parallel_size:
            num_warmup_microbatches = total_num_microbatches
            all_warmup_microbatches = True
        else:
            num_warmup_microbatches = (pipeline_parallel_size - pipeline_parallel_rank - 1) * 2
            num_warmup_microbatches += (num_model_chunks - 1) * pipeline_parallel_size
            num_warmup_microbatches = min(num_warmup_microbatches, total_num_microbatches)
    num_microbatches_remaining = total_num_microbatches - num_warmup_microbatches

    # Checkpoint the activations of partial Transformer layers in a number of micro-batches
    # within the maximum outstanding micro-batch backpropagations.
    # Micro-batches with the ids less than 'num_microbatches_with_partial_activation_checkpoints'
    # checkpoint partial Transformer layers (or skip checkpointing) and
    # the rest of micro-batches within a window of micro-batches checkpoint
    # all Transformer layers. The window of micro-batches is set by the maximum
    # outstanding backpropagations and becomes smaller at later pipeline stages.
    # Please refer the appendix C in https://arxiv.org/pdf/2205.05198.pdf
    max_outstanding_backprops = None
    if config.num_microbatches_with_partial_activation_checkpoints is not None:
        max_outstanding_backprops = num_warmup_microbatches + 1

    # Synchronize params for first two model chunks
    if config.param_sync_func is not None:
        config.param_sync_func(model[0].parameters())
        config.param_sync_func(model[1].parameters())

    def get_model_chunk_id(microbatch_id, forward):
        """Helper method to get the model chunk ID given the iteration number."""
        microbatch_id_in_group = microbatch_id % (pipeline_parallel_size * num_model_chunks)
        model_chunk_id = microbatch_id_in_group // pipeline_parallel_size
        if not forward:
            model_chunk_id = num_model_chunks - model_chunk_id - 1
        return model_chunk_id

    def is_first_microbatch_for_model_chunk(microbatch_id: int) -> bool:
        """Check if an iteration is the first for a model chunk."""
        microbatch_group_size = pipeline_parallel_size * num_model_chunks
        num_microbatch_groups = total_num_microbatches // microbatch_group_size
        microbatch_group_id = microbatch_id // microbatch_group_size
        microbatch_id_in_group = microbatch_id % microbatch_group_size
        if microbatch_group_id == 0:
            return microbatch_id_in_group % pipeline_parallel_size == 0
        else:
            return False

    def is_last_microbatch_for_model_chunk(microbatch_id: int) -> bool:
        """Check if an iteration is the last for a model chunk."""
        microbatch_group_size = pipeline_parallel_size * num_model_chunks
        num_microbatch_groups = total_num_microbatches // microbatch_group_size
        microbatch_group_id = microbatch_id // microbatch_group_size
        microbatch_id_in_group = microbatch_id % microbatch_group_size
        if microbatch_group_id == num_microbatch_groups - 1:
            return microbatch_id_in_group % pipeline_parallel_size == pipeline_parallel_size - 1
        else:
            return False

    def forward_step_helper(microbatch_id, checkpoint_activations_microbatch):
        """Helper method to run forward step with model split into chunks
        (run set_virtual_pipeline_model_parallel_rank() before calling
        forward_step())."""
        model_chunk_id = get_model_chunk_id(microbatch_id, forward=True)
        parallel_state.set_virtual_pipeline_model_parallel_rank(model_chunk_id)

        # launch param synchronization for next model chunk
        # Note: Asynchronous communication tends to slow down compute.
        # To reduce idling from mismatched microbatch times, we launch
        # asynchronous communication at the same time across the
        # pipeline-parallel group.
        if config.param_sync_func is not None:
            param_sync_microbatch_id = microbatch_id + pipeline_parallel_rank
            if (
                param_sync_microbatch_id < total_num_microbatches
                and is_first_microbatch_for_model_chunk(param_sync_microbatch_id)
            ):
                param_sync_chunk_id = get_model_chunk_id(param_sync_microbatch_id, forward=True) + 1
                if 1 < param_sync_chunk_id < num_model_chunks:
                    config.param_sync_func(model[param_sync_chunk_id].parameters())

        # forward step
        if parallel_state.is_pipeline_first_stage():
            if len(input_tensors[model_chunk_id]) == len(output_tensors[model_chunk_id]):
                input_tensors[model_chunk_id].append(None)
        input_tensor = input_tensors[model_chunk_id][-1]
        output_tensor = forward_step(
            forward_step_func,
            data_iterator[model_chunk_id],
            model[model_chunk_id],
            num_microbatches,
            input_tensor,
            forward_data_store,
            config,
            collect_non_loss_data,
            checkpoint_activations_microbatch,
        )
        output_tensors[model_chunk_id].append(output_tensor)

        # if forward-only, no need to save tensors for a backward pass
        if forward_only:
            input_tensors[model_chunk_id].pop()
            output_tensors[model_chunk_id].pop()

        return output_tensor

    def backward_step_helper(microbatch_id):
        """Helper method to run backward step with model split into chunks
        (run set_virtual_pipeline_model_parallel_rank() before calling
        backward_step())."""
        model_chunk_id = get_model_chunk_id(microbatch_id, forward=False)
        parallel_state.set_virtual_pipeline_model_parallel_rank(model_chunk_id)

        # launch grad synchronization (default)
        if config.grad_sync_func is None and is_last_microbatch_for_model_chunk(microbatch_id):
            enable_grad_sync()
            synchronized_model_chunks.add(model_chunk_id)

        if parallel_state.is_pipeline_last_stage():
            if len(output_tensor_grads[model_chunk_id]) == 0:
                output_tensor_grads[model_chunk_id].append(None)
        input_tensor = input_tensors[model_chunk_id].pop(0)
        output_tensor = output_tensors[model_chunk_id].pop(0)
        output_tensor_grad = output_tensor_grads[model_chunk_id].pop(0)
        input_tensor_grad = backward_step(
            input_tensor, output_tensor, output_tensor_grad, model_type, config
        )

        # launch grad synchronization (custom grad sync)
        # Note: Asynchronous communication tends to slow down compute.
        # To reduce idling from mismatched microbatch times, we launch
        # asynchronous communication at the same time across the
        # pipeline-parallel group.
        if config.grad_sync_func is not None:
            grad_sync_microbatch_id = microbatch_id - pipeline_parallel_rank
            if grad_sync_microbatch_id >= 0 and is_last_microbatch_for_model_chunk(
                grad_sync_microbatch_id
            ):
                grad_sync_chunk_id = get_model_chunk_id(grad_sync_microbatch_id, forward=False)
                enable_grad_sync()
                config.grad_sync_func(model[grad_sync_chunk_id].parameters())
                synchronized_model_chunks.add(grad_sync_chunk_id)
        disable_grad_sync()

        return input_tensor_grad

    # Run warmup forward passes.
    parallel_state.set_virtual_pipeline_model_parallel_rank(0)
    input_tensors[0].append(p2p_communication.recv_forward(tensor_shape, config))

    fwd_wait_handles = None
    bwd_wait_handles = None

    for k in range(num_warmup_microbatches):

        if fwd_wait_handles is not None:
            for req in fwd_wait_handles:
                req.wait()

        # Decide to checkpoint all layers' activations of the current micro-batch
        if max_outstanding_backprops is not None:
            checkpoint_activations_microbatch = (
                k % max_outstanding_backprops
                >= config.num_microbatches_with_partial_activation_checkpoints
            )
        else:
            checkpoint_activations_microbatch = None

        output_tensor = forward_step_helper(k, checkpoint_activations_microbatch)

        # Determine if tensor should be received from previous stage.
        next_forward_model_chunk_id = get_model_chunk_id(k + 1, forward=True)
        recv_prev = True
        if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
            if next_forward_model_chunk_id == 0:
                recv_prev = False
        if k == (total_num_microbatches - 1):
            recv_prev = False

        # Don't send tensor downstream if on last stage.
        if parallel_state.is_pipeline_last_stage():
            output_tensor = None

        # Send and receive tensors as appropriate (send tensors computed
        # in this iteration; receive tensors for next iteration).
        if not config.overlap_p2p_comm:
            if (
                k == (num_warmup_microbatches - 1)
                and not forward_only
                and not all_warmup_microbatches
            ):
                input_tensor_grad = None
                recv_next = True
                if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
                    recv_next = False
                (
                    input_tensor,
                    output_tensor_grad,
                ) = p2p_communication.send_forward_backward_recv_forward_backward(
                    output_tensor,
                    input_tensor_grad,
                    recv_prev=recv_prev,
                    recv_next=recv_next,
                    tensor_shape=tensor_shape,
                    config=config,
                )
                output_tensor_grads[num_model_chunks - 1].append(output_tensor_grad)
            else:
                input_tensor = p2p_communication.send_forward_recv_forward(
                    output_tensor, recv_prev=recv_prev, tensor_shape=tensor_shape, config=config
                )
            input_tensors[next_forward_model_chunk_id].append(input_tensor)
        else:
            input_tensor, fwd_wait_handles = p2p_communication.send_forward_recv_forward(
                output_tensor,
                recv_prev=recv_prev,
                tensor_shape=tensor_shape,
                config=config,
                overlap_p2p_comm=True,
            )

            if (
                k == (num_warmup_microbatches - 1)
                and not forward_only
                and not all_warmup_microbatches
            ):
                input_tensor_grad = None
                recv_next = True
                if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
                    recv_next = False

                (
                    output_tensor_grad,
                    bwd_wait_handles,
                ) = p2p_communication.send_backward_recv_backward(
                    input_tensor_grad,
                    recv_next=recv_next,
                    tensor_shape=tensor_shape,
                    config=config,
                    overlap_p2p_comm=True,
                )

                output_tensor_grads[num_model_chunks - 1].append(output_tensor_grad)
            input_tensors[next_forward_model_chunk_id].append(input_tensor)

        deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)

    # Run 1F1B in steady state.
    for k in range(num_microbatches_remaining):
        # Forward pass.
        forward_k = k + num_warmup_microbatches

        # Decide to checkpoint all layers' activations of the current micro-batch
        if max_outstanding_backprops is not None:
            checkpoint_activations_microbatch = (
                forward_k % max_outstanding_backprops
                >= config.num_microbatches_with_partial_activation_checkpoints
            )
        else:
            checkpoint_activations_microbatch = None

        if config.overlap_p2p_comm:
            if fwd_wait_handles is not None:
                for req in fwd_wait_handles:
                    req.wait()

            deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)

            output_tensor = forward_step_helper(forward_k, checkpoint_activations_microbatch)

            # Determine if current stage has anything to send in either direction,
            # otherwise set tensor to None.
            forward_model_chunk_id = get_model_chunk_id(forward_k, forward=True)
            parallel_state.set_virtual_pipeline_model_parallel_rank(forward_model_chunk_id)

            # Last virtual stage no activation tensor to send
            if parallel_state.is_pipeline_last_stage():
                output_tensor = None

            # Determine if peers are sending, and where in data structure to put
            # received tensors.
            recv_prev = True
            if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
                # First stage is ahead of last stage by (pipeline_parallel_size - 1).
                next_forward_model_chunk_id = get_model_chunk_id(
                    forward_k - (pipeline_parallel_size - 1), forward=True
                )
                if next_forward_model_chunk_id == (num_model_chunks - 1):
                    recv_prev = False
                next_forward_model_chunk_id += 1
            else:
                next_forward_model_chunk_id = get_model_chunk_id(forward_k + 1, forward=True)

            # If last iteration, don't receive; we already received one extra
            # before the start of the for loop.
            if k == (num_microbatches_remaining - 1):
                recv_prev = False

            # Send activation tensor to the next stage and receive activation tensor from the
            # previous stage
            input_tensor, fwd_wait_handles = p2p_communication.send_forward_recv_forward(
                output_tensor,
                recv_prev=recv_prev,
                tensor_shape=tensor_shape,
                config=config,
                overlap_p2p_comm=True,
            )
            # assert fwd_wait_handles is not None

            if bwd_wait_handles is not None:
                for req in bwd_wait_handles:
                    req.wait()

            # Backward pass.
            backward_k = k
            input_tensor_grad = backward_step_helper(backward_k)

            backward_model_chunk_id = get_model_chunk_id(backward_k, forward=False)
            parallel_state.set_virtual_pipeline_model_parallel_rank(backward_model_chunk_id)

            # First virtual stage no activation gradient tensor to send
            if parallel_state.is_pipeline_first_stage():
                input_tensor_grad = None

            # Determine if the current virtual stage has an activation gradient tensor to receive
            recv_next = True
            if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
                # Last stage is ahead of first stage by (pipeline_parallel_size - 1).
                next_backward_model_chunk_id = get_model_chunk_id(
                    backward_k - (pipeline_parallel_size - 1), forward=False
                )
                if next_backward_model_chunk_id == 0:
                    recv_next = False
                next_backward_model_chunk_id -= 1
            else:
                next_backw
# ... truncated ...


            params = []
            for model_chunk_id in range(num_model_chunks):
                if model_chunk_id not in synchronized_model_chunks:
                    params.extend(model[model_chunk_id].parameters())
                    synchronized_model_chunks.add(model_chunk_id)
            if params:
                config.grad_sync_func(params)

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism, layernorm all-reduce for sequence parallelism, and
        # embedding all-reduce for pipeline parallelism).
        config.finalize_model_grads_func(model)

    return forward_data_store

def get_tensor_shapes(
    *,
    rank: int,
    model_type: ModelType,
    seq_length: int,
    micro_batch_size: int,
    decoder_seq_length: int,
    config,
):
    # Determine right tensor sizes (based on position of rank with respect to split
    # rank) and model size.
    # Send two tensors if model is T5 and rank is in decoder stage:
    #     first tensor is decoder (pre-transpose),
    #     second tensor is encoder (post-transpose).
    # If model is T5 and rank is at the boundary:
    #     send one tensor (post-transpose from encoder).
    # Otherwise, send one tensor (pre-transpose).
    tensor_shapes = []

    if config.sequence_parallel:
        seq_length = seq_length // parallel_state.get_tensor_model_parallel_world_size()
        if model_type == ModelType.encoder_and_decoder:
            decoder_seq_length = (
                decoder_seq_length // parallel_state.get_tensor_model_parallel_world_size()
            )

    if model_type == ModelType.encoder_and_decoder:
        if parallel_state.is_pipeline_stage_before_split(rank):
            tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
        else:
            tensor_shapes.append((decoder_seq_length, micro_batch_size, config.hidden_size))
            tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
    else:
        tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
    return tensor_shapes

def recv_forward(tensor_shapes, config):
    input_tensors = []
    if parallel_state.is_pipeline_first_stage():
        for tensor_shape in tensor_shapes:
            input_tensors.append(None)
    else:
        for tensor_shape in tensor_shapes:
            if tensor_shape is None:
                input_tensors.append(None)
            else:
                input_tensors.append(p2p_communication.recv_forward(tensor_shape, config))
    return input_tensors

def recv_backward(tensor_shapes, config):
    output_tensor_grads = []
    if parallel_state.is_pipeline_last_stage():
        for tensor_shape in tensor_shapes:


    tensor_shapes = []

    if config.sequence_parallel:
        seq_length = seq_length // parallel_state.get_tensor_model_parallel_world_size()
        if model_type == ModelType.encoder_and_decoder:
            decoder_seq_length = (
                decoder_seq_length // parallel_state.get_tensor_model_parallel_world_size()
            )

    if model_type == ModelType.encoder_and_decoder:
        if parallel_state.is_pipeline_stage_before_split(rank):
            tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
        else:
            tensor_shapes.append((decoder_seq_length, micro_batch_size, config.hidden_size))
            tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
    else:
        tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
    return tensor_shapes

def recv_forward(tensor_shapes, config):
    input_tensors = []
    if parallel_state.is_pipeline_first_stage():
        for tensor_shape in tensor_shapes:
            input_tensors.append(None)
    else:
        for tensor_shape in tensor_shapes:
            if tensor_shape is None:
                input_tensors.append(None)
            else:
                input_tensors.append(p2p_communication.recv_forward(tensor_shape, config))
    return input_tensors

def recv_backward(tensor_shapes, config):
    output_tensor_grads = []
    if parallel_state.is_pipeline_last_stage():
        for tensor_shape in tensor_shapes:
            output_tensor_grads.append(None)
    else:
        for tensor_shape in tensor_shapes:
            if tensor_shape is None:
                output_tensor_grads.append(None)
            else:
                output_tensor_grads.append(p2p_communication.recv_backward(tensor_shape, config))
    return output_tensor_grads

def send_forward(output_tensors, tensor_shapes, config):
    if not parallel_state.is_pipeline_last_stage():
        if not isinstance(output_tensors, list):
            output_tensors = [output_tensors]


            tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
    else:
        tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
    return tensor_shapes

def recv_forward(tensor_shapes, config):
    input_tensors = []
    if parallel_state.is_pipeline_first_stage():
        for tensor_shape in tensor_shapes:
            input_tensors.append(None)
    else:
        for tensor_shape in tensor_shapes:
            if tensor_shape is None:
                input_tensors.append(None)
            else:
                input_tensors.append(p2p_communication.recv_forward(tensor_shape, config))
    return input_tensors

def recv_backward(tensor_shapes, config):
    output_tensor_grads = []
    if parallel_state.is_pipeline_last_stage():
        for tensor_shape in tensor_shapes:
            output_tensor_grads.append(None)
    else:
        for tensor_shape in tensor_shapes:
            if tensor_shape is None:
                output_tensor_grads.append(None)
            else:
                output_tensor_grads.append(p2p_communication.recv_backward(tensor_shape, config))
    return output_tensor_grads

def send_forward(output_tensors, tensor_shapes, config):
    if not parallel_state.is_pipeline_last_stage():
        if not isinstance(output_tensors, list):
            output_tensors = [output_tensors]
        for (output_tensor, tensor_shape) in zip(output_tensors, tensor_shapes):
            if tensor_shape is None:
                continue
            p2p_communication.send_forward(output_tensor, config)

def send_backward(input_tensor_grads, tensor_shapes, config):
    if not parallel_state.is_pipeline_first_stage():
        if not isinstance(input_tensor_grads, list):
            input_tensor_grads = [input_tensor_grads]
        for (input_tensor_grad, tensor_shape) in zip(input_tensor_grads, tensor_shapes):
            if tensor_shape is None:
                continue
            p2p_communication.send_backward(input_tensor_grad, config)


                input_tensors.append(None)
            else:
                input_tensors.append(p2p_communication.recv_forward(tensor_shape, config))
    return input_tensors

def recv_backward(tensor_shapes, config):
    output_tensor_grads = []
    if parallel_state.is_pipeline_last_stage():
        for tensor_shape in tensor_shapes:
            output_tensor_grads.append(None)
    else:
        for tensor_shape in tensor_shapes:
            if tensor_shape is None:
                output_tensor_grads.append(None)
            else:
                output_tensor_grads.append(p2p_communication.recv_backward(tensor_shape, config))
    return output_tensor_grads

def send_forward(output_tensors, tensor_shapes, config):
    if not parallel_state.is_pipeline_last_stage():
        if not isinstance(output_tensors, list):
            output_tensors = [output_tensors]
        for (output_tensor, tensor_shape) in zip(output_tensors, tensor_shapes):
            if tensor_shape is None:
                continue
            p2p_communication.send_forward(output_tensor, config)

def send_backward(input_tensor_grads, tensor_shapes, config):
    if not parallel_state.is_pipeline_first_stage():
        if not isinstance(input_tensor_grads, list):
            input_tensor_grads = [input_tensor_grads]
        for (input_tensor_grad, tensor_shape) in zip(input_tensor_grads, tensor_shapes):
            if tensor_shape is None:
                continue
            p2p_communication.send_backward(input_tensor_grad, config)

def send_forward_recv_backward(output_tensors, tensor_shapes, config):
    if not isinstance(output_tensors, list):
        output_tensors = [output_tensors]
    output_tensor_grads = []
    for (output_tensor, tensor_shape) in zip(output_tensors, tensor_shapes):
        if tensor_shape is None:
            output_tensor_grads.append(None)
            continue


            output_tensor_grads.append(None)
    else:
        for tensor_shape in tensor_shapes:
            if tensor_shape is None:
                output_tensor_grads.append(None)
            else:
                output_tensor_grads.append(p2p_communication.recv_backward(tensor_shape, config))
    return output_tensor_grads

def send_forward(output_tensors, tensor_shapes, config):
    if not parallel_state.is_pipeline_last_stage():
        if not isinstance(output_tensors, list):
            output_tensors = [output_tensors]
        for (output_tensor, tensor_shape) in zip(output_tensors, tensor_shapes):
            if tensor_shape is None:
                continue
            p2p_communication.send_forward(output_tensor, config)

def send_backward(input_tensor_grads, tensor_shapes, config):
    if not parallel_state.is_pipeline_first_stage():
        if not isinstance(input_tensor_grads, list):
            input_tensor_grads = [input_tensor_grads]
        for (input_tensor_grad, tensor_shape) in zip(input_tensor_grads, tensor_shapes):
            if tensor_shape is None:
                continue
            p2p_communication.send_backward(input_tensor_grad, config)

def send_forward_recv_backward(output_tensors, tensor_shapes, config):
    if not isinstance(output_tensors, list):
        output_tensors = [output_tensors]
    output_tensor_grads = []
    for (output_tensor, tensor_shape) in zip(output_tensors, tensor_shapes):
        if tensor_shape is None:
            output_tensor_grads.append(None)
            continue
        output_tensor_grad = p2p_communication.send_forward_recv_backward(
            output_tensor, tensor_shape, config
        )
        output_tensor_grads.append(output_tensor_grad)
    return output_tensor_grads

def send_backward_recv_forward(input_tensor_grads, tensor_shapes, config):
    if not isinstance(input_tensor_grads, list):
        input_tensor_grads = [input_tensor_grads]


def send_forward(output_tensors, tensor_shapes, config):
    if not parallel_state.is_pipeline_last_stage():
        if not isinstance(output_tensors, list):
            output_tensors = [output_tensors]
        for (output_tensor, tensor_shape) in zip(output_tensors, tensor_shapes):
            if tensor_shape is None:
                continue
            p2p_communication.send_forward(output_tensor, config)

def send_backward(input_tensor_grads, tensor_shapes, config):
    if not parallel_state.is_pipeline_first_stage():
        if not isinstance(input_tensor_grads, list):
            input_tensor_grads = [input_tensor_grads]
        for (input_tensor_grad, tensor_shape) in zip(input_tensor_grads, tensor_shapes):
            if tensor_shape is None:
                continue
            p2p_communication.send_backward(input_tensor_grad, config)

def send_forward_recv_backward(output_tensors, tensor_shapes, config):
    if not isinstance(output_tensors, list):
        output_tensors = [output_tensors]
    output_tensor_grads = []
    for (output_tensor, tensor_shape) in zip(output_tensors, tensor_shapes):
        if tensor_shape is None:
            output_tensor_grads.append(None)
            continue
        output_tensor_grad = p2p_communication.send_forward_recv_backward(
            output_tensor, tensor_shape, config
        )
        output_tensor_grads.append(output_tensor_grad)
    return output_tensor_grads

def send_backward_recv_forward(input_tensor_grads, tensor_shapes, config):
    if not isinstance(input_tensor_grads, list):
        input_tensor_grads = [input_tensor_grads]
    input_tensors = []
    for (input_tensor_grad, tensor_shape) in zip(input_tensor_grads, tensor_shapes):
        if tensor_shape is None:
            input_tensors.append(None)
            continue
        input_tensor = p2p_communication.send_backward_recv_forward(
            input_tensor_grad, tensor_shape, config
        )
        input_tensors.append(input_tensor)
    return input_tensors

def forward_backward_pipelining_without_interleaving(
    *,
    forward_step_func,


            if tensor_shape is None:
                continue
            p2p_communication.send_backward(input_tensor_grad, config)

def send_forward_recv_backward(output_tensors, tensor_shapes, config):
    if not isinstance(output_tensors, list):
        output_tensors = [output_tensors]
    output_tensor_grads = []
    for (output_tensor, tensor_shape) in zip(output_tensors, tensor_shapes):
        if tensor_shape is None:
            output_tensor_grads.append(None)
            continue
        output_tensor_grad = p2p_communication.send_forward_recv_backward(
            output_tensor, tensor_shape, config
        )
        output_tensor_grads.append(output_tensor_grad)
    return output_tensor_grads

def send_backward_recv_forward(input_tensor_grads, tensor_shapes, config):
    if not isinstance(input_tensor_grads, list):
        input_tensor_grads = [input_tensor_grads]
    input_tensors = []
    for (input_tensor_grad, tensor_shape) in zip(input_tensor_grads, tensor_shapes):
        if tensor_shape is None:
            input_tensors.append(None)
            continue
        input_tensor = p2p_communication.send_backward_recv_forward(
            input_tensor_grad, tensor_shape, config
        )
        input_tensors.append(input_tensor)
    return input_tensors

def forward_backward_pipelining_without_interleaving(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,
    micro_batch_size: int,
    decoder_seq_length: int = None,
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
):
    """Run non-interleaved 1F1B schedule, with communication between pipeline
    stages.

    Returns dictionary with losses if the last stage, empty dict otherwise."""

    if isinstance(model, list):


        )
        output_tensor_grads.append(output_tensor_grad)
    return output_tensor_grads

def send_backward_recv_forward(input_tensor_grads, tensor_shapes, config):
    if not isinstance(input_tensor_grads, list):
        input_tensor_grads = [input_tensor_grads]
    input_tensors = []
    for (input_tensor_grad, tensor_shape) in zip(input_tensor_grads, tensor_shapes):
        if tensor_shape is None:
            input_tensors.append(None)
            continue
        input_tensor = p2p_communication.send_backward_recv_forward(
            input_tensor_grad, tensor_shape, config
        )
        input_tensors.append(input_tensor)
    return input_tensors

def forward_backward_pipelining_without_interleaving(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,
    micro_batch_size: int,
    decoder_seq_length: int = None,
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
):
    """Run non-interleaved 1F1B schedule, with communication between pipeline
    stages.

    Returns dictionary with losses if the last stage, empty dict otherwise."""

    if isinstance(model, list):
        assert (
            len(model) == 1
        ), "non-interleaved pipeline parallelism does not support model chunking"
        model = model[0]
    if isinstance(data_iterator, list):
        assert (
            len(data_iterator) == 1
        ), "non-pipeline-parallel schedule does not support model   "
        data_iterator = data_iterator[0]

    config = get_model_config(model)
    if config.overlap_p2p_comm:
        raise ValueError(
            "Non-interleaved pipeline parallelism does not support overlapping p2p communication"
        )

    if config.timers is not None:
        config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)

    # Disable async grad reductions
    no_sync_func = config.no_sync_func
    if no_sync_func is None:
        no_sync_func = contextlib.nullcontext
    no_sync_context = None

    def disable_grad_sync():
        """Disable asynchronous grad reductions"""
        nonlocal no_sync_context
        if no_sync_context is None:
            no_sync_context = no_sync_func()
            no_sync_context.__enter__()

    def enable_grad_sync():
        """Enable asynchronous grad reductions"""
        nonlocal no_sync_context
        if no_sync_context is not None:
            no_sync_context.__exit__(None, None, None)
            no_sync_context = None

    disable_grad_sync()

    # Compute number of warmup microbatches.
    num_warmup_microbatches = (
        parallel_state.get_pipeline_model_parallel_world_size()
        - parallel_state.get_pipeline_model_parallel_rank()
        - 1
    )
    num_warmup_microbatches = min(num_warmup_microbatches, num_microbatches)
    num_microbatches_remaining = num_microbatches - num_warmup_microbatches

    # Checkpoint the activations of partial Transformer layers in a number of micro-batches
    # within the maximum outstanding micro-batch backpropagations.
    # Micro-batches with the ids less than 'num_microbatches_with_partial_activation_checkpoints'
    # checkpoint partial Transformer layers (or skip checkpointing) and
    # the rest of micro-batches within a window of micro-batches checkpoint
    # all Transformer layers. The window of micro-batches is set by the maximum
    # outstanding backpropagations and becomes smaller at later pipeline stages.
    # Please refer the appendix C in https://arxiv.org/pdf/2205.05198.pdf
    max_outstanding_backprops = None
    if config.num_microbatches_with_partial_activation_checkpoints is not None:
        max_outstanding_backprops = num_warmup_microbatches + 1

    model_type = get_model_type(model)

    rank = parallel_state.get_pipeline_model_parallel_rank()
    recv_tensor_shapes = get_tensor_shapes(
        rank=rank - 1,
        model_type=model_type,
        seq_length=seq_length,
        micro_batch_size=micro_batch_size,
        decoder_seq_length=decoder_seq_length,
        config=config,
    )
    send_tensor_shapes = get_tensor_shapes(
        rank=rank,
        model_type=model_type,
        seq_length=seq_length,
        micro_batch_size=micro_batch_size,
        decoder_seq_length=decoder_seq_length,
        config=config,
    )

    # Input, output tensors only need to be saved when doing backward passes
    input_tensors = None
    output_tensors = None
    if not forward_only:
        input_tensors = []
        output_tensors = []
    forward_data_store = []

    # Run warmup forward passes.
    for i in range(num_warmup_microbatches):
        # Decide to checkpoint all layers' activations of the current micro-batch
        if max_outstanding_backprops is not None:
            checkpoint_activations_microbatch = (
                i % max_outstanding_backprops
                >= config.num_microbatches_with_partial_activation_checkpoints
            )
        else:
            checkpoint_activations_microbatch = None

        input_tensor = recv_forward(recv_tensor_shapes, config)
        output_tensor = forward_step(
            forward_step_func,
            data_iterator,
            model,
            num_microbatches,
            input_tensor,
            forward_data_store,
            config,
            collect_non_loss_data,
            checkpoint_activations_microbatch,
        )
        send_forward(output_tensor, send_tensor_shapes, config)

        if not forward_only:
            input_tensors.append(input_tensor)
            output_tensors.append(output_tensor)
            deallocate_output_tensor(output_tensor[0], config.deallocate_pipeline_outputs)

    # Before running 1F1B, need to receive first forward tensor.
    # If all microbatches are run in warmup / cooldown phase, then no need to
    # receive this tensor here.
    if num_microbatches_remaining > 0:
        input_tensor = recv_forward(recv_tensor_shapes, config)

    # Run 1F1B in steady state.
    for i in range(num_microbatches_remaining):
        last_iteration = i == (num_microbatches_remaining - 1)

        # Decide to checkpoint all layers' activations of the current micro-batch
        if max_outstanding_backprops is not None:
            checkpoint_activations_microbatch = (
                (i + num_warmup_microbatches) % max_outstanding_backprops
            ) >= config.num_microbatches_with_partial_activation_checkpoints
        else:
            checkpoint_activations_microbatch = None

        output_tensor = forward_step(
            forward_step_func,
            data_iterator,
            model,
            num_microbatches,
            input_tensor,
            forward_data_store,
            config,
            collect_non_loss_data,
            checkpoint_activations_microbatch,
        )

        if forward_only:
            send_forward(output_tensor, send_tensor_shapes, config)

            if not last_iteration:
                input_tensor = recv_forward(recv_tensor_shapes, config)

        else:
            output_tensor_grad = send_forward_recv_backward(
                output_tensor, send_tensor_shapes, config
            )

            # Add input_tensor and output_tensor to end of list.
            input_tensors.append(input_tensor)
            output_tensors.append(output_tensor)
            deallocate_output_tensor(output_tensor[0], config.deallocate_pipeline_outputs)

            # Pop input_tensor and output_tensor from the start of the list for
            # the backward pass.
            input_tensor = input_tensors.pop(0)
            output_tensor = output_tensors.pop(0)

            # Enable grad sync for the last microbatch in the batch if the full
            # backward pass completes in the 1F1B stage.
            if num_warmup_microbatches == 0 and last_iteration:
                if config.grad_sync_func is None or rank == 0:
                    enable_grad_sync()

            input_tensor_grad = backward_step(
                input_tensor, output_tensor, output_tensor_grad, model_type, config
            )

            if last_iteration:
                input_tensor = None
                send_backward(input_tensor_grad, recv_tensor_shapes, config)
            else:
                input_tensor = send_backward_recv_forward(
                    input_tensor_grad, recv_tensor_shapes, config
                )

    # Run cooldown backward passes.
    if not forward_only:
        for i in range(num_warmup_microbatches):

            # Enable async grad reduction in the last backward pass
            # Note: If grad sync function is provided, only enable
            # async grad reduction in first pipeline stage. Other
            # pipeline stages do grad reduction during pipeline
            # bubble.
            if i == num_warmup_microbatches - 1:
                if config.grad_sync_func is None or rank == 0:
                    enable_grad_sync()

            input_tensor = input_tensors.pop(0)
            output_tensor = output_tensors.pop(0)

            output_tensor_grad = recv_backward(send_tensor_shapes, config)

            input_tensor_grad = backward_step(
                input_tensor, output_tensor, output_tensor_grad, model_type, config
            )

            send_backward(input_tensor_grad, recv_tensor_shapes, config)

        # Launch any remaining grad reductions.
        if no_sync_context is not None:
            enable_grad_sync()
            if config.grad_sync_func is not None:
                config.grad_sync_func(model.parameters())

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism, layernorm all-reduce for sequence parallelism, and
        # embedding all-reduce for pipeline parallelism).
        config.finalize_model_grads_func([model])

    return forward_data_store

def early_exit_forward_backward_no_pipelining(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,  # unused
    micro_batch_size: int,  # unused
    decoder_seq_length: int = None,  # unused
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
    early_exit_loss_weight: EarlyExitLossWeight = None,
):
    """Run forward and backward passes with no pipeline parallelism
    (no inter-stage communication).

    Returns dictionary with losses.


            send_backward(input_tensor_grad, recv_tensor_shapes, config)

        # Launch any remaining grad reductions.
        if no_sync_context is not None:
            enable_grad_sync()
            if config.grad_sync_func is not None:
                config.grad_sync_func(model.parameters())

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism, layernorm all-reduce for sequence parallelism, and
        # embedding all-reduce for pipeline parallelism).
        config.finalize_model_grads_func([model])

    return forward_data_store

def early_exit_forward_backward_no_pipelining(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,  # unused
    micro_batch_size: int,  # unused
    decoder_seq_length: int = None,  # unused
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
    early_exit_loss_weight: EarlyExitLossWeight = None,
):
    """Run forward and backward passes with no pipeline parallelism
    (no inter-stage communication).

    Returns dictionary with losses.

    See get_forward_backward_func() for argument details
    """

    if isinstance(model, list):
        assert len(model) == 1, "non-pipeline-parallel schedule does not support model chunking"
        model = model[0]
    if isinstance(data_iterator, list):
        assert (
            len(data_iterator) == 1
        ), "non-pipeline-parallel schedule does not support model chunking"
        data_iterator = data_iterator[0]

    config = get_model_config(model)
    if config.timers is not None:
        config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)

    no_sync_func = config.no_sync_func
    if no_sync_func is None:
        no_sync_func = contextlib.nullcontext

    if early_exit_loss_weight:
        early_exit_loss_weight.update()

    forward_data_store = []
    backward_data_store = []
    input_tensor, output_tensor_grad = None, None
    with no_sync_func():
        for i in range(num_microbatches - 1):
            output_tensor, early_exit_output = early_exit_forward_step(
                forward_step_func,
                data_iterator,
                model,
                num_microbatches,
                input_tensor,
                forward_data_store,
                config,
                collect_non_loss_data,
            )
            if not forward_only:
                exit_loss = cal_early_exit_loss(early_exit_output, backward_data_store, num_microbatches, early_exit_loss_weight)
                early_exit_backward_step(input_tensor, output_tensor, output_tensor_grad, config, early_exit_loss=exit_loss)

    # Run computation for last microbatch out of context handler (want to
    # synchronize gradients).
    output_tensor, early_exit_output = early_exit_forward_step(
        forward_step_func,
        data_iterator,
        model,
        num_microbatches,
        input_tensor,
        forward_data_store,
        config,
        collect_non_loss_data,
    )

    if not forward_only:
        exit_loss = cal_early_exit_loss(early_exit_output, backward_data_store, num_microbatches, early_exit_loss_weight)
        early_exit_backward_step(input_tensor, output_tensor, output_tensor_grad, config, early_exit_loss=exit_loss)

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism and layernorm all-reduce for sequence parallelism).
        config.finalize_model_grads_func([model])
    if len(backward_data_store) == len(forward_data_store):
        backward_data_store = [{**f, **b} for (f, b) in zip(forward_data_store, backward_data_store)]
    return backward_data_store

def forward_backward_pipelining_for_early_exit_tuning(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,
    micro_batch_size: int,
    decoder_seq_length: int = None,
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
    early_exit_loss_weight: EarlyExitLossWeight = None,
):
    if isinstance(model, list):
        assert (
            len(model) == 1
        ), "non-interleaved pipeline parallelism does not support model chunking"
        model = model[0]


        config,
        collect_non_loss_data,
    )

    if not forward_only:
        exit_loss = cal_early_exit_loss(early_exit_output, backward_data_store, num_microbatches, early_exit_loss_weight)
        early_exit_backward_step(input_tensor, output_tensor, output_tensor_grad, config, early_exit_loss=exit_loss)

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism and layernorm all-reduce for sequence parallelism).
        config.finalize_model_grads_func([model])
    if len(backward_data_store) == len(forward_data_store):
        backward_data_store = [{**f, **b} for (f, b) in zip(forward_data_store, backward_data_store)]
    return backward_data_store

def forward_backward_pipelining_for_early_exit_tuning(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,
    micro_batch_size: int,
    decoder_seq_length: int = None,
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
    early_exit_loss_weight: EarlyExitLossWeight = None,
):
    if isinstance(model, list):
        assert (
            len(model) == 1
        ), "non-interleaved pipeline parallelism does not support model chunking"
        model = model[0]
    if isinstance(data_iterator, list):
        assert (
            len(data_iterator) == 1
        ), "non-pipeline-parallel schedule does not support model chunking"
        data_iterator = data_iterator[0]

    config = get_model_config(model)
    if config.overlap_p2p_comm:
        raise ValueError(
            "Non-interleaved pipeline parallelism does not support overlapping p2p communication"
        )

    if config.timers is not None:
        config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)

    # Disable async grad reductions
    no_sync_func = config.no_sync_func
    if no_sync_func is None:
        no_sync_func = contextlib.nullcontext
    no_sync_context = None
    has_early_exit = parallel_state.has_early_exit()

    def disable_grad_sync():
        """Disable asynchronous grad reductions"""
        nonlocal no_sync_context
        if no_sync_context is None:
            no_sync_context = no_sync_func()
            no_sync_context.__enter__()

    def enable_grad_sync():
        """Enable asynchronous grad reductions"""
        nonlocal no_sync_context
        if no_sync_context is not None:
            no_sync_context.__exit__(None, None, None)
            no_sync_context = None

    disable_grad_sync()

    if early_exit_loss_weight:
        early_exit_loss_weight.update()

    model_type = get_model_type(model)

    rank = parallel_state.get_pipeline_model_parallel_rank()
    recv_tensor_shapes = get_tensor_shapes(
        rank=rank - 1,
        model_type=model_type,
        seq_length=seq_length,
        micro_batch_size=micro_batch_size,
        decoder_seq_length=decoder_seq_length,
        config=config,
    )
    send_tensor_shapes = get_tensor_shapes(
        rank=rank,
        model_type=model_type,
        seq_length=seq_length,
        micro_batch_size=micro_batch_size,
        decoder_seq_length=decoder_seq_length,
        config=config,
    )

    forward_data_store = []

    # Run warmup forward passes.
    for i in range(num_microbatches):
        with torch.no_grad():
            input_tensor = recv_forward(recv_tensor_shapes, config)
            output_tensor, early_exit_loss_func = early_exit_forward_step(
                forward_step_func,
                data_iterator,
                model,
                num_microbatches,
                input_tensor,
                forward_data_store,
                config,
                collect_non_loss_data,
            )
            send_forward(output_tensor, send_tensor_shapes, config)

        if has_early_exit and not forward_only:
            if i == num_microbatches - 1:
                if config.grad_sync_func is None or rank == 0:
                    enable_grad_sync()
            exit_loss = cal_early_exit_loss(early_exit_loss_func, forward_data_store, num_microbatches, early_exit_loss_weight)
            early_exit_backward_step(input_tensor, output_tensor, [None], config,
                early_exit_loss=exit_loss
            )

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism, layernorm all-reduce for sequence parallelism, and
        # embedding all-reduce for pipeline parallelism).
        config.finalize_model_grads_func([model])

    return forward_data_store

def early_exit_forward_backward_pipelining(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,
    micro_batch_size: int,
    decoder_seq_length: int = None,
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
    early_exit_loss_weight: EarlyExitLossWeight = None,
):
    if isinstance(model, list):
        assert (
            len(model) == 1
        ), "non-interleaved pipeline parallelism does not support model chunking"
        model = model[0]


            if i == num_microbatches - 1:
                if config.grad_sync_func is None or rank == 0:
                    enable_grad_sync()
            exit_loss = cal_early_exit_loss(early_exit_loss_func, forward_data_store, num_microbatches, early_exit_loss_weight)
            early_exit_backward_step(input_tensor, output_tensor, [None], config,
                early_exit_loss=exit_loss
            )

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism, layernorm all-reduce for sequence parallelism, and
        # embedding all-reduce for pipeline parallelism).
        config.finalize_model_grads_func([model])

    return forward_data_store

def early_exit_forward_backward_pipelining(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,
    micro_batch_size: int,
    decoder_seq_length: int = None,
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
    early_exit_loss_weight: EarlyExitLossWeight = None,
):
    if isinstance(model, list):
        assert (
            len(model) == 1
        ), "non-interleaved pipeline parallelism does not support model chunking"
        model = model[0]
    if isinstance(data_iterator, list):
        assert (
            len(data_iterator) == 1
        ), "non-pipeline-parallel schedule does not support model chunking"
        data_iterator = data_iterator[0]

    config = get_model_config(model)
    if config.overlap_p2p_comm:
        raise ValueError(
            "Non-interleaved pipeline parallelism does not support overlapping p2p communication"
        )

    if config.timers is not None:
        config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)

    # Disable async grad reductions
    no_sync_func = config.no_sync_func
    if no_sync_func is None:
        no_sync_func = contextlib.nullcontext
    no_sync_context = None
    has_early_exit = parallel_state.has_early_exit()

    def disable_grad_sync():
        """Disable asynchronous grad reductions"""
        nonlocal no_sync_context
        if no_sync_context is None:
            no_sync_context = no_sync_func()
            no_sync_context.__enter__()

    def enable_grad_sync():
        """Enable asynchronous grad reductions"""
        nonlocal no_sync_context
        if no_sync_context is not None:
            no_sync_context.__exit__(None, None, None)
            no_sync_context = None

    disable_grad_sync()

    if early_exit_loss_weight:
        early_exit_loss_weight.update()

    # Compute number of warmup microbatches.
    num_warmup_microbatches = (
        parallel_state.get_pipeline_model_parallel_world_size()
        - parallel_state.get_pipeline_model_parallel_rank()
        - 1
    )
    num_warmup_microbatches = min(num_warmup_microbatches, num_microbatches)
    num_microbatches_remaining = num_microbatches - num_warmup_microbatches

    # Checkpoint the activations of partial Transformer layers in a number of micro-batches
    # within the maximum outstanding micro-batch backpropagations.
    # Micro-batches with the ids less than 'num_microbatches_with_partial_activation_checkpoints'
    # checkpoint partial Transformer layers (or skip checkpointing) and
    # the rest of micro-batches within a window of micro-batches checkpoint
    # all Transformer layers. The window of micro-batches is set by the maximum
    # outstanding backpropagations and becomes smaller at later pipeline stages.
    # Please refer the appendix C in https://arxiv.org/pdf/2205.05198.pdf
    max_outstanding_backprops = None
    if config.num_microbatches_with_partial_activation_checkpoints is not None:
        max_outstanding_backprops = num_warmup_microbatches + 1

    model_type = get_model_type(model)

    rank = parallel_state.get_pipeline_model_parallel_rank()
    recv_tensor_shapes = get_tensor_shapes(
        rank=rank - 1,
        model_type=model_type,
        seq_length=seq_length,
        micro_batch_size=micro_batch_size,
        decoder_seq_length=decoder_seq_length,
        config=config,
    )
    send_tensor_shapes = get_tensor_shapes(
        rank=rank,
        model_type=model_type,
        seq_length=seq_length,
        micro_batch_size=micro_batch_size,
        decoder_seq_length=decoder_seq_length,
        config=config,
    )

    # Input, output tensors only need to be saved when doing backward passes
    input_tensors = None
    output_tensors = None
    if not forward_only:
        input_tensors = []
        output_tensors = []
    if has_early_exit:
        early_exit_loss_funcs = []

    forward_data_store = []

    # Run warmup forward passes.
    for i in range(num_warmup_microbatches):
        # Decide to checkpoint all layers' activations of the current micro-batch
        if max_outstanding_backprops is not None:
            checkpoint_activations_microbatch = (
                i % max_outstanding_backprops
                >= config.num_microbatches_with_partial_activation_checkpoints
            )
        else:
            checkpoint_activations_microbatch = None

        input_tensor = recv_forward(recv_tensor_shapes, config)
        output_tensor, early_exit_output = early_exit_forward_step(
            forward_step_func,
            data_iterator,
            model,
            num_microbatches,
            input_tensor,
            forward_data_store,
            config,
            collect_non_loss_data,
            checkpoint_activations_microbatch,
        )
        send_forward(output_tensor, send_tensor_shapes, config)

        if not forward_only:
            if not has_early_exit:
                deallocate_output_tensor(output_tensor[0], config.deallocate_pipeline_outputs)
            input_tensors.append(input_tensor)
            output_tensors.append(output_tensor)
            if has_early_exit:
                early_exit_loss_funcs.append(early_exit_output)

    # Before running 1F1B, need to receive first forward tensor.
    # If all microbatches are run in warmup / cooldown phase, then no need to
    # receive this tensor here.
    if num_microbatches_remaining > 0:
        input_tensor = recv_forward(recv_tensor_shapes, config)

    # Run 1F1B in steady state.
    for i in range(num_microbatches_remaining):
        last_iteration = i == (num_microbatches_remaining - 1)

        # Decide to checkpoint all layers' activations of the current micro-batch
        if max_outstanding_backprops is not None:
            checkpoint_activations_microbatch = (
                (i + num_warmup_microbatches) % max_outstanding_backprops
            ) >= config.num_microbatches_with_partial_activation_checkpoints
        else:
            checkpoint_activations_microbatch = None
        output_tensor, early_exit_output = early_exit_forward_step(
            forward_step_func,
            data_iterator,
            model,
            num_microbatches,
            input_tensor,
            forward_data_store,
            config,
            collect_non_loss_data,
            checkpoint_activations_microbatch,
        )

        if forward_only:
            send_forward(output_tensor, send_tensor_shapes, config)
            if not last_iteration:
                input_tensor = recv_forward(recv_tensor_shapes, config)
        else:
            output_tensor_grad = send_forward_recv_backward(
                output_tensor, send_tensor_shapes, config
            )
            if not has_early_exit:
                deallocate_output_tensor(output_tensor[0], config.deallocate_pipeline_outputs)
            # Add input_tensor and output_tensor to end of list.
            input_tensors.append(input_tensor)
            output_tensors.append(output_tensor)
            if has_early_exit:
                early_exit_loss_funcs.append(early_exit_output)

            # Pop input_tensor and output_tensor from the start of the list for
            # the backward pass.
            input_tensor = input_tensors.pop(0)
            output_tensor = output_tensors.pop(0)

            if has_early_exit:
                exit_loss = cal_early_exit_loss(early_exit_loss_funcs.pop(0), forward_data_store, num_microbatches, early_exit_loss_weight)
            else:
                exit_loss = None

            input_tensor_grad = early_exit_backward_step(
                input_tensor, output_tensor, output_tensor_grad, config,
                early_exit_loss=exit_loss
            )

            if last_iteration:
                input_tensor = None
                send_backward(input_tensor_grad, recv_tensor_shapes, config)
            else:
                input_tensor = send_backward_recv_forward(
                    input_tensor_grad, recv_tensor_shapes, config
                )

    # Run cooldown backward passes.
    if not forward_only:
        for i in range(num_warmup_microbatches):
            # Enable async grad reduction in the last backward pass
            # Note: If grad sync function is provided, only enable
            # async grad reduction in first pipeline stage. Other
            # pipeline stages do grad reduction during pipeline
            # bubble.
            if i == num_warmup_microbatches - 1:
                if config.grad_sync_func is None or rank == 0:
                    enable_grad_sync()

            input_tensor = input_tensors.pop(0)
            output_tensor = output_tensors.pop(0)
            if has_early_exit:
                exit_loss = cal_early_exit_loss(early_exit_loss_funcs.pop(0), forward_data_store, num_microbatches, early_exit_loss_weight)
            else:
                exit_loss = None
            output_tensor_grad = recv_backward(send_tensor_shapes, config)

            input_tensor_grad = early_exit_backward_step(
                input_tensor, output_tensor, output_tensor_grad, config,
                early_exit_loss=exit_loss
            )

            send_backward(input_tensor_grad, recv_tensor_shapes, config)

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism, layernorm all-reduce for sequence parallelism, and
        # embedding all-reduce for pipeline parallelism).
        config.finalize_model_grads_func([model])

    return forward_data_store

def early_exit_forward_backward_pipelining_with_bubble_filling(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,
    micro_batch_size: int,
    decoder_seq_length: int = None,
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
    num_fill_warmup_microbatches: int = 0,
    num_fill_cooldown_microbatches: int = 0,
    backward_forward_ratio: float = 2.0,
    early_exit_loss_weight: EarlyExitLossWeight = None,
):
    if isinstance(model, list):
        assert (


            input_tensor_grad = early_exit_backward_step(
                input_tensor, output_tensor, output_tensor_grad, config,
                early_exit_loss=exit_loss
            )

            send_backward(input_tensor_grad, recv_tensor_shapes, config)

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism, layernorm all-reduce for sequence parallelism, and
        # embedding all-reduce for pipeline parallelism).
        config.finalize_model_grads_func([model])

    return forward_data_store

def early_exit_forward_backward_pipelining_with_bubble_filling(
    *,
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,
    micro_batch_size: int,
    decoder_seq_length: int = None,
    forward_only: bool = False,
    collect_non_loss_data: bool = False,
    num_fill_warmup_microbatches: int = 0,
    num_fill_cooldown_microbatches: int = 0,
    backward_forward_ratio: float = 2.0,
    early_exit_loss_weight: EarlyExitLossWeight = None,
):
    if isinstance(model, list):
        assert (
            len(model) == 1
        ), "non-interleaved pipeline parallelism does not support model chunking"
        model = model[0]
    if isinstance(data_iterator, list):
        assert (
            len(data_iterator) == 1
        ), "non-pipeline-parallel schedule does not support model chunking"
        data_iterator = data_iterator[0]

    config = get_model_config(model)
    if config.overlap_p2p_comm:
        raise ValueError(
            "Non-interleaved pipeline parallelism does not support overlapping p2p communication"
        )

    if config.timers is not None:
        config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)

    # Disable async grad reductions
    no_sync_func = config.no_sync_func
    if no_sync_func is None:
        no_sync_func = contextlib.nullcontext
    no_sync_context = None
    has_early_exit = parallel_state.has_early_exit()

    def disable_grad_sync():
        """Disable asynchronous grad reductions"""
        nonlocal no_sync_context
        if no_sync_context is None:
            no_sync_context = no_sync_func()
            no_sync_context.__enter__()

    def enable_grad_sync():
        """Enable asynchronous grad reductions"""
        nonlocal no_sync_context
        if no_sync_context is not None:
            no_sync_context.__exit__(None, None, None)
            no_sync_context = None

    disable_grad_sync()

    if early_exit_loss_weight:
        early_exit_loss_weight.update()

    # Compute number of warmup microbatches.
    num_warmup_microbatches = (
        parallel_state.get_pipeline_model_parallel_world_size()
        - parallel_state.get_pipeline_model_parallel_rank()
    )

    # todo @(pxc): check mirco_batch num globally
    assert num_warmup_microbatches < num_microbatches

    num_microbatches_remaining = num_microbatches - num_warmup_microbatches - num_fill_warmup_microbatches

    num_microbatches_for_loss = num_microbatches - num_fill_warmup_microbatches

    cur_warmup_bubble_size = backward_forward_ratio * (num_warmup_microbatches - 1) - parallel_state.get_pipeline_model_parallel_rank()
    next_warmup_bubble_size = max(0, cur_warmup_bubble_size - backward_forward_ratio - 1.0)

    cur_num_partial_forward_microbatches = min(num_fill_warmup_microbatches, int(cur_warmup_bubble_size / (1.0 + backward_forward_ratio)))
    next_num_partial_forward_microbatches = min(num_fill_warmup_microbatches, int(next_warmup_bubble_size / (1.0 + backward_forward_ratio)))

    cur_cooldown_bubble_size = backward_forward_ratio * parallel_state.get_pipeline_model_parallel_rank()
    pre_cooldown_bubble_size = max(0, cur_cooldown_bubble_size - backward_forward_ratio)

    cur_num_partial_backward_microbatches = min(num_fill_cooldown_microbatches, int(cur_cooldown_bubble_size / (1.0 + backward_forward_ratio)))
    pre_num_partial_backward_microbatches = min(num_fill_cooldown_microbatches, int(pre_cooldown_bubble_size / (1.0 + backward_forward_ratio)))

    num_cooldown_microbatches = max(0, num_warmup_microbatches - num_fill_cooldown_microbatches + cur_num_partial_backward_microbatches)

    # Checkpoint the activations of partial Transformer layers in a number of micro-batches
    # within the maximum outstanding micro-batch backpropagations.
    # Micro-batches with the ids less than 'num_microbatches_with_partial_activation_checkpoints'
    # checkpoint partial Transformer layers (or skip checkpointing) and
    # the rest of micro-batches within a window of micro-batches checkpoint
    # all Transformer layers. The window of micro-batches is set by the maximum
    # outstanding backpropagations and becomes smaller at later pipeline stages.
    # Please refer the appendix C in https://arxiv.org/pdf/2205.05198.pdf
    max_outstanding_backprops = None
    if config.num_microbatches_with_partial_activation_checkpoints is not None:
        max_outstanding_backprops = num_warmup_microbatches
    model_type = get_model_type(model)

    rank = parallel_state.get_pipeline_model_parallel_rank()
    recv_tensor_shapes = get_tensor_shapes(
        rank=rank - 1,
        model_type=model_type,
        seq_length=seq_length,
        micro_batch_size=micro_batch_size,
        decoder_seq_length=decoder_seq_length,
        config=config,
    )
    send_tensor_shapes = get_tensor_shapes(
        rank=rank,
        model_type=model_type,
        seq_length=seq_length,
        micro_batch_size=micro_batch_size,
        decoder_seq_length=decoder_seq_length,
        config=config,
    )

    # Input, output tensors only need to be saved when doing backward passes
    input_tensors = None
    output_tensors = None
    if not forward_only:
        input_tensors = []
        output_tensors = []
    if has_early_exit:
        early_exit_loss_funcs = []

    forward_data_store = []

    # Run warmup forward passes.
    for i in range(num_warmup_microbatches):
        # Decide to checkpoint all layers' activations of the current micro-batch
        if max_outstanding_backprops is not None:
            checkpoint_activations_microbatch = (
                i % max_outstanding_backprops
                >= config.num_microbatches_with_partial_activation_checkpoints
            )
        else:
            checkpoint_activations_microbatch = None

        input_tensor = recv_forward(recv_tensor_shapes, config)
        output_tensor, early_exit_output = early_exit_forward_step(
            forward_step_func,
            data_iterator,
            model,
            num_microbatches_for_loss,
            input_tensor,
            forward_data_store,
            config,
            collect_non_loss_data,
            checkpoint_activations_microbatch,
        )
        send_forward(output_tensor, send_tensor_shapes, config)

        if not forward_only:
            if not has_early_exit:
                deallocate_output_tensor(output_tensor[0], config.deallocate_pipeline_outputs)
            input_tensors.append(input_tensor)
            output_tensors.append(output_tensor)
            if has_early_exit:
                early_exit_loss_funcs.append(early_exit_output)

    warmup_input_tensors = []
    warmup_output_tensors = []
    warmup_early_exit_outputs = []

    # Fill warmup bubbles
    for i in range(num_fill_warmup_microbatches):
        if (has_early_exit or parallel_state.post_stage_has_early_exit()) and \
                (i < cur_num_partial_forward_microbatches):
            is_last = i + 1 > next_num_partial_forward_microbatches
            input_tensor = recv_forward(recv_tensor_shapes, config)
            output_tensor, early_exit_output = early_exit_forward_step(
                forward_step_func,
                data_iterator,
                model,
                num_microbatches_for_loss,
                input_tensor,
                forward_data_store,
                config,
                collect_non_loss_data,
                checkpoint_activations_microbatch,
            )
            if not is_last:
                send_forward(output_tensor, send_tensor_shapes, config)
            if not forward_only:
                if not has_early_exit:
                    deallocate_output_tensor(output_tensor[0], config.deallocate_pipeline_outputs)
                warmup_input_tensors.append(input_tensor)
                warmup_output_tensors.append(output_tensor)
                warmup_early_exit_outputs.append(early_exit_output)
        else:
            # align data iterator of all pipeline stages
            next(data_iterator)

    if not forward_only:
        for i in reversed(range(cur_num_partial_forward_microbatches)):
            # compute first backward without recv backward
            is_last = i + 1 > next_num_partial_forward_microbatches
            if is_last:
                output_tensor_grad = None
            else:
                output_tensor_grad = recv_backward(
                    send_tensor_shapes, config
                )

            input_tensor = warmup_input_tensors.pop(-1)
            output_tensor = warmup_output_tensors.pop(-1)
            exit_output = warmup_early_exit_outputs.pop(-1)
            if exit_output:
                exit_loss = cal_early_exit_loss(
                    exit_output, forward_data_store,
                    num_microbatches_for_loss, early_exit_loss_weight)
                input_tensor_grad = early_exit_backward_step(
                    input_tensor, output_tensor, output_tensor_grad, config,
                    early_exit_loss=exit_loss
                )
                send_backward(input_tensor_grad, recv_tensor_shapes, config)
            elif output_tensor_grad:
                input_tensor_grad = backward_step(
                    input_tensor, output_tensor, output_tensor_grad, model_type, config
                )
                send_backward(input_tensor_grad, recv_tensor_shapes, config)

    # Run 1F1B in steady state.
    for i in range(num_microbatches_remaining):
        # Decide to checkpoint all layers' activations of the current micro-batch
        if max_outstanding_backprops is not None:
            checkpoint_activations_microbatch = (
                (i + num_warmup_microbatches) % max_outstanding_backprops
            ) >= config.num_microbatches_with_partial_activation_checkpoints
        else:
            checkpoint_activations_microbatch = None

        if not forward_only:
            output_tensor_grad = recv_backward(
                send_tensor_shapes, config
            )
            # Pop input_tensor and output_tensor from the start of the list for
            # the backward pass.
            input_tensor = input_tensors.pop(0)
            output_tensor = output_tensors.pop(0)

            if has_early_exit:
                exit_loss = cal_early_exit_loss(
                    early_exit_loss_funcs.pop(0), forward_data_store,
                    num_microbatches_for_loss, early_exit_loss_weight)
            else:
                exit_loss = None

            input_tensor_grad = early_exit_backward_step(
                input_tensor, output_tensor, output_tensor_grad, config,
                early_exit_loss=exit_loss
            )
            input_tensor = recv_forward(recv_tensor_shapes, config)
            send_backward(input_tensor_grad, recv_tensor_shapes, config)
        else:
            input_tensor = recv_forward(recv_tensor_shapes, config)

        output_tensor, early_exit_output = early_exit_forward_step(
            forward_step_func,
            data_iterator,
            model,
            num_microbatches_for_loss,
            input_tensor,
            forward_data_store,
            config,
            collect_non_loss_data,
            checkpoint_activations_microbatch,
        )
        send_forward(output_tensor, send_tensor_shapes, config)
        if not forward_only:
            if not has_early_exit:
                deallocate_output_tensor(output_tensor[0], config.deallocate_pipeline_outputs)
            # Add input_tensor and output_tensor to end of list.
            input_tensors.append(input_tensor)
            output_tensors.append(output_tensor)
            if has_early_exit:
                early_exit_loss_funcs.append(early_exit_output)

    if not forward_only:
        # Run cooldown backward passes.
        for i in range(num_cooldown_microbatches):
            is_partial_backward = (num_cooldown_microbatches - i) <= (cur_num_partial_backward_microbatches - pre_num_partial_backward_microbatches)
            if i == num_cooldown_microbatches - 1:
                if config.grad_sync_func is None or rank == 0:
                    enable_grad_sync()

            input_tensor = input_tensors.pop(0)
            output_tensor = output_tensors.pop(0)
            if has_early_exit:
                exit_loss = cal_early_exit_loss(early_exit_loss_funcs.pop(0), forward_data_store, num_microbatches_for_loss, early_exit_loss_weight)
            else:
                exit_loss = None
            output_tensor_grad = recv_backward(send_tensor_shapes, config)

            input_tensor_grad = early_exit_backward_step(
                input_tensor, output_tensor, output_tensor_grad, config,
                early_exit_loss=exit_loss
            )
            if not is_partial_backward:
                send_backward(input_tensor_grad, recv_tensor_shapes, config)

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism, layernorm all-reduce for sequence parallelism, and
        # embedding all-reduce for pipeline parallelism).
        config.finalize_model_grads_func([model])

    return forward_data_store

def cal_early_exit_loss(early_exit_loss_funcs, forward_data_store, num_microbatches, early_exit_loss_weight):
    exit_loss_dict = {}
    exit_losses = []
    for layer_num, exit_loss_func in early_exit_loss_funcs.items():
        loss = exit_loss_func(log_dict=exit_loss_dict)
        loss_weight = early_exit_loss_weight.get_weight(layer_num)
        exit_losses.append(loss.multiply_(loss_weight))
        exit_loss_dict[f'exit weight [{layer_num}]'] = early_exit_loss_weight.get_weight(layer_num)
    forward_data_store.append(exit_loss_dict)
    return torch.sum(torch.stack(exit_losses), dim=0).div(num_microbatches)

def early_exit_forward_step(
    forward_step_func,
    data_iterator,
    model,
    num_microbatches,
    input_tensor,


            input_tensor_grad = early_exit_backward_step(
                input_tensor, output_tensor, output_tensor_grad, config,
                early_exit_loss=exit_loss
            )
            if not is_partial_backward:
                send_backward(input_tensor_grad, recv_tensor_shapes, config)

    if config.timers is not None:
        config.timers('forward-backward').stop()

    if config.finalize_model_grads_func is not None and not forward_only:
        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism, layernorm all-reduce for sequence parallelism, and
        # embedding all-reduce for pipeline parallelism).
        config.finalize_model_grads_func([model])

    return forward_data_store

def cal_early_exit_loss(early_exit_loss_funcs, forward_data_store, num_microbatches, early_exit_loss_weight):
    exit_loss_dict = {}
    exit_losses = []
    for layer_num, exit_loss_func in early_exit_loss_funcs.items():
        loss = exit_loss_func(log_dict=exit_loss_dict)
        loss_weight = early_exit_loss_weight.get_weight(layer_num)
        exit_losses.append(loss.multiply_(loss_weight))
        exit_loss_dict[f'exit weight [{layer_num}]'] = early_exit_loss_weight.get_weight(layer_num)
    forward_data_store.append(exit_loss_dict)
    return torch.sum(torch.stack(exit_losses), dim=0).div(num_microbatches)

def early_exit_forward_step(
    forward_step_func,
    data_iterator,
    model,
    num_microbatches,
    input_tensor,
    forward_data_store,
    config,
    collect_non_loss_data=False,
    checkpoint_activations_microbatch=None,
):
    """Forward step for early exit model.

    If first stage, input tensor is obtained from data_iterator, otherwise
    passed-in input_tensor is used.

    Returns output tensor."""
    if config.timers is not None:


        # Finalize model grads (perform full grad all-reduce / reduce-scatter for
        # data parallelism, layernorm all-reduce for sequence parallelism, and
        # embedding all-reduce for pipeline parallelism).
        config.finalize_model_grads_func([model])

    return forward_data_store

def cal_early_exit_loss(early_exit_loss_funcs, forward_data_store, num_microbatches, early_exit_loss_weight):
    exit_loss_dict = {}
    exit_losses = []
    for layer_num, exit_loss_func in early_exit_loss_funcs.items():
        loss = exit_loss_func(log_dict=exit_loss_dict)
        loss_weight = early_exit_loss_weight.get_weight(layer_num)
        exit_losses.append(loss.multiply_(loss_weight))
        exit_loss_dict[f'exit weight [{layer_num}]'] = early_exit_loss_weight.get_weight(layer_num)
    forward_data_store.append(exit_loss_dict)
    return torch.sum(torch.stack(exit_losses), dim=0).div(num_microbatches)

def early_exit_forward_step(
    forward_step_func,
    data_iterator,
    model,
    num_microbatches,
    input_tensor,
    forward_data_store,
    config,
    collect_non_loss_data=False,
    checkpoint_activations_microbatch=None,
):
    """Forward step for early exit model.

    If first stage, input tensor is obtained from data_iterator, otherwise
    passed-in input_tensor is used.

    Returns output tensor."""
    if config.timers is not None:
        config.timers('forward-compute', log_level=2).start()

    unwrap_output_tensor = False
    if not isinstance(input_tensor, list):
        input_tensor = [input_tensor]
        unwrap_output_tensor = True

    set_input_tensor = get_attr_wrapped_model(model, "set_input_tensor")
    set_input_tensor(input_tensor)

    if config.enable_autocast:
        context_manager = torch.autocast("cuda", dtype=config.autocast_dtype)
    else:
        context_manager = contextlib.nullcontext()
    with context_manager:
        if checkpoint_activations_microbatch is None:
            lm_output, loss_func = forward_step_func(data_iterator, model)
        else:
            lm_output, loss_func = forward_step_func(
                data_iterator, model, checkpoint_activations_microbatch
            )
    early_exit_output = None
    if parallel_state.has_early_exit():
        output_tensor, early_exit_output = lm_output
    else:
        output_tensor = lm_output
    loss_dict = {}

    if parallel_state.is_pipeline_last_stage() and not parallel_state.is_tune_exit():
        output_tensor = loss_func(output_tensor=output_tensor,
                                  log_dict=loss_dict,
                                  log_key='lm loss')
        output_tensor.div_(num_microbatches)

    if loss_dict:
        forward_data_store.append(loss_dict)

    if config.timers is not None:
        config.timers('forward-compute').stop()

    if unwrap_output_tensor:
        return output_tensor, early_exit_output

    return [output_tensor], early_exit_output

def early_exit_backward_step(input_tensor, output_tensor, output_tensor_grad, config, early_exit_loss=None):
    """Backward step through passed-in output tensor.

    If last stage, output_tensor_grad is None, otherwise gradient of loss
    with respect to stage's output tensor.

    Returns gradient of loss with respect to input tensor (None if first
    stage)."""

    if config.timers is not None:
        config.timers('backward-compute', log_level=2).start()

    # Retain the grad on the input_tensor.
    unwrap_input_tensor_grad = False
    if not isinstance(input_tensor, list):
        input_tensor = [input_tensor]
        unwrap_input_tensor_grad = True
    for x in input_tensor:


    loss_dict = {}

    if parallel_state.is_pipeline_last_stage() and not parallel_state.is_tune_exit():
        output_tensor = loss_func(output_tensor=output_tensor,
                                  log_dict=loss_dict,
                                  log_key='lm loss')
        output_tensor.div_(num_microbatches)

    if loss_dict:
        forward_data_store.append(loss_dict)

    if config.timers is not None:
        config.timers('forward-compute').stop()

    if unwrap_output_tensor:
        return output_tensor, early_exit_output

    return [output_tensor], early_exit_output

def early_exit_backward_step(input_tensor, output_tensor, output_tensor_grad, config, early_exit_loss=None):
    """Backward step through passed-in output tensor.

    If last stage, output_tensor_grad is None, otherwise gradient of loss
    with respect to stage's output tensor.

    Returns gradient of loss with respect to input tensor (None if first
    stage)."""

    if config.timers is not None:
        config.timers('backward-compute', log_level=2).start()

    # Retain the grad on the input_tensor.
    unwrap_input_tensor_grad = False
    if not isinstance(input_tensor, list):
        input_tensor = [input_tensor]
        unwrap_input_tensor_grad = True
    for x in input_tensor:
        if x is not None:
            x.retain_grad()

    if not isinstance(output_tensor, list):
        output_tensor = [output_tensor]
    if not isinstance(output_tensor_grad, list):
        output_tensor_grad = [output_tensor_grad]

    # Backward pass.
    if output_tensor_grad[0] is None and config.grad_scale_func is not None:
        output_tensor[0] = config.grad_scale_func(output_tensor[0])
    if early_exit_loss is not None:
        if output_tensor_grad[0] is not None:
            fake_loss = early_exit_loss + torch.sum(output_tensor[0] * output_tensor_grad[0])
        elif output_tensor[0].numel() == 1:
            fake_loss = early_exit_loss + output_tensor[0]
        else:
            fake_loss = early_exit_loss
        custom_backward(fake_loss, None)
    elif config.deallocate_pipeline_outputs:
        custom_backward(output_tensor[0], output_tensor_grad[0])
    else:
        torch.autograd.backward(output_tensor[0], grad_tensors=output_tensor_grad[0])

    # Collect the grad of the input_tensor.
    input_tensor_grad = [None]
    if input_tensor is not None:
        input_tensor_grad = []
        for x in input_tensor:
            if x is None:
                input_tensor_grad.append(None)
            else:
                input_tensor_grad.append(x.grad)

    if unwrap_input_tensor_grad:
        input_tensor_grad = input_tensor_grad[0]

    if config.timers is not None:
        config.timers('backward-compute').stop()

    return input_tensor_grad


                                            num_fill_cooldown_microbatches=args.num_fill_cooldown_microbatches,
                                            early_exit_loss_weight=early_exit_loss_weight)
            elif args.tune_exit:
                forward_backward_func = partial(forward_backward_pipelining_for_early_exit_tuning, early_exit_loss_weight=early_exit_loss_weight)
            else:
                forward_backward_func = partial(early_exit_forward_backward_pipelining, early_exit_loss_weight=early_exit_loss_weight)
        elif parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None:
            forward_backward_func = forward_backward_pipelining_with_interleaving
        else:
            forward_backward_func = forward_backward_pipelining_without_interleaving
    else:
        if len(args.exit_layer_nums) > 0:
            forward_backward_func = partial(early_exit_forward_backward_no_pipelining, early_exit_loss_weight=early_exit_loss_weight)
        else:
            forward_backward_func = forward_backward_no_pipelining
    return forward_backward_func

class EarlyExitLossWeight():

    def __init__(self, exit_layer_loss_weight, exit_layer_loss_weight_init,
                 exit_layer_weight_warmup_iters, exit_layer_weight_warmup_style):
        args = get_args()
        self.warmup = exit_layer_weight_warmup_iters > 0 and args.curr_iteration < exit_layer_weight_warmup_iters
        if self.warmup:
            self.warmup_iters = exit_layer_weight_warmup_iters
            self.exit_layer_loss_weight = {layer_num: weight for layer_num, weight in exit_layer_loss_weight_init.items()}
            self.exit_layer_loss_weight_init = exit_layer_loss_weight_init
            self.exit_layer_loss_weight_delta = {
                layer_num: exit_layer_loss_weight[layer_num] - exit_layer_loss_weight_init[layer_num]
                    for layer_num in exit_layer_loss_weight.keys()
            }
            if exit_layer_weight_warmup_style == 'cosine':
                self.update_func = self.cosine_warmup
            else: # linear
                self.update_func = self.linear_warmup
        else:
            self.exit_layer_loss_weight = exit_layer_loss_weight

    def cosine_warmup(self, inc_ratio):
        for layer_num in self.exit_layer_loss_weight.keys():
            self.exit_layer_loss_weight[layer_num] = 0.5 * (math.cos(math.pi * (inc_ratio + 1.0)) + 1.0) \
                    * self.exit_layer_loss_weight_delta[layer_num] + self.exit_layer_loss_weight_init[layer_num]

    def linear_warmup(self, inc_ratio):
        for layer_num in self.exit_layer_loss_weight.keys():
            self.exit_layer_loss_weight[layer_num] = inc_ratio * self.exit_layer_loss_weight_delta[layer_num] \
                + self.exit_layer_loss_weight_init[layer_num]

    def get_weight(self, layer):
        return self.exit_layer_loss_weight[layer]

    def update(self):
        if self.warmup:
            iteration = get_args().curr_iteration
            if iteration <= self.warmup_iters:
                self.update_func(float(iteration) / self.warmup_iters)
        return


    def __init__(self, exit_layer_loss_weight, exit_layer_loss_weight_init,
                 exit_layer_weight_warmup_iters, exit_layer_weight_warmup_style):
        args = get_args()
        self.warmup = exit_layer_weight_warmup_iters > 0 and args.curr_iteration < exit_layer_weight_warmup_iters
        if self.warmup:
            self.warmup_iters = exit_layer_weight_warmup_iters
            self.exit_layer_loss_weight = {layer_num: weight for layer_num, weight in exit_layer_loss_weight_init.items()}
            self.exit_layer_loss_weight_init = exit_layer_loss_weight_init
            self.exit_layer_loss_weight_delta = {
                layer_num: exit_layer_loss_weight[layer_num] - exit_layer_loss_weight_init[layer_num]
                    for layer_num in exit_layer_loss_weight.keys()
            }
            if exit_layer_weight_warmup_style == 'cosine':
                self.update_func = self.cosine_warmup
            else: # linear
                self.update_func = self.linear_warmup
        else:
            self.exit_layer_loss_weight = exit_layer_loss_weight

    def cosine_warmup(self, inc_ratio):
        for layer_num in self.exit_layer_loss_weight.keys():
            self.exit_layer_loss_weight[layer_num] = 0.5 * (math.cos(math.pi * (inc_ratio + 1.0)) + 1.0) \
                    * self.exit_layer_loss_weight_delta[layer_num] + self.exit_layer_loss_weight_init[layer_num]

    def linear_warmup(self, inc_ratio):
        for layer_num in self.exit_layer_loss_weight.keys():
            self.exit_layer_loss_weight[layer_num] = inc_ratio * self.exit_layer_loss_weight_delta[layer_num] \
                + self.exit_layer_loss_weight_init[layer_num]

    def get_weight(self, layer):
        return self.exit_layer_loss_weight[layer]

    def update(self):
        if self.warmup:
            iteration = get_args().curr_iteration
            if iteration <= self.warmup_iters:
                self.update_func(float(iteration) / self.warmup_iters)
        return

def deallocate_output_tensor(out, deallocate_pipeline_outputs=False):
    '''Pseudo-deallocate (i.e., set to scalar) the output tensor's '.data' field.

    This method should be called right after the output tensor has been


        if self.warmup:
            self.warmup_iters = exit_layer_weight_warmup_iters
            self.exit_layer_loss_weight = {layer_num: weight for layer_num, weight in exit_layer_loss_weight_init.items()}
            self.exit_layer_loss_weight_init = exit_layer_loss_weight_init
            self.exit_layer_loss_weight_delta = {
                layer_num: exit_layer_loss_weight[layer_num] - exit_layer_loss_weight_init[layer_num]
                    for layer_num in exit_layer_loss_weight.keys()
            }
            if exit_layer_weight_warmup_style == 'cosine':
                self.update_func = self.cosine_warmup
            else: # linear
                self.update_func = self.linear_warmup
        else:
            self.exit_layer_loss_weight = exit_layer_loss_weight

    def cosine_warmup(self, inc_ratio):
        for layer_num in self.exit_layer_loss_weight.keys():
            self.exit_layer_loss_weight[layer_num] = 0.5 * (math.cos(math.pi * (inc_ratio + 1.0)) + 1.0) \
                    * self.exit_layer_loss_weight_delta[layer_num] + self.exit_layer_loss_weight_init[layer_num]

    def linear_warmup(self, inc_ratio):
        for layer_num in self.exit_layer_loss_weight.keys():
            self.exit_layer_loss_weight[layer_num] = inc_ratio * self.exit_layer_loss_weight_delta[layer_num] \
                + self.exit_layer_loss_weight_init[layer_num]

    def get_weight(self, layer):
        return self.exit_layer_loss_weight[layer]

    def update(self):
        if self.warmup:
            iteration = get_args().curr_iteration
            if iteration <= self.warmup_iters:
                self.update_func(float(iteration) / self.warmup_iters)
        return

def deallocate_output_tensor(out, deallocate_pipeline_outputs=False):
    '''Pseudo-deallocate (i.e., set to scalar) the output tensor's '.data' field.

    This method should be called right after the output tensor has been
    sent to the next pipeline stage. At this point, the output tensor is
    only useful for its '.grad_fn' field, and not its '.data'.
    '''
    if (out is None) or (not deallocate_pipeline_outputs):
        return


                layer_num: exit_layer_loss_weight[layer_num] - exit_layer_loss_weight_init[layer_num]
                    for layer_num in exit_layer_loss_weight.keys()
            }
            if exit_layer_weight_warmup_style == 'cosine':
                self.update_func = self.cosine_warmup
            else: # linear
                self.update_func = self.linear_warmup
        else:
            self.exit_layer_loss_weight = exit_layer_loss_weight

    def cosine_warmup(self, inc_ratio):
        for layer_num in self.exit_layer_loss_weight.keys():
            self.exit_layer_loss_weight[layer_num] = 0.5 * (math.cos(math.pi * (inc_ratio + 1.0)) + 1.0) \
                    * self.exit_layer_loss_weight_delta[layer_num] + self.exit_layer_loss_weight_init[layer_num]

    def linear_warmup(self, inc_ratio):
        for layer_num in self.exit_layer_loss_weight.keys():
            self.exit_layer_loss_weight[layer_num] = inc_ratio * self.exit_layer_loss_weight_delta[layer_num] \
                + self.exit_layer_loss_weight_init[layer_num]

    def get_weight(self, layer):
        return self.exit_layer_loss_weight[layer]

    def update(self):
        if self.warmup:
            iteration = get_args().curr_iteration
            if iteration <= self.warmup_iters:
                self.update_func(float(iteration) / self.warmup_iters)
        return

def deallocate_output_tensor(out, deallocate_pipeline_outputs=False):
    '''Pseudo-deallocate (i.e., set to scalar) the output tensor's '.data' field.

    This method should be called right after the output tensor has been
    sent to the next pipeline stage. At this point, the output tensor is
    only useful for its '.grad_fn' field, and not its '.data'.
    '''
    if (out is None) or (not deallocate_pipeline_outputs):
        return
    assert isinstance(out, torch.Tensor), "expected Tensor, found %s." % type(out).__name__
    assert out._base is None, "counter-productive to free a view of another tensor."
    out.data = torch.empty((1,), device=out.device, dtype=out.dtype,)