d9d 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,227 @@
+from ..component.program import (
+    PipelineProgramBuilder,
+    ScheduleStyle,
+    add_communication_ops,
+    build_stage_to_host_rank_topology,
+)
+from ..component.runtime import (
+    ActionBase,
+    BackwardFullInputComputeAction,
+    BackwardWeightComputeAction,
+    ForwardComputeAction,
+)
+
+
+class ZeroBubbleVPipelineProgramBuilder(PipelineProgramBuilder):
+    """
+    Builder for the Zero Bubble V (ZBV) Pipeline Schedule.
+
+    This schedule is designed for V-shape topologies (2 stages per rank) and
+    utilizes the Zero Bubble optimizations by splitting backward passes.
+
+    It requires exactly two stages
+    per rank organized in a V-shape topology and splits backward passes into
+    Input and Weight gradients to optimize pipeline throughput.
+
+    References:
+        https://arxiv.org/pdf/2401.10241, Section 6
+    """
+
+    def __init__(self):
+        """Constructs the ZBV builder."""
+
+    def compose(
+            self, num_microbatches: int, pp_size: int
+    ) -> dict[int, list[ActionBase]]:
+        num_stages = self.num_stages_per_rank * pp_size
+
+        # 1. Topology
+        # V-style: Rank 0 gets Stage 0 & Stage N-1. Rank 1 gets Stage 1 & Stage N-2...
+        stage_to_rank = build_stage_to_host_rank_topology(
+            pp_size=pp_size, num_stages=num_stages, style=ScheduleStyle.v
+        )
+
+        actions: dict[int, list[ActionBase]] = {}
+
+        for rank in range(pp_size):
+            actions[rank] = self._generate_rank_schedule(
+                rank=rank,
+                pp_size=pp_size,
+                num_stages=num_stages,
+                target_microbatches=num_microbatches,
+            )
+
+        # 2. Inject Communications
+        return add_communication_ops(
+            compute_actions=actions,
+            stage_to_rank=stage_to_rank,
+            num_stages=num_stages
+        )
+
+    def _generate_rank_schedule(  # noqa: C901
+            self,
+            rank: int,
+            pp_size: int,
+            num_stages: int,
+            target_microbatches: int,
+    ) -> list[ActionBase]:
+        # ZBV logic assumes the pipeline is fully saturated to define the loop bounds.
+        # We simulate enough steps to cover the topology startup, then filter
+        # down to the user's requested microbatches at the end.
+        simulated_n_micro = max(2 * pp_size - 1, target_microbatches)
+
+        rank_ops: list[ActionBase] = []
+
+        # -- Stage Identification (V-Shape) --
+        # s0: The "Forward-going" chunk (e.g., Stage 0 for Rank 0)
+        # s1: The "Backward-coming" chunk (e.g., Stage N-1 for Rank 0)
+        s0 = rank
+        s1 = num_stages - 1 - rank
+
+        # -- Counters --
+        # Track next microbatch index for each operation type on each chunk.
+        # F: Forward, I: Backward Input, W: Backward Weight
+        f0_cnt = 0
+        b0_cnt = 0  # Input Grad Counter (Chunk 0)
+        w0_cnt = 0  # Weight Grad Counter (Chunk 0)
+
+        f1_cnt = 0
+        b1_cnt = 0  # Input Grad Counter (Chunk 1)
+        w1_cnt = 0  # Weight Grad Counter (Chunk 1)
+
+        # -- Helpers --
+
+        def emit_f(stage: int, idx: int):
+            rank_ops.append(ForwardComputeAction(stage_idx=stage, microbatch_idx=idx))
+
+        def emit_i_and_w(stage: int, idx: int):
+            rank_ops.append(
+                BackwardFullInputComputeAction(
+                    stage_idx=stage, microbatch_idx=idx, full_backward=False
+                )
+            )
+            rank_ops.append(
+                BackwardWeightComputeAction(stage_idx=stage, microbatch_idx=idx)
+            )
+
+        def emit_i(stage: int, idx: int):
+            rank_ops.append(
+                BackwardFullInputComputeAction(
+                    stage_idx=stage, microbatch_idx=idx, full_backward=False
+                )
+            )
+
+        def emit_w(stage: int, idx: int):
+            rank_ops.append(
+                BackwardWeightComputeAction(stage_idx=stage, microbatch_idx=idx)
+            )
+
+        # -- Phase 1: Warmup 1 (Chunk 0 Forwards) --
+        warmup_n1 = 2 * (pp_size - rank) - 1
+        for _ in range(warmup_n1):
+            emit_f(s0, f0_cnt)
+            f0_cnt += 1
+
+        # -- Phase 2: Warmup 2 (Interleave F1, F0) --
+        warmup_n2 = rank
+        for _ in range(warmup_n2):
+            emit_f(s1, f1_cnt)
+            f1_cnt += 1
+            emit_f(s0, f0_cnt)
+            f0_cnt += 1
+
+        # -- Phase 3: Warmup 3 (F1, then B1 I+W) --
+        warmup_n3 = pp_size - rank
+        for _ in range(warmup_n3):
+            emit_f(s1, f1_cnt)
+            f1_cnt += 1
+
+            emit_i_and_w(s1, b1_cnt)
+            b1_cnt += 1
+            w1_cnt += 1
+
+        # -- Phase 4: Stable State --
+        while f1_cnt < f0_cnt or f0_cnt < simulated_n_micro:
+            # Emit F0 if within bounds
+            if f0_cnt < simulated_n_micro:
+                emit_f(s0, f0_cnt)
+                f0_cnt += 1
+
+            # Emit B0 (I+W)
+            emit_i_and_w(s0, b0_cnt)
+            b0_cnt += 1
+            w0_cnt += 1
+
+            # Emit F1
+            emit_f(s1, f1_cnt)
+            f1_cnt += 1
+
+            # Emit B1 (I+W)
+            emit_i_and_w(s1, b1_cnt)
+            b1_cnt += 1
+            w1_cnt += 1
+
+        # -- Phase 5: Cooldown 1 (Splitting I and W) --
+        # In cooldown, the I and W streams diverge to fill bubbles.
+        cooldown_n1 = rank
+        for _ in range(cooldown_n1):
+            emit_i(s0, b0_cnt)
+            b0_cnt += 1
+
+            emit_i(s1, b1_cnt)
+            b1_cnt += 1
+
+        # -- Phase 6: Cooldown 2 (I0, then W0) --
+        cooldown_n2 = pp_size - rank
+        for _ in range(cooldown_n2):
+            # Input Grad Chunk 0
+            emit_i(s0, b0_cnt)
+            b0_cnt += 1
+
+            # Weight Grad Chunk 0 (delayed from previous steps)
+            emit_w(s0, w0_cnt)
+            w0_cnt += 1
+
+        # -- Phase 7: Flush Remaining Weights --
+
+        # Flush W1
+        while w1_cnt < b1_cnt:
+            emit_w(s1, w1_cnt)
+            w1_cnt += 1
+
+        # Flush W0
+        while w0_cnt < b0_cnt:
+            emit_w(s0, w0_cnt)
+            w0_cnt += 1
+
+        # -- Integrity Check --
+        if not (w0_cnt == b0_cnt == f0_cnt):
+            raise RuntimeError(
+                f"ZBV Schedule Failed (Chunk 0): F={f0_cnt}, I={b0_cnt}, W={w0_cnt}"
+            )
+        if not (w1_cnt == b1_cnt == f1_cnt):
+            raise RuntimeError(
+                f"ZBV Schedule Failed (Chunk 1): F={f1_cnt}, I={b1_cnt}, W={w1_cnt}"
+            )
+
+        # -- Post-Process: Filter to Target Microbatches --
+        # Remove any actions involving simulated microbatches beyond the user's request.
+        final_ops: list[ActionBase] = []
+        for action in rank_ops:
+            if isinstance(action, (ForwardComputeAction,
+                                   BackwardFullInputComputeAction,
+                                   BackwardWeightComputeAction)):
+                if action.microbatch_idx < target_microbatches:
+                    final_ops.append(action)
+            else:
+                final_ops.append(action)
+
+        return final_ops
+
+    @property
+    def num_stages_per_rank(self) -> int:
+        return 2
+
+    @property
+    def topology_style(self) -> ScheduleStyle:
+        return ScheduleStyle.v

d9d/pipelining/infra/stage/__init__.py

@@ -0,0 +1,5 @@
+from .stage import PipelineStage
+
+__all__ = [
+    "PipelineStage"
+]

d9d/pipelining/infra/stage/communications.py

@@ -0,0 +1,274 @@
+import dataclasses
+
+import torch
+import torch.distributed as dist
+
+
+@dataclasses.dataclass(kw_only=True, slots=True)
+class ReceiveStageInput:
+    """
+    Instruction to receive a specific tensor from a previous stage (or next stage during backward).
+
+    Attributes:
+        name: A unique identifier for the communication operation.
+        from_stage: The stage index sending the data.
+        buffer: The pre-allocated tensor buffer where data will be received.
+    """
+
+    name: str
+    from_stage: int
+    buffer: torch.Tensor
+
+
+@dataclasses.dataclass
+class StartStageInput:
+    """
+    Instruction indicating that the input for this stage does not come from communication
+    (e.g., this is the first stage receiving data loader inputs).
+    """
+
+
+StageInput = ReceiveStageInput | StartStageInput
+
+
+@dataclasses.dataclass(kw_only=True, slots=True)
+class SendStageOutput:
+    """
+    Instruction to send a specific tensor to a next stage (or previous if backward).
+
+    Attributes:
+        to_stage: The stage index receiving the data.
+    """
+
+    to_stage: int
+
+
+@dataclasses.dataclass
+class EndStageOutput:
+    """
+    Instruction indicating that the output of this stage is not sent anywhere
+    (e.g., this is the last stage computing loss).
+    """
+
+
+StageOutput = SendStageOutput | EndStageOutput
+
+
+class StageCommunicationHandler:
+    """
+    Manages Point-to-Point (P2P) communication descriptors for a specific data flow direction within a pipeline stage.
+
+    This class handles the creation of P2P operations (send/recv) across multiple microbatches,
+    managing buffers and mapping logical stage indices to physical ranks.
+    """
+
+    def __init__(
+            self,
+
+            name: str,
+            stage_index: int,
+            num_microbatches: int,
+
+            input_stage_index: int | None,
+            input_args: dict[str, torch.Tensor],
+
+            output_stage_index: int | None,
+            output_args: dict[str, torch.Tensor],
+
+            stage_idx_to_host_rank: dict[int, int],
+            group: dist.ProcessGroup
+    ):
+        """
+        Constructs a StageCommunicationHandler object.
+
+        Args:
+            name: Name prefix for this handler (e.g., 'fwd', 'bwd').
+            stage_index: The logical index of the current stage.
+            num_microbatches: Total number of microbatches ("chunks") to schedule.
+            input_stage_index: The logical index of the stage providing inputs, or None if inputs are local.
+            input_args: Metadata (shapes/dtypes) for input tensors.
+            output_stage_index: The logical index of the stage consuming outputs, or None if outputs are terminal.
+            output_args: Metadata (shapes/dtypes) for output tensors.
+            stage_idx_to_host_rank: Mapping from logical stage indices to physical world ranks.
+            group: The process group strictly for pipeline communication.
+        """
+
+        self._input_handlers = self._build_inputs(
+            name=name,
+            stage_index=stage_index,
+            num_microbatches=num_microbatches,
+            input_stage_index=input_stage_index,
+            input_args=input_args
+        )
+        self._output_handlers = self._build_outputs(
+            output_stage_index=output_stage_index,
+            output_args=output_args
+        )
+
+        self._stage_idx_to_host_rank = stage_idx_to_host_rank
+        self._group = group
+
+    @staticmethod
+    def _build_inputs(
+            name: str,
+            stage_index: int,
+            num_microbatches: int,
+            input_stage_index: int | None,
+            input_args: dict[str, torch.Tensor]
+    ) -> dict[int, dict[str, StageInput]]:
+        handlers: dict[int, dict[str, StageInput]] = {}
+
+        for chunk_id in range(num_microbatches):
+            handlers[chunk_id] = {}
+            for input_name, input_tensor_meta in input_args.items():
+                if input_stage_index is None:
+                    handlers[chunk_id][input_name] = StartStageInput()
+                else:
+                    handlers[chunk_id][input_name] = ReceiveStageInput(
+                        name=f"{name}_recv_from_{input_stage_index}_to_{stage_index}[{chunk_id}][{input_name}]",
+                        from_stage=input_stage_index,
+                        buffer=torch.empty(
+                            input_tensor_meta.size(),
+                            dtype=input_tensor_meta.dtype,
+                            layout=input_tensor_meta.layout,
+                            device="cuda"  # force device
+                        )
+                    )
+        return handlers
+
+    @staticmethod
+    def _build_outputs(
+            output_stage_index: int | None,
+            output_args: dict[str, torch.Tensor]
+    ) -> dict[str, StageOutput]:
+        handlers: dict[str, StageOutput] = {}
+
+        for output_name in output_args:
+            if output_stage_index is None:
+                handlers[output_name] = EndStageOutput()
+            else:
+                handlers[output_name] = SendStageOutput(
+                    to_stage=output_stage_index
+                )
+        return handlers
+
+    def set_input_requires_grad_(self, requires_grad: bool):
+        """
+        Sets the `requires_grad` flag for all internal input buffers.
+
+        Typically used to enable gradient flow from backward stages to forward stages.
+
+        Args:
+            requires_grad: Whether the buffers should require gradients.
+        """
+
+        for inputs in self._input_handlers.values():
+            for info in inputs.values():
+                if isinstance(info, ReceiveStageInput):
+                    info.buffer.requires_grad_(requires_grad)
+
+    def set_inputs_local(self, inputs: dict[str, torch.Tensor], microbatch_index: int):
+        """
+        Manually fills the input buffer for a specific microbatch with local data.
+
+        This is used when the stage is the first in the pipeline or receives data
+        from a dataloader rather than via network communication.
+
+        Args:
+            inputs: Dictionary of input tensors.
+            microbatch_index: The microbatch identifier.
+        """
+
+        for input_name, input_value in inputs.items():
+            handler = self._input_handlers[microbatch_index][input_name]
+            if not isinstance(handler, ReceiveStageInput):
+                raise RuntimeError("Tried to set a buffer of no-receive stage input")
+            prev_requires_grad = handler.buffer.requires_grad
+            handler.buffer = input_value.detach().requires_grad_(
+                prev_requires_grad)
+
+    def get_inputs(self, microbatch_index: int) -> dict[str, torch.Tensor]:
+        """
+        Retrieves the input tensors for a specific microbatch from the internal buffers.
+
+        Args:
+            microbatch_index: The microbatch identifier.
+
+        Returns:
+            Dictionary mapping input names to tensors.
+        """
+        outputs: dict[str, torch.Tensor] = {}
+
+        for input_name, input_info in self._input_handlers[microbatch_index].items():
+            if not isinstance(input_info, ReceiveStageInput):
+                raise RuntimeError("Tried to get a buffer of no receive stage input")
+            outputs[input_name] = input_info.buffer
+
+        return outputs
+
+    def create_receive_ops(self, microbatch_index: int) -> list[dist.P2POp]:
+        """
+        Generates the PyTorch P2P receive operations for a specific microbatch.
+
+        Args:
+            microbatch_index: The microbatch identifier.
+
+        Returns:
+            A list of `dist.P2POp` objects configured for `dist.irecv`.
+        """
+
+        ops = []
+
+        inputs = self._input_handlers[microbatch_index]
+        # sort ops by parameter names to ensure receive ops are ordered the same for send and recv
+        for _input_name, input_info in sorted(inputs.items(), key=lambda x: x[0]):
+            match input_info:
+                case StartStageInput():
+                    pass
+                case ReceiveStageInput():
+                    peer_rank = self._stage_idx_to_host_rank[input_info.from_stage]
+                    peer_global_rank = dist.get_global_rank(self._group, peer_rank)
+                    op = dist.P2POp(dist.irecv, input_info.buffer, peer_global_rank, self._group)
+                    ops.append(op)
+                case _:
+                    raise ValueError()
+
+        return ops
+
+    def create_send_ops(self, send_contents: dict[str, torch.Tensor]) -> list[dist.P2POp]:
+        """
+        Generates the PyTorch P2P send operations for the provided tensors.
+
+        Args:
+            send_contents: Dictionary of tensors to send.
+
+        Returns:
+            A list of `dist.P2POp` objects configured for `dist.isend`.
+        """
+
+        ops = []
+
+        # sort ops by parameter names to ensure receive ops are ordered the same for send and recv
+        for output_name, output_info in sorted(self._output_handlers.items(), key=lambda x: x[0]):
+            output_tensor = send_contents[output_name]
+
+            match output_info:
+                case EndStageOutput():
+                    pass
+                case SendStageOutput():
+                    peer_rank = self._stage_idx_to_host_rank[output_info.to_stage]
+                    peer_global_rank = dist.get_global_rank(self._group, peer_rank)
+                    op = dist.P2POp(dist.isend, output_tensor, peer_global_rank, self._group)
+                    ops.append(op)
+                case _:
+                    raise ValueError()
+
+        return ops
+
+    def reset(self):
+        """Resets the internal state, specifically clearing gradients on input buffers."""
+
+        for inp_handlers in self._input_handlers.values():
+            for inp_handler in inp_handlers.values():
+                if isinstance(inp_handler, ReceiveStageInput):
+                    inp_handler.buffer.grad = None