d9d 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,361 @@
+import abc
+import dataclasses
+from enum import StrEnum
+from typing import Any
+
+import torch
+
+from d9d.pipelining.infra.stage import PipelineStage
+
+from .communications import PipelineCommunicationHandler
+from .loss import PipelineLossHandler
+
+
+@dataclasses.dataclass(kw_only=True, slots=True)
+class ActionContext:
+    """
+    Holds the runtime context required to execute a pipeline action.
+
+    Attributes:
+        pipeline_inputs_microbatches: The global inputs sharded by microbatch.
+        pipeline_kwargs_microbatches: The global keyword arguments sharded by microbatch.
+        stages: A mapping of stage indices to their active PipelineStage instances.
+        communications: The handler for P2P communications.
+        loss: The handler for loss computation, or None if not available.
+    """
+
+    pipeline_inputs_microbatches: tuple[dict[str, torch.Tensor], ...]
+    pipeline_kwargs_microbatches: tuple[dict[str, Any], ...]
+
+    stages: dict[int, PipelineStage]
+    communications: PipelineCommunicationHandler
+    loss: PipelineLossHandler | None
+
+
+class ActionWorkType(StrEnum):
+    """
+    Classifies the type of work performed by an action.
+
+    Attributes:
+        compute: Indicates the action involves computation components (forward, backward).
+        communicate: Indicates the action involves network I/O components (send, receive).
+    """
+
+    compute = "compute"
+    communicate = "communicate"
+
+
+class ActionBase(abc.ABC):
+    """
+    Abstract base class for all pipeline schedule actions.
+
+    An action represents an atomic unit of work in a pipeline schedule,
+    such as computing a microbatch or sending/receiving a tensor.
+    """
+
+    @abc.abstractmethod
+    def apply(self, ctx: ActionContext):
+        """
+        Executes the action logic using the provided context.
+
+        Args:
+            ctx: The runtime context containing stages, data, and communication handlers.
+        """
+
+        ...
+
+    @property
+    @abc.abstractmethod
+    def work_type(self) -> ActionWorkType:
+        """Returns the classification of work this action performs."""
+        ...
+
+    @property
+    @abc.abstractmethod
+    def has_backward_work(self) -> bool:
+        """Returns True if this action involves backward pass computations."""
+        ...
+
+    @abc.abstractmethod
+    def __str__(self) -> str:
+        """Returns a short string representation of the action for logging/visualization."""
+        ...
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class ForwardSendAction(ActionBase):
+    """
+    Action to schedule a forward pass tensor send operation.
+
+    Attributes:
+        stage_idx: The integer index of the pipeline stage initiating the send operation.
+        microbatch_idx: The integer index of the microbatch being sent.
+    """
+
+    stage_idx: int
+    microbatch_idx: int
+
+    def apply(self, ctx: ActionContext):
+        ctx.communications.schedule_fwd_send(self.stage_idx, self.microbatch_idx)
+
+    @property
+    def work_type(self) -> ActionWorkType:
+        return ActionWorkType.communicate
+
+    @property
+    def has_backward_work(self) -> bool:
+        return False
+
+    def __str__(self) -> str:
+        return f"{self.stage_idx}SEND_F{self.microbatch_idx}"
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class BackwardSendAction(ActionBase):
+    """
+    Action to schedule a backward pass gradient send operation.
+
+    Attributes:
+        stage_idx: The integer index of the pipeline stage initiating the send operation.
+        microbatch_idx: The integer index of the microbatch being sent.
+    """
+
+    stage_idx: int
+    microbatch_idx: int
+
+    def apply(self, ctx: ActionContext):
+        ctx.communications.schedule_bwd_send(self.stage_idx, self.microbatch_idx)
+
+    @property
+    def work_type(self) -> ActionWorkType:
+        return ActionWorkType.communicate
+
+    @property
+    def has_backward_work(self) -> bool:
+        return True
+
+    def __str__(self) -> str:
+        return f"{self.stage_idx}SEND_B{self.microbatch_idx}"
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class ForwardReceiveAction(ActionBase):
+    """
+    Action to schedule a forward pass tensor receive operation.
+
+    Attributes:
+        stage_idx: The integer index of the pipeline stage expecting the receive operation.
+        microbatch_idx: The integer index of the microbatch being received.
+    """
+
+    stage_idx: int
+    microbatch_idx: int
+
+    def apply(self, ctx: ActionContext):
+        ctx.communications.schedule_fwd_recv(self.stage_idx, self.microbatch_idx)
+
+    @property
+    def work_type(self) -> ActionWorkType:
+        return ActionWorkType.communicate
+
+    @property
+    def has_backward_work(self) -> bool:
+        return True
+
+    def __str__(self) -> str:
+        return f"{self.stage_idx}RECV_F{self.microbatch_idx}"
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class BackwardReceiveAction(ActionBase):
+    """
+    Action to schedule a backward pass gradient receive operation.
+
+    Attributes:
+        stage_idx: The integer index of the pipeline stage expecting the receive operation.
+        microbatch_idx: The integer index of the microbatch being received.
+    """
+
+    stage_idx: int
+    microbatch_idx: int
+
+    def apply(self, ctx: ActionContext):
+        ctx.communications.schedule_bwd_recv(self.stage_idx, self.microbatch_idx)
+
+    @property
+    def work_type(self) -> ActionWorkType:
+        return ActionWorkType.communicate
+
+    @property
+    def has_backward_work(self) -> bool:
+        return True
+
+    def __str__(self) -> str:
+        return f"{self.stage_idx}RECV_B{self.microbatch_idx}"
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class ForwardComputeAction(ActionBase):
+    """
+    Action to perform forward computation for a specific microbatch.
+
+    Attributes:
+        stage_idx: The integer index of the pipeline stage.
+        microbatch_idx: The integer index of the microbatch to compute.
+    """
+
+    stage_idx: int
+    microbatch_idx: int
+
+    def apply(self, ctx: ActionContext):
+        # todo check unsharded
+        stage = ctx.stages[self.stage_idx]
+
+        if not stage.info.is_current_stage_first and self.stage_idx - 1 not in ctx.stages:
+            ctx.communications.wait_fwd_recv(self.stage_idx, self.microbatch_idx)
+
+        stage.forward_one_chunk(
+            microbatch_index=self.microbatch_idx,
+            pipeline_inputs=ctx.pipeline_inputs_microbatches[self.microbatch_idx],
+            pipeline_kwargs=ctx.pipeline_kwargs_microbatches[self.microbatch_idx]
+        )
+        result = stage.get_local_fwd_output(self.microbatch_idx)
+
+        if stage.info.is_current_stage_last and ctx.loss is not None:
+            ctx.loss.compute_loss(result, self.microbatch_idx)
+
+        if not stage.info.is_current_stage_last and self.stage_idx + 1 in ctx.stages:
+            ctx.stages[self.stage_idx + 1].set_local_fwd_input(
+                inputs=result,
+                microbatch_index=self.microbatch_idx
+            )
+
+    @property
+    def work_type(self) -> ActionWorkType:
+        return ActionWorkType.compute
+
+    @property
+    def has_backward_work(self) -> bool:
+        return False
+
+    def __str__(self) -> str:
+        return f"{self.stage_idx}F{self.microbatch_idx}"
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class BackwardFullInputComputeAction(ActionBase):
+    """
+    Action to perform backward computation with respect to inputs.
+
+    Attributes:
+        stage_idx: The integer index of the pipeline stage.
+        microbatch_idx: The integer index of the microbatch to compute.
+        full_backward: If True, performs a full backward pass including inputs
+            and weights. If False, may only compute gradients w.r.t inputs
+            (depending on schedule implementation).
+    """
+
+    stage_idx: int
+    microbatch_idx: int
+    full_backward: bool
+
+    def apply(self, ctx: ActionContext):
+        # todo unshard
+        stage = ctx.stages[self.stage_idx]
+
+        if not stage.info.is_current_stage_last and self.stage_idx + 1 not in ctx.stages:
+            ctx.communications.wait_bwd_recv(self.stage_idx, self.microbatch_idx)
+
+        if stage.info.is_current_stage_last and ctx.loss is not None:
+            loss = ctx.loss.acquire_loss(self.microbatch_idx)
+        else:
+            loss = None
+
+        stage.backward_one_chunk(
+            microbatch_index=self.microbatch_idx,
+            full_backward=self.full_backward,
+            loss=loss
+        )
+
+        if not stage.info.is_current_stage_first and self.stage_idx - 1 in ctx.stages:
+            ctx.stages[self.stage_idx - 1].set_local_bwd_input(
+                microbatch_index=self.microbatch_idx,
+                inputs=stage.pop_local_bwd_output(self.microbatch_idx)
+            )
+
+    @property
+    def work_type(self) -> ActionWorkType:
+        return ActionWorkType.compute
+
+    @property
+    def has_backward_work(self) -> bool:
+        return True
+
+    def __str__(self) -> str:
+        letter = "B" if self.full_backward else "I"
+        return f"{self.stage_idx}{letter}{self.microbatch_idx}"
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class BackwardWeightComputeAction(ActionBase):
+    """
+    Action to perform gradient accumulation on weights.
+
+    Attributes:
+        stage_idx: The integer index of the pipeline stage.
+        microbatch_idx: The integer index of the microbatch to compute.
+    """
+
+    stage_idx: int
+    microbatch_idx: int
+
+    def apply(self, ctx: ActionContext):
+        # todo unshard
+        stage = ctx.stages[self.stage_idx]
+
+        stage.backward_weight_one_chunk(
+            microbatch_index=self.microbatch_idx
+        )
+
+    @property
+    def work_type(self) -> ActionWorkType:
+        return ActionWorkType.compute
+
+    @property
+    def has_backward_work(self) -> bool:
+        return True
+
+    def __str__(self) -> str:
+        return f"{self.stage_idx}W{self.microbatch_idx}"
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class ComposeAction(ActionBase):
+    """
+    Composite action scheduling multiple sub-actions sequentially.
+
+    Used for forward/backward overlapping.
+
+    Attributes:
+        actions: A tuple of sub-actions to be executed sequentially.
+    """
+
+    actions: tuple[ActionBase, ...]
+
+    def apply(self, ctx: ActionContext):
+        for act in self.actions:
+            act.apply(ctx)
+
+    @property
+    def work_type(self) -> ActionWorkType:
+        sub_work_types = {x.work_type for x in self.actions}
+        if len(sub_work_types) != 1:
+            raise ValueError("")
+        return next(iter(sub_work_types))
+
+    @property
+    def has_backward_work(self) -> bool:
+        return any(x.has_backward_work for x in self.actions)
+
+    def __str__(self) -> str:
+        return "|".join(map(str, self.actions))

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

@@ -0,0 +1,101 @@
+import torch.distributed as dist
+
+from d9d.pipelining.infra.stage import PipelineStage
+
+
+def _schedule_batched_p2p(ops: list[dist.P2POp]) -> list[dist.Work]:
+    if not len(ops):
+        return []
+
+    return dist.batch_isend_irecv(ops)
+
+
+def _wait_batched_p2p(work: list[dist.Work]):
+    for work_item in work:
+        work_item.wait()
+
+
+class PipelineCommunicationHandler:
+    """Manages point-to-point communications between pipeline stages."""
+
+    def __init__(self, stages: dict[int, PipelineStage]):
+        """
+        Constructs the communication handler.
+
+        Args:
+            stages: Mapping of stage indices to PipelineStage instances.
+        """
+
+        self._stages = stages
+
+        self._forward_receive_ops: dict[tuple[int, int], list[dist.Work]] = {}
+        self._backward_receive_ops: dict[tuple[int, int], list[dist.Work]] = {}
+
+        self._send_ops: list[list[dist.Work]] = []
+
+    def schedule_fwd_send(self, stage_idx: int, microbatch_idx: int):
+        """Schedules non-blocking connection to send forward pass outputs."""
+
+        stage = self._stages[stage_idx]
+        work = _schedule_batched_p2p(stage.get_fwd_send_ops(microbatch_idx))
+        self._send_ops.append(work)
+
+    def schedule_bwd_send(self, stage_idx: int, microbatch_idx: int):
+        """Schedules non-blocking connection to send backward pass outputs."""
+
+        stage = self._stages[stage_idx]
+        work = _schedule_batched_p2p(stage.get_bwd_send_ops(microbatch_idx))
+        self._send_ops.append(work)
+
+    def schedule_fwd_recv(self, stage_idx: int, microbatch_idx: int):
+        """
+        Schedules non-blocking connection to receive forward pass inputs.
+
+        Raises:
+            ValueError: If a receive op is already pending for this stage/microbatch.
+        """
+        stage = self._stages[stage_idx]
+        key = (stage_idx, microbatch_idx)
+
+        if key in self._forward_receive_ops:
+            raise ValueError()
+
+        work = _schedule_batched_p2p(stage.get_fwd_recv_ops(microbatch_idx))
+        self._forward_receive_ops[key] = work
+
+    def wait_fwd_recv(self, stage_idx: int, microbatch_idx: int):
+        """Blocks until the forward pass receive operation completes."""
+        key = (stage_idx, microbatch_idx)
+        _wait_batched_p2p(self._forward_receive_ops.pop(key))
+
+    def schedule_bwd_recv(self, stage_idx: int, microbatch_idx: int):
+        """
+        Schedules non-blocking connection to receive backward pass inputs.
+
+        Raises:
+            ValueError: If a receive op is already pending for this stage/microbatch.
+        """
+
+        stage = self._stages[stage_idx]
+        key = (stage_idx, microbatch_idx)
+
+        if key in self._backward_receive_ops:
+            raise ValueError()
+
+        work = _schedule_batched_p2p(stage.get_bwd_recv_ops(microbatch_idx))
+
+        self._backward_receive_ops[key] = work
+
+    def wait_bwd_recv(self, stage_idx: int, microbatch_idx: int):
+        """Blocks until the backward pass receive operation completes."""
+
+        key = (stage_idx, microbatch_idx)
+        _wait_batched_p2p(self._backward_receive_ops.pop(key))
+
+    def wait_send_all(self):
+        """Blocks until all pending send operations are completed."""
+
+        while self._send_ops:
+            ops = self._send_ops.pop()
+            for op in ops:
+                op.wait()

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

@@ -0,0 +1,113 @@
+from typing import Any
+
+import torch
+from torch.autograd.profiler import record_function
+
+from d9d.core.dist_context import REGULAR_DOMAIN, DistributedContext
+from d9d.core.sharding import ShardingSpec, shard_spec_on_dim, shard_tree
+from d9d.pipelining.api import PipelineSchedule, PipelineShardingSpec
+from d9d.pipelining.infra.stage import PipelineStage
+
+from .action import ActionBase, ActionContext
+from .communications import PipelineCommunicationHandler
+from .loss import LossFn, PipelineLossHandler
+
+
+class PipelineScheduleExecutor(PipelineSchedule):
+    """Executes a defined pipeline schedule by interpreting a sequence of actions."""
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            stages: list[PipelineStage],
+            num_microbatches: int,
+            loss_fn: LossFn | None,
+            program: dict[int, list[ActionBase]]
+    ):
+        """
+        Constructs the schedule executor.
+
+        Args:
+            dist_context: The distributed context.
+            stages: List of stages managed by this executor.
+            num_microbatches: Number of microbatches the global batch is split.
+            loss_fn: Function to compute loss.
+            program: The execution plan mapping rank ID to a list of actions.
+        """
+
+        self._dist_ctx = dist_context
+        self._stages = {stage.info.current_stage: stage for stage in stages}
+        self._num_microbatches = num_microbatches
+        self._program = program
+
+        self._has_backward = any(any(
+            action.has_backward_work for action in sub_program
+        ) for sub_program in program.values())
+
+        self._comm_handler = PipelineCommunicationHandler(self._stages)
+        if loss_fn is None:
+            self._loss_handler = None
+        else:
+            self._loss_handler = PipelineLossHandler(loss_fn)
+
+        self._input_data_sharding_spec: ShardingSpec | None = None
+        self._input_kwargs_sharding_spec: ShardingSpec | None = None
+
+    def configure_buffers(
+            self,
+            inputs: dict[str, torch.Tensor],
+            kwargs: dict[str, Any],
+            sharding_spec: PipelineShardingSpec | None
+    ):
+        if sharding_spec is None or sharding_spec.input_data is None:
+            self._input_data_sharding_spec = shard_spec_on_dim(inputs, dim=0)
+        if sharding_spec is None or sharding_spec.input_kwargs is None:
+            self._input_kwargs_sharding_spec = shard_spec_on_dim(kwargs, dim=0)
+
+        for stage in self._stages.values():
+            stage.configure_buffers(
+                num_microbatches=self._num_microbatches,
+                pipeline_inputs=inputs,
+                has_backward=self._has_backward
+            )
+
+    def step(self, inputs: dict[str, torch.Tensor], kwargs: dict[str, Any]):
+        if self._input_data_sharding_spec is None or self._input_kwargs_sharding_spec is None:
+            raise ValueError("Please configure sharding specs first")
+
+        self._dist_ctx.logger.debug("Begin pipeline step")
+        pp_group = self._dist_ctx.mesh_for(REGULAR_DOMAIN).get_group("pp")
+
+        for stage in self._stages.values():
+            stage.reset()
+
+        # Shard inputs and kwargs to microbatches
+        inputs_shard = shard_tree(
+            inputs,
+            num_shards=self._num_microbatches,
+            sharding_spec=self._input_data_sharding_spec,
+            enforce_even_split=True
+        )
+        kwargs_shard = shard_tree(
+            kwargs,
+            num_shards=self._num_microbatches,
+            sharding_spec=self._input_kwargs_sharding_spec,
+            enforce_even_split=True
+        )
+
+        my_program = self._program[pp_group.rank()]
+
+        for action in my_program:
+            with record_function(str(action)):
+                self._dist_ctx.logger.debug(f"Running pipeline action {action}")
+                action.apply(ActionContext(
+                    loss=self._loss_handler,
+                    stages=self._stages,
+                    communications=self._comm_handler,
+                    pipeline_inputs_microbatches=inputs_shard,
+                    pipeline_kwargs_microbatches=kwargs_shard
+                ))
+
+        self._dist_ctx.logger.debug("Waiting for potentially hanging PP send comms")
+        self._comm_handler.wait_send_all()  # finalize just in case
+        self._dist_ctx.logger.debug("End pipeline step")

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

@@ -0,0 +1,55 @@
+from collections.abc import Callable
+
+import torch
+
+LossFn = Callable[[dict[str, torch.Tensor], int], torch.Tensor]
+
+
+class PipelineLossHandler:
+    """Manages loss computation and state caching across forward and backward passes."""
+
+    def __init__(self, loss_fn: LossFn):
+        """
+        Constructs the loss handler.
+
+        Args:
+            loss_fn: The callable that computes loss from model outputs.
+        """
+
+        self._loss_fn = loss_fn
+        self._cached_values: dict[int, torch.Tensor] = {}
+
+    def compute_loss(self, forward_result: dict[str, torch.Tensor], microbatch_index: int) -> torch.Tensor:
+        """
+        Computes loss for a given microbatch result and caches it.
+
+        Args:
+            forward_result: The output from the last stage of the model.
+            microbatch_index: The index of the microbatch being processed.
+
+        Returns:
+            The computed loss tensor.
+        """
+
+        result = self._loss_fn(forward_result, microbatch_index)
+        self._cached_values[microbatch_index] = result
+        return result
+
+    def acquire_loss(self, microbatch_index: int) -> torch.Tensor:
+        """
+        Retrieves the cached loss tensor for the backward pass and removes it from the cache.
+
+        Args:
+            microbatch_index: The index of the microbatch.
+
+        Returns:
+            The previously computed loss tensor.
+
+        Raises:
+            ValueError: If the loss for this microbatch hasn't been computed yet.
+        """
+
+        if microbatch_index not in self._cached_values:
+            raise ValueError()
+
+        return self._cached_values[microbatch_index]

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

@@ -0,0 +1,15 @@
+"""
+Pipeline Schedule Implementations
+"""
+
+from .bfs import LoopedBFSPipelineProgramBuilder
+from .dualpipev import DualPipeVPipelineProgramBuilder
+from .interleaved import Interleaved1F1BPipelineProgramBuilder
+from .zerobubblev import ZeroBubbleVPipelineProgramBuilder
+
+__all__ = [
+    "DualPipeVPipelineProgramBuilder",
+    "Interleaved1F1BPipelineProgramBuilder",
+    "LoopedBFSPipelineProgramBuilder",
+    "ZeroBubbleVPipelineProgramBuilder"
+]