def distribute_layers_for_pipeline_stage(
        num_layers: int,
        num_virtual_layers_pre: int,
        num_virtual_layers_post: int,
        stage: PipelineStageInfo
) -> tuple[int, int]:
    """
    Calculates the layer index range for a specific pipeline stage.

    This function distributes a given number of layers across multiple pipeline
    stages as evenly as possible. It accounts for additional, non-layer
    computational load on the first and last stages (e.g., embeddings and the
    LM head) by using the concept of 'virtual layers' to reserve capacity.

    Args:
        num_layers: The total number of primary model layers to be distributed
            (e.g., the transformer blocks).
        num_virtual_layers_pre: The number of 'virtual' layers representing the
            computational cost of modules on the *first* stage, before the main
            layers (e.g., token and positional embeddings).
        num_virtual_layers_post: The number of 'virtual' layers representing the
            computational cost of modules on the *last* stage, after the main
            layers (e.g., the final layer normalization and LM head).
        stage: An object containing total stages and current stage index.

    Returns:
        A tuple (start_index, end_index), representing the slice of layers for
            the given stage. The start_index is inclusive and the end_index is
            exclusive.

    Raises:
        ValueError: If the pipeline configuration results in a stage having zero
            or negative layers assigned (pipeline too long for the model size).
    """

    num_layers_virtual = num_layers + num_virtual_layers_pre + num_virtual_layers_post

    base_layers_per_stage = num_layers_virtual // stage.num_stages
    extra_layers = num_layers_virtual % stage.num_stages

    layer_count_per_stage = []

    for proposed_stage_i in range(stage.num_stages):
        proposed_stage = PipelineStageInfo(num_stages=stage.num_stages, current_stage=proposed_stage_i)
        layers = base_layers_per_stage + 1 if proposed_stage_i < extra_layers else base_layers_per_stage

        adjustment = 0
        if proposed_stage.is_current_stage_first:
            adjustment += num_virtual_layers_pre
        if proposed_stage.is_current_stage_last:
            adjustment += num_virtual_layers_post

        actual_layers = layers - adjustment

        if actual_layers <= 0:
            raise ValueError(f"Tried to distribute layers, but got {actual_layers} on "
                             f"stage {proposed_stage.current_stage}. Perhaps the pipeline is too long for this model?")

        layer_count_per_stage.append(actual_layers)

    start_layer_id = sum(layer_count_per_stage[:stage.current_stage])
    num_layers_in_stage = layer_count_per_stage[stage.current_stage]

    return start_layer_id, start_layer_id + num_layers_in_stage

def build_schedule(
        dist_context: DistributedContext,
        n_microbatches: int,
        schedule_config: AnyPipelineScheduleConfig,
        model_provider: Callable[[PipelineStageInfo], nn.Module],
        loss_fn: Callable[[dict[str, torch.Tensor], int], torch.Tensor] | None,
        sharding_spec: PipelineShardingSpec
) -> tuple[PipelineScheduleInfo, list[nn.Module]]:
    """
    Constructs the pipeline schedule and instantiates model stages.

    This function coordinates the creation of the distributed pipeline. It:
    1.  Selects the appropriate `PipelineProgramBuilder` based on the config.
    2.  Calculates the global stage topology mapping stages to ranks.
    3.  Instantiates the local model stages for the current rank using `model_provider`.
    4.  Wraps models in `PipelineStage` containers.
    5.  Generates the execution program (action list).
    6.  Builds the runtime executor.

    Args:
        dist_context: The distributed context.
        n_microbatches: Number of microbatches per global step.
        schedule_config: Configuration object determining the schedule strategy.
        model_provider: A factory function that accepts stage info and returns an `nn.Module`
            for that specific stage.
        loss_fn: Optional loss function. Required if training (backward pass needed).
        sharding_spec: Specification for how data and states are sharded.

    Returns:
        A tuple containing:
        1.  `PipelineScheduleInfo`: The executable schedule and metadata.
        2.  `list[nn.Module]`: The local PyTorch modules created for this rank.
    """

    program_builder = PIPELINE_PROGRAM_REGISTRY.program_for(schedule_config)
    mesh = dist_context.mesh_for(REGULAR_DOMAIN)["pp"]

    num_stages = program_builder.num_stages_per_rank * mesh.size()

    stage_to_host = build_stage_to_host_rank_topology(
        num_stages=num_stages,
        pp_size=mesh.size(),
        style=program_builder.topology_style
    )
    host_to_stage = invert_stage_to_host_rank_topology(stage_to_host)
    this_rank_stages = host_to_stage[mesh.get_local_rank()]

    stages = []
    modules = []
    has_first_stage = False
    has_last_stage = False

    for stage_idx in this_rank_stages:
        stage_info = PipelineStageInfo(
            num_stages=num_stages,
            current_stage=stage_idx
        )

        if stage_info.is_current_stage_first:
            has_first_stage = True
        if stage_info.is_current_stage_last:
            has_last_stage = True

        model = model_provider(stage_info)
        modules.append(model)
        stage = PipelineStage(
            info=stage_info,
            module=model,
            group=mesh.get_group(),
            stage_to_host_topology=stage_to_host
        )
        stages.append(stage)

    program = program_builder.compose(num_microbatches=n_microbatches, pp_size=mesh.size())
    schedule = PipelineScheduleExecutor(
        dist_context=dist_context,
        stages=stages,
        num_microbatches=n_microbatches,
        loss_fn=loss_fn,
        program=program,
        sharding_spec=sharding_spec
    )

    return PipelineScheduleInfo(
        schedule=schedule,
        has_first_stage=has_first_stage,
        has_last_stage=has_last_stage
    ), modules