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 (passed 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 which 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 (ModuleDict keys keep the global layer index in parameter names).
        self.layers = nn.ModuleDict({
            str(layer): TransformerBlock(...)
            for layer in range(self.start_layer, self.end_layer)
        })

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

        # Only build the head on the 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)

        # The last stage produces the PipelineOutput; every other stage produces a StageTransfer.
        if self.stage.is_current_stage_last:
            return MyOutput(logits=self.head(x))
        return MyTransfer(hidden_states=x)

    # --- Protocol Implementation ---
    # Receives one microbatch of PipelineInput (not a global batch), so shapes are used as they are.
    # Returns a StageTransfer of TensorSpec descriptors and allocates no tensors.
    # The engine never calls this for the `incoming` edge of the first stage or the `outgoing` edge
    # of the last stage, so these edges need no special case.
    def stage_transfer_spec(self, pipeline_input: MyInput, boundary: StageBoundary) -> MyTransfer:
        batch_size, seq_len = pipeline_input.input_ids.shape
        return MyTransfer(
            hidden_states=TensorSpec(
                shape=(batch_size, seq_len, self.config.hidden_size), dtype=torch.bfloat16
            )
        )

pipelining:
  schedule:
    schedule: 1f1b
    num_stages_per_rank: 2
    zero_bubble: true

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 computes the loss 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 the configuration.
dist_context: DistributedContext = ...
model_config = ...
schedule_config = PipelineSchedule1F1BConfig(
    num_stages_per_rank=4,  # 4 virtual stages per rank
    zero_bubble=True  # Use the ZB1P variant
)

# 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 the model stages (the callback is passed 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. Run the step.
# The callback is passed to each step, so it can use per-step data (here, the targets).
schedule_info.schedule.step(
    inputs_microbatches=inputs_microbatches,
    shared_microbatches=shared_microbatches,
    callback=loss_handler.compute_loss,
)