d9d 0.1.0__py3-none-any.whl

d9d/pipelining/factory/config.py

@@ -0,0 +1,89 @@
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+
+class PipelineScheduleInferenceConfig(BaseModel):
+    """
+    Configuration for inference-only pipeline execution.
+
+    This schedule runs all forward passes sequentially without any backward passes.
+    """
+
+    schedule: Literal["inference"] = "inference"
+
+
+class PipelineScheduleGPipeConfig(BaseModel):
+    """
+    Configuration for GPipe execution.
+
+    This assumes a single stage per rank and processes all microbatches for the
+    forward pass before switching to the backward pass.
+    """
+
+    schedule: Literal["gpipe"] = "gpipe"
+
+
+class PipelineScheduleLoopedBFSConfig(BaseModel):
+    """
+    Configuration for Looped Breadth-First Search execution.
+
+    Similar to GPipe, but supports multiple stages per rank (virtualization).
+    It executes all available work for a specific stage before moving to the next.
+    """
+
+    schedule: Literal["looped_bfs"] = "looped_bfs"
+
+    num_stages_per_rank: int
+
+
+class PipelineSchedule1F1BConfig(BaseModel):
+    """
+    Configuration for Interleaved 1F1B and Interleaved Zero Bubble execution.
+
+    Supports assigning multiple stages per rank and sharding backward to dI and dW
+    to reduce pipeline bubbles.
+    """
+
+    schedule: Literal["1f1b"] = "1f1b"
+
+    num_stages_per_rank: int
+    zero_bubble: bool
+
+
+class PipelineScheduleZeroBubbleVConfig(BaseModel):
+    """
+    Configuration for Zero Bubble V (ZBV) execution.
+
+    A specialized V-shape topology schedule that splits backward passes into
+    Input and Weight gradients to maximize overlap. Requires exactly 2 stages per rank.
+    """
+    schedule: Literal["zero_bubble_v"] = "zero_bubble_v"
+
+
+class PipelineScheduleDualPipeVConfig(BaseModel):
+    """
+    Configuration for DualPipeV execution.
+
+    A bidirectional pipeline schedule for high-throughput training, utilizing
+    V-shape topology and reciprocal forward/backward scheduling.
+    """
+
+    schedule: Literal["dual_pipe_v"] = "dual_pipe_v"
+
+
+AnyPipelineScheduleConfig = Annotated[
+    PipelineScheduleInferenceConfig |
+    PipelineScheduleGPipeConfig |
+    PipelineScheduleLoopedBFSConfig |
+    PipelineSchedule1F1BConfig |
+    PipelineScheduleZeroBubbleVConfig |
+    PipelineScheduleDualPipeVConfig,
+    Field(discriminator="schedule")
+]
+"""Union of all supported pipeline schedule configuration types.
+
+This type alias uses a Pydantic discriminator on the ``schedule`` field to allow
+polymorphic validation and serialization of specific schedule configs (e.g.
+Inference, GPipe, 1F1B, ZeroBubble, etc.).
+"""

d9d/pipelining/factory/factory.py

@@ -0,0 +1,114 @@
+import dataclasses
+from collections.abc import Callable
+
+import torch
+from torch import nn
+
+from ...core.dist_context import REGULAR_DOMAIN, DistributedContext
+from ..api import PipelineSchedule, PipelineStageInfo
+from ..infra.schedule.component.program import (
+    build_stage_to_host_rank_topology,
+    invert_stage_to_host_rank_topology,
+)
+from ..infra.schedule.component.runtime import PipelineScheduleExecutor
+from ..infra.stage import PipelineStage
+from .config import (
+    AnyPipelineScheduleConfig,
+)
+from .registry import PIPELINE_PROGRAM_REGISTRY
+
+
+@dataclasses.dataclass(kw_only=True)
+class PipelineScheduleInfo:
+    """Contains the built pipeline schedule and rank-specific metadata."""
+
+    schedule: PipelineSchedule
+    has_first_stage: bool
+    has_last_stage: bool
+
+
+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,
+) -> 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).
+
+    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
+    )
+
+    return PipelineScheduleInfo(
+        schedule=schedule,
+        has_first_stage=has_first_stage,
+        has_last_stage=has_last_stage
+    ), modules

d9d/pipelining/factory/registry.py

@@ -0,0 +1,82 @@
+from collections.abc import Callable
+from typing import TypeVar, cast
+
+from d9d.pipelining.factory import (
+    AnyPipelineScheduleConfig,
+    PipelineSchedule1F1BConfig,
+    PipelineScheduleDualPipeVConfig,
+    PipelineScheduleGPipeConfig,
+    PipelineScheduleInferenceConfig,
+    PipelineScheduleLoopedBFSConfig,
+    PipelineScheduleZeroBubbleVConfig,
+)
+from d9d.pipelining.infra.schedule.component.program import PipelineProgramBuilder
+from d9d.pipelining.infra.schedule.program import (
+    DualPipeVPipelineProgramBuilder,
+    Interleaved1F1BPipelineProgramBuilder,
+    LoopedBFSPipelineProgramBuilder,
+    ZeroBubbleVPipelineProgramBuilder,
+)
+
+TConfig = TypeVar("TConfig", bound=AnyPipelineScheduleConfig)
+
+TRegistryDict = dict[
+    type[AnyPipelineScheduleConfig],
+    Callable[[AnyPipelineScheduleConfig], PipelineProgramBuilder]
+]
+
+TBoundRegistryFn = Callable[[TConfig], PipelineProgramBuilder]
+
+
+class PipelineProgramRegistry:
+    def __init__(self) -> None:
+        self._registry: TRegistryDict = {}
+
+    def register_program(
+            self, config_cls: type[TConfig]
+    ) -> Callable[[TBoundRegistryFn], TBoundRegistryFn]:
+        def decorator(func: TBoundRegistryFn) -> TBoundRegistryFn:
+            config_cls_any = cast(type[AnyPipelineScheduleConfig], config_cls)
+            self._registry[config_cls_any] = func
+            return func
+
+        return decorator
+
+    def program_for(self, config: AnyPipelineScheduleConfig) -> PipelineProgramBuilder:
+        program_fn = self._registry[type(config)]
+        program = program_fn(config)
+        return program
+
+
+PIPELINE_PROGRAM_REGISTRY = PipelineProgramRegistry()
+
+
+@PIPELINE_PROGRAM_REGISTRY.register_program(PipelineScheduleGPipeConfig)
+def _build_gpipe(_: PipelineScheduleGPipeConfig) -> PipelineProgramBuilder:
+    return LoopedBFSPipelineProgramBuilder(num_stages_per_rank=1, inference_mode=False)
+
+
+@PIPELINE_PROGRAM_REGISTRY.register_program(PipelineScheduleInferenceConfig)
+def _build_inference(_: PipelineScheduleInferenceConfig) -> PipelineProgramBuilder:
+    return LoopedBFSPipelineProgramBuilder(num_stages_per_rank=1, inference_mode=True)
+
+
+@PIPELINE_PROGRAM_REGISTRY.register_program(PipelineScheduleLoopedBFSConfig)
+def _build_looped_bfs(cfg: PipelineScheduleLoopedBFSConfig) -> PipelineProgramBuilder:
+    return LoopedBFSPipelineProgramBuilder(num_stages_per_rank=cfg.num_stages_per_rank, inference_mode=False)
+
+
+@PIPELINE_PROGRAM_REGISTRY.register_program(PipelineSchedule1F1BConfig)
+def _build_1f1b(cfg: PipelineSchedule1F1BConfig) -> PipelineProgramBuilder:
+    return Interleaved1F1BPipelineProgramBuilder(num_stages_per_rank=cfg.num_stages_per_rank,
+                                                 enable_zero_bubble=cfg.zero_bubble)
+
+
+@PIPELINE_PROGRAM_REGISTRY.register_program(PipelineScheduleDualPipeVConfig)
+def _build_dual_pipe_v(_: PipelineScheduleDualPipeVConfig) -> PipelineProgramBuilder:
+    return DualPipeVPipelineProgramBuilder()
+
+
+@PIPELINE_PROGRAM_REGISTRY.register_program(PipelineScheduleZeroBubbleVConfig)
+def _build_zero_bubble_v(_: PipelineScheduleZeroBubbleVConfig) -> PipelineProgramBuilder:
+    return ZeroBubbleVPipelineProgramBuilder()

d9d/pipelining/infra/schedule/component/program/__init__.py

@@ -0,0 +1,22 @@
+"""
+Pipeline Schedule Building Components.
+
+This package provides the core building blocks and compiler passes used to generate
+execution schedules for distributed pipelines.
+"""
+
+from .base import PipelineProgramBuilder
+from .communications import add_communication_ops
+from .topology import (
+    ScheduleStyle,
+    build_stage_to_host_rank_topology,
+    invert_stage_to_host_rank_topology,
+)
+
+__all__ = [
+    "PipelineProgramBuilder",
+    "ScheduleStyle",
+    "add_communication_ops",
+    "build_stage_to_host_rank_topology",
+    "invert_stage_to_host_rank_topology"
+]

d9d/pipelining/infra/schedule/component/program/base.py

@@ -0,0 +1,35 @@
+import abc
+
+from ..program.topology import ScheduleStyle
+from ..runtime import ActionBase
+
+
+class PipelineProgramBuilder(abc.ABC):
+    """Abstract interface for building pipeline execution schedules."""
+
+    @abc.abstractmethod
+    def compose(self, num_microbatches: int, pp_size: int) -> dict[int, list[ActionBase]]:
+        """
+        Generates the execution program for all ranks in the pipeline.
+
+        Args:
+            num_microbatches: Number of microbatches per step.
+            pp_size: Number of pipeline parallel ranks.
+
+        Returns:
+            A dictionary mapping rank indices to their list of sequential actions.
+        """
+        ...
+
+    @property
+    @abc.abstractmethod
+    def num_stages_per_rank(self) -> int:
+        """Returns the number of model stages designated for each rank."""
+
+        ...
+
+    @property
+    @abc.abstractmethod
+    def topology_style(self) -> ScheduleStyle:
+        """Returns the topology style strategy used to assign stages to ranks."""
+        ...

d9d/pipelining/infra/schedule/component/program/communications.py

@@ -0,0 +1,203 @@
+import copy
+import dataclasses
+
+from ..runtime.action import (
+    ActionBase,
+    BackwardFullInputComputeAction,
+    BackwardReceiveAction,
+    BackwardSendAction,
+    ComposeAction,
+    ForwardComputeAction,
+    ForwardReceiveAction,
+    ForwardSendAction,
+)
+
+
+def _get_sub_actions(action: ActionBase) -> tuple[ActionBase, ...]:
+    if isinstance(action, ComposeAction):
+        return action.actions
+    return (action,)
+
+
+def _check_action_communication_dependencies_fulfilled(
+        action: ActionBase,
+        rank_events: set[ActionBase],
+        num_stages: int
+) -> bool:
+    match action:
+        case ForwardComputeAction():
+            if action.stage_idx == 0:
+                return True
+            if ForwardReceiveAction(action.stage_idx, action.microbatch_idx) in rank_events:
+                return True
+            if ForwardComputeAction(action.stage_idx - 1, action.microbatch_idx) in rank_events:
+                return True
+            return False
+        case BackwardFullInputComputeAction():
+            if action.stage_idx == num_stages - 1:
+                return True
+            if BackwardReceiveAction(action.stage_idx, action.microbatch_idx) in rank_events:
+                return True
+
+            next_full = BackwardFullInputComputeAction(
+                action.stage_idx + 1,
+                action.microbatch_idx,
+                full_backward=True
+            )
+            next_inp = BackwardFullInputComputeAction(
+                action.stage_idx + 1,
+                action.microbatch_idx,
+                full_backward=False
+            )
+
+            if next_full in rank_events or next_inp in rank_events:
+                return True
+            return False
+        case _:
+            return True
+
+
+def check_action_communication_dependencies_fulfilled(
+        action: ActionBase,
+        rank_events: set[ActionBase],
+        num_stages: int
+) -> bool:
+    """
+    Checks if data dependencies (Receive or Local Compute) are met for an action.
+
+    This function determines if a compute action is allowed to run based on
+    whether its inputs are available in `rank_events`. Inputs are available
+    if they were either computed locally by a previous stage or received
+    from a remote rank.
+
+    Args:
+        action: The action to check.
+        rank_events: A set of actions already completed on this rank.
+        num_stages: Total number of stages in the pipeline.
+
+    Returns:
+        True if all dependencies are satisfied, False otherwise.
+    """
+
+    return all(
+        _check_action_communication_dependencies_fulfilled(sub, rank_events, num_stages)
+        for sub in _get_sub_actions(action)
+    )
+
+
+@dataclasses.dataclass(kw_only=True)
+class _CommunicationPackage:
+    send: ActionBase
+    recv: ActionBase
+    sends_to_rank: int
+
+
+def _create_communications_for_action(
+        action: ActionBase,
+        num_stages: int,
+        stage_to_rank: dict[int, int],
+) -> _CommunicationPackage | None:
+    match action:
+        case ForwardComputeAction():
+            if action.stage_idx == num_stages - 1:
+                return None
+
+            curr_rank, next_rank = stage_to_rank[action.stage_idx], stage_to_rank[action.stage_idx + 1]
+            if curr_rank == next_rank:
+                return None
+
+            return _CommunicationPackage(
+                send=ForwardSendAction(action.stage_idx, action.microbatch_idx),
+                recv=ForwardReceiveAction(action.stage_idx + 1, action.microbatch_idx),
+                sends_to_rank=next_rank
+            )
+        case BackwardFullInputComputeAction():
+            if action.stage_idx == 0:
+                return None
+
+            curr_rank, prev_rank = stage_to_rank[action.stage_idx], stage_to_rank[action.stage_idx - 1]
+            if curr_rank == prev_rank:
+                return None
+
+            return _CommunicationPackage(
+                send=BackwardSendAction(action.stage_idx, action.microbatch_idx),
+                recv=BackwardReceiveAction(action.stage_idx - 1, action.microbatch_idx),
+                sends_to_rank=prev_rank
+            )
+        case _:
+            return None
+
+
+def add_communication_ops(
+        compute_actions: dict[int, list[ActionBase]],
+        stage_to_rank: dict[int, int],
+        num_stages: int,
+) -> dict[int, list[ActionBase]]:
+    """
+    Injects communication actions into a computation-only schedule.
+
+    This function iterates through the provided compute schedule and simulates execution.
+    When a compute action produces a result needed by a different rank, it injects
+    Send/Receive pairs. It also reorders actions to ensure that Receive
+    operations occur before the Computes that depend on them, preventing deadlocks.
+
+    Args:
+        compute_actions: Initial schedule containing only compute operations.
+        stage_to_rank: Mapping from stage index to rank index.
+        num_stages: Total number of pipeline stages.
+
+    Returns:
+        A new schedule dictionary including both compute and communication actions.
+
+    Raises:
+        RuntimeError: If the schedule simulation enters a deadlock state.
+    """
+
+    compute_actions = copy.deepcopy(compute_actions)
+
+    full_actions: dict[int, list[ActionBase]] = {rank: [] for rank in compute_actions}
+    completed_events: dict[int, set[ActionBase]] = {rank: set() for rank in compute_actions}
+
+    while compute_actions:
+        progress = False
+
+        for rank in sorted(compute_actions.keys()):
+            if not compute_actions[rank]:
+                del compute_actions[rank]
+                continue
+
+            current_action = compute_actions[rank][0]
+            sub_actions = _get_sub_actions(current_action)
+
+            # Check readiness
+            if not check_action_communication_dependencies_fulfilled(
+                    current_action, completed_events[rank], num_stages
+            ):
+                continue
+
+            # Execute
+            full_actions[rank].append(current_action)
+            compute_actions[rank].pop(0)
+            progress = True
+
+            for sub_action in sub_actions:
+                completed_events[rank].add(sub_action)
+
+                comm_pkg = _create_communications_for_action(
+                    sub_action,
+                    num_stages=num_stages,
+                    stage_to_rank=stage_to_rank
+                )
+                if comm_pkg:
+                    # Add Send locally
+                    full_actions[rank].append(comm_pkg.send)
+                    completed_events[rank].add(comm_pkg.send)
+
+                    # Add Recv remotely and unblock target
+                    full_actions[comm_pkg.sends_to_rank].append(comm_pkg.recv)
+                    completed_events[comm_pkg.sends_to_rank].add(comm_pkg.recv)
+
+        if not progress and compute_actions:
+            raise RuntimeError("Deadlock in schedule simulation")
+
+    return full_actions

d9d/pipelining/infra/schedule/component/program/topology.py

@@ -0,0 +1,78 @@
+from collections import defaultdict
+from enum import StrEnum
+
+
+class ScheduleStyle(StrEnum):
+    """
+    Defines the strategy for mapping logical stages to physical ranks.
+
+    Attributes:
+        loop: Assigns stages in a round-robin circular fashion (mod pp_size).
+        v: Assigns stages in a zig-zag V-shape pattern. Useful for interleaved 1F1B schedules.
+    """
+
+    loop = "loop"
+    v = "v"
+
+
+def build_stage_to_host_rank_topology(
+    pp_size: int, num_stages: int, style: ScheduleStyle
+) -> dict[int, int]:
+    """
+    Constructs the mapping from stage index to rank index.
+
+    Args:
+        pp_size: Number of pipeline parallel ranks.
+        num_stages: Total number of model stages.
+        style: The topology style to use for assignment.
+
+    Returns:
+        A dictionary mapping stage IDs to Rank IDs.
+
+    Raises:
+        ValueError: If the style is unknown or if V-style parameters are invalid
+            (num_stages must be divisible by pp_size).
+    """
+
+    match style:
+        case ScheduleStyle.loop:
+            return {stage_index: stage_index % pp_size for stage_index in range(num_stages)}
+        case ScheduleStyle.v:
+            if num_stages % pp_size != 0:
+                raise ValueError(
+                    f"num_stages {num_stages} must be evenly divisible by pp_size {pp_size} for V schedules"
+                )
+
+            result = {}
+            rank_index = 0
+            for stage_index in range(num_stages):
+                result[stage_index] = rank_index
+                if (stage_index + 1) % pp_size == 0:
+                    continue
+                if (stage_index // pp_size) % 2 == 0:
+                    rank_index += 1
+                else:
+                    rank_index -= 1
+            return result
+        case _:
+            raise ValueError()
+
+
+def invert_stage_to_host_rank_topology(
+        stage_to_host: dict[int, int]
+) -> dict[int, list[int]]:
+    """
+    Inverts the topology mapping to list execution stages per rank.
+
+    Args:
+        stage_to_host: Mapping from stage index to rank index.
+
+    Returns:
+        A dictionary where keys are Rank IDs and values are lists of Stage IDs
+        managed by that rank.
+    """
+
+    host_to_stage = defaultdict(list)
+    for stage_idx, host in stage_to_host.items():
+        host_to_stage[host].append(stage_idx)
+    return dict(host_to_stage)

d9d/pipelining/infra/schedule/component/runtime/__init__.py

@@ -0,0 +1,29 @@
+"""
+Pipelining Runtime Package.
+"""
+
+from .action import (
+    ActionBase,
+    BackwardFullInputComputeAction,
+    BackwardReceiveAction,
+    BackwardSendAction,
+    BackwardWeightComputeAction,
+    ComposeAction,
+    ForwardComputeAction,
+    ForwardReceiveAction,
+    ForwardSendAction,
+)
+from .executor import PipelineScheduleExecutor
+
+__all__ = [
+    "ActionBase",
+    "BackwardFullInputComputeAction",
+    "BackwardReceiveAction",
+    "BackwardSendAction",
+    "BackwardWeightComputeAction",
+    "ComposeAction",
+    "ForwardComputeAction",
+    "ForwardReceiveAction",
+    "ForwardSendAction",
+    "PipelineScheduleExecutor",
+]