Pipeline Parallelism

import dataclasses
import torch
from torch import nn
from d9d.pipelining.api import (
    ModuleSupportsPipelining,
    PipelineStageInfo,
    StageBoundary,
    TensorSpec,
    distribute_layers_for_pipeline_stage,
)


@dataclasses.dataclass
class MyInput:              # PipelineInput
    input_ids: torch.Tensor


@dataclasses.dataclass
class MyTransfer:           # StageTransfer (the only role crossing P2P)
    hidden_states: torch.Tensor


@dataclasses.dataclass
class MyShared:             # SharedInput (broadcast to every stage)
    position_ids: torch.Tensor


@dataclasses.dataclass
class MyOutput:             # PipelineOutput
    logits: torch.Tensor


class MyModelChunk(
    nn.Module,
    ModuleSupportsPipelining[MyInput, MyTransfer, MyShared, MyOutput],
):
    def __init__(self, stage: PipelineStageInfo, config):
        super().__init__()
        self.stage = stage
        self.config = config

        # 1. Determine what layers live here
        self.start_layer, self.end_layer = distribute_layers_for_pipeline_stage(
            config.n_layers, num_virtual_layers_pre=1, num_virtual_layers_post=1, stage=stage
        )

        # 2. Build sub-modules (using ModuleDict - for compatibility)
        self.layers = nn.ModuleDict({
            str(layer): TransformerBlock(...)
            for layer in range(self.start_layer, self.end_layer)
        })

        # Only build embeddings on first stage
        if stage.is_current_stage_first:
            self.embed = nn.Embedding(...)

        # Only build head on last stage
        if stage.is_current_stage_last:
            self.head = nn.Linear(...)

    def forward(self, inputs: MyInput | MyTransfer, shared: MyShared) -> MyTransfer | MyOutput:
        # Branch on the stage position, never on which argument is None.
        if self.stage.is_current_stage_first:
            x = self.embed(inputs.input_ids)
        else:
            x = inputs.hidden_states

        # Run local layers
        for layer_idx in range(self.start_layer, self.end_layer):
            x = self.layers[str(layer_idx)](x)

        # Last stage produces the PipelineOutput; everyone else produces a StageTransfer.
        if self.stage.is_current_stage_last:
            return MyOutput(logits=self.head(x))
        return MyTransfer(hidden_states=x)

    # --- Protocol Implementation ---
    # Receives a single representative microbatch of PipelineInput (not a global batch), so shapes are
    # used as-is, and returns a StageTransfer of TensorSpec descriptors — no tensors are allocated.
    # The engine never calls this for the first stage's `incoming` nor the last stage's `outgoing`
    # edge, so terminal boundaries need no special-casing.
    def stage_transfer_spec(self, pipeline_input: MyInput, boundary: StageBoundary) -> MyTransfer:
        micro_batch_size, seq_len = pipeline_input.input_ids.shape
        return MyTransfer(
            hidden_states=TensorSpec(
                shape=(micro_batch_size, seq_len, self.config.hidden_dim), dtype=torch.bfloat16
            )
        )

from torch import Tensor
import torch.nn.functional as F

from d9d.core.dist_context import DistributedContext
from d9d.pipelining.factory import build_schedule, PipelineSchedule1F1BConfig


# 0. Define an object that manages loss calculation per microbatch. It receives the PipelineOutput
#    produced by the last stage and reads its fields by attribute.
class MyLossHandler:
    def __init__(self, targets_microbatches: list[Tensor]):
        self._targets = targets_microbatches

    def compute_loss(self, outputs: MyOutput, microbatch_idx: int):
        # Implement any custom logic here
        current_target = self._targets[microbatch_idx]
        return F.cross_entropy(outputs.logits.view(-1, outputs.logits.shape[-1]), current_target.view(-1))


# 1. Define configuration
dist_context: DistributedContext = ...
model_config = ...
schedule_config = PipelineSchedule1F1BConfig(
    num_stages_per_rank=4,  # 4 Virtual stages per rank
    zero_bubble=True  # Enable ZB1P optimization
)

# 2. Build the per-microbatch inputs (the "pack"). The number of microbatches is decided here, per
#    step. Each stage receives one PipelineInput and one SharedInput per microbatch.
inputs_microbatches = tuple(MyInput(input_ids=mb) for mb in my_input_microbatches)
shared_microbatches = tuple(MyShared(position_ids=pos) for pos in my_position_microbatches)
targets_microbatches = [...]  # one target tensor per microbatch

# 3. Build the schedule and model shards (the callback is supplied per step, not here)
loss_handler = MyLossHandler(targets_microbatches)
schedule_info, modules = build_schedule(
    dist_context=dist_context,
    schedule_config=schedule_config,
    model_provider=lambda stage: MyModelChunk(stage, model_config),  # Factory function
)

# 4. Execution
# The callback is passed to each step, so it can close over per-step data (here, the targets).
schedule_info.schedule.step(
    inputs_microbatches=inputs_microbatches,
    shared_microbatches=shared_microbatches,
    callback=loss_handler.compute_loss,
)