d9d 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,86 @@
+from ..component.program import (
+    PipelineProgramBuilder,
+    ScheduleStyle,
+    add_communication_ops,
+    build_stage_to_host_rank_topology,
+)
+from ..component.runtime import (
+    ActionBase,
+    BackwardFullInputComputeAction,
+    ForwardComputeAction,
+)
+
+
+class LoopedBFSPipelineProgramBuilder(PipelineProgramBuilder):
+    """
+    Builder for the Breadth-First Pipeline Parallelism schedule.
+
+    This schedule runs all available forward microbatches for local stages first.
+    If configured for training, it then runs backwards in reverse topological order.
+
+    References:
+        https://arxiv.org/pdf/2211.05953
+    """
+
+    def __init__(self, num_stages_per_rank: int, inference_mode: bool = False):
+        """
+        Constructs the LoopedBFS builder.
+
+        Args:
+            num_stages_per_rank: Number of stages per rank.
+            inference_mode: If True, only forward passes are scheduled. If False,
+                both forward and backward passes are scheduled.
+        """
+        self._num_stages_per_rank = num_stages_per_rank
+        self._inference_mode = inference_mode
+
+    def compose(self, num_microbatches: int, pp_size: int) -> dict[int, list[ActionBase]]:
+        num_stages = self._num_stages_per_rank * pp_size
+        stage_to_rank = build_stage_to_host_rank_topology(
+            pp_size=pp_size,
+            num_stages=num_stages,
+            style=ScheduleStyle.loop
+        )
+
+        compute_actions: dict[int, list[ActionBase]] = {r: [] for r in range(pp_size)}
+
+        for rank in range(pp_size):
+            my_stages = [s for s in range(num_stages) if stage_to_rank[s] == rank]
+
+            # Schedule all Forwards
+            # In Breadth-First loops, we finish all microbatches for the current stage
+            # before moving to the next stage assigned to this rank.
+            for stage_idx in my_stages:
+                for mb_idx in range(num_microbatches):
+                    compute_actions[rank].append(
+                        ForwardComputeAction(
+                            stage_idx=stage_idx,
+                            microbatch_idx=mb_idx
+                        )
+                    )
+
+            # Schedule all Backwards (Reverse order) - Only if training
+            if not self._inference_mode:
+                for stage_idx in reversed(my_stages):
+                    for mb_idx in reversed(range(num_microbatches)):
+                        compute_actions[rank].append(
+                            BackwardFullInputComputeAction(
+                                stage_idx=stage_idx,
+                                microbatch_idx=mb_idx,
+                                full_backward=True
+                            )
+                        )
+
+        return add_communication_ops(
+            compute_actions=compute_actions,
+            stage_to_rank=stage_to_rank,
+            num_stages=num_stages
+        )
+
+    @property
+    def num_stages_per_rank(self) -> int:
+        return self._num_stages_per_rank
+
+    @property
+    def topology_style(self) -> ScheduleStyle:
+        return ScheduleStyle.loop

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

@@ -0,0 +1,234 @@
+from collections import deque
+
+from ..component.program import (
+    PipelineProgramBuilder,
+    ScheduleStyle,
+    add_communication_ops,
+    build_stage_to_host_rank_topology,
+)
+from ..component.runtime import (
+    ActionBase,
+    BackwardFullInputComputeAction,
+    BackwardWeightComputeAction,
+    ComposeAction,
+    ForwardComputeAction,
+)
+
+
+class DualPipeVPipelineProgramBuilder(PipelineProgramBuilder):
+    """
+    Builder for the DualPipeV Pipeline Parallelism schedule.
+
+    DualPipeV is a specialized bi-directional pipeline schedule designed for high
+    throughput training. It requires exactly 2 stages per pipeline rank (V-shape)
+    and utilizes split backward passes (Input gradients vs Weight gradients)
+    to fill pipeline bubbles.
+
+    References:
+        https://github.com/deepseek-ai/DualPipe
+        https://hackmd.io/@ufotalent/r1lVXsa9Jg
+    """
+
+    def __init__(self):
+        """
+        Constructs the DualPipeV builder.
+        """
+
+    @staticmethod
+    def _build_for_rank(  # noqa: C901
+            rank: int, stage_to_rank: dict[int, int], num_microbatches: int, pp_size: int
+    ) -> list[ActionBase]:
+        compute_actions: list[ActionBase] = []
+
+        # Identify local stages: s0 is Phase 0, s1 is Phase 1
+        my_stages = sorted([s for s, r in stage_to_rank.items() if r == rank])
+        s0, s1 = my_stages[0], my_stages[1]
+
+        # Track microbatch indices for each stage and operation type
+        # f_idx: Next Forward microbatch
+        # b_idx: Next Backward microbatch (Input or Full)
+        f_idx = {s0: 0, s1: 0}
+        b_idx = {s0: 0, s1: 0}
+
+        # Queue for Zero Bubble optimization: stores (stage, mb_idx) for deferred weight grads
+        weight_queue: deque[tuple[int, int]] = deque()
+
+        # --- Helper Functions for Action Emission ---
+
+        def _add_f(stage: int):
+            compute_actions.append(
+                ForwardComputeAction(stage_idx=stage, microbatch_idx=f_idx[stage])
+            )
+            f_idx[stage] += 1
+
+        def _add_b_full(stage: int):
+            compute_actions.append(
+                BackwardFullInputComputeAction(
+                    stage_idx=stage,
+                    microbatch_idx=b_idx[stage],
+                    full_backward=True,
+                )
+            )
+            b_idx[stage] += 1
+
+        def _add_b_input(stage: int):
+            mb = b_idx[stage]
+            compute_actions.append(
+                BackwardFullInputComputeAction(
+                    stage_idx=stage,
+                    microbatch_idx=mb,
+                    full_backward=False,
+                )
+            )
+            weight_queue.append((stage, mb))
+            b_idx[stage] += 1
+
+        def _pop_w():
+            if not weight_queue:
+                return
+            s, mb = weight_queue.popleft()
+            compute_actions.append(
+                BackwardWeightComputeAction(stage_idx=s, microbatch_idx=mb)
+            )
+
+        def _add_overlap_f_b(stage_f: int, stage_b: int, b_is_full: bool):
+            """Emit overlapped Forward and Backward actions."""
+            mb_f = f_idx[stage_f]
+            mb_b = b_idx[stage_b]
+
+            act_f = ForwardComputeAction(stage_idx=stage_f, microbatch_idx=mb_f)
+
+            act_b = BackwardFullInputComputeAction(
+                stage_idx=stage_b, microbatch_idx=mb_b, full_backward=b_is_full
+            )
+            if not b_is_full:
+                weight_queue.append((stage_b, mb_b))
+
+            f_idx[stage_f] += 1
+            b_idx[stage_b] += 1
+
+            # Note: d9d infra treats ComposeAction sequentially in simulation,
+            # but runtime may overlap them.
+            compute_actions.append(ComposeAction(actions=(act_f, act_b)))
+
+        # Step 1: nF0 (Startup Phase 0)
+        step_1 = (pp_size - rank - 1) * 2
+        for _ in range(step_1):
+            _add_f(s0)
+
+        # Step 2: nF0F1 (Forward fill)
+        step_2 = rank + 1
+        for _ in range(step_2):
+            _add_f(s0)
+            _add_f(s1)
+
+        # Step 3: nI1W1F1 (Mixed Phase with Zero Bubble)
+        step_3 = pp_size - rank - 1
+        for _ in range(step_3):
+            _add_b_input(s1)  # Backward Input Phase 1
+            _pop_w()  # Weight Phase (accumulated from prev)
+            _add_f(s1)  # Forward Phase 1
+
+        # Step 4: The Main Loop (Interleaved Forward/Backward)
+        step_4 = num_microbatches - 2 * pp_size + rank + 1
+        for i in range(step_4):
+            # Sub-step A: F0 & B1
+            if i == 0 and rank == pp_size - 1:
+                # Specific case for last rank on first iter: do not overlap
+                _add_f(s0)
+                _add_b_full(s1)
+            else:
+                # Overlap F0 and B1 (usually full backward unless we were in ZB mode,
+                # but DualPipeV main loop defaults to full for simplicity unless tuned)
+                # DeepSeek impl uses standard backward here (zb=False).
+                _add_overlap_f_b(stage_f=s0, stage_b=s1, b_is_full=True)
+
+            # Sub-step B: F1 & B0
+            # Overlap F1 and B0 (Full)
+            _add_overlap_f_b(stage_f=s1, stage_b=s0, b_is_full=True)
+
+        # Step 5: Cooldown F1/B0
+        step_5 = pp_size - rank - 1
+        for _ in range(step_5):
+            _add_b_full(s1)
+            _add_overlap_f_b(stage_f=s1, stage_b=s0, b_is_full=True)
+
+        # Step 6: Cooldown B1/B0 with Zero Bubble ramp-up
+        step_6 = rank + 1
+        enable_zb = False
+        for i in range(step_6):
+            # Phase 1 Backward
+            if i == step_6 // 2 and rank % 2 == 1:
+                enable_zb = True
+
+            if enable_zb:
+                _add_b_input(s1)
+            else:
+                _add_b_full(s1)
+
+            # Phase 0 Backward
+            if i == step_6 // 2 and rank % 2 == 0:
+                enable_zb = True
+
+            if enable_zb:
+                _add_b_input(s0)
+            else:
+                _add_b_full(s0)
+
+        # Step 7: Zero Bubble Weights + B0
+        step_7 = pp_size - rank - 1
+        for _ in range(step_7):
+            _pop_w()
+            # DeepSeek source explicitly uses enable_zb=True here for chunk 0
+            _add_b_input(s0)
+
+        # Step 8: Flush Weights
+        step_8 = rank + 1
+        for _ in range(step_8):
+            _pop_w()
+
+        return compute_actions
+
+    def compose(
+            self, num_microbatches: int, pp_size: int
+    ) -> dict[int, list[ActionBase]]:
+        num_stages = self.num_stages_per_rank * pp_size
+
+        if num_microbatches < num_stages:
+            raise ValueError(
+                f"DualPipeV requires num_microbatches ({num_microbatches}) >= "
+                f"num_stages ({num_stages})."
+            )
+
+        # Ranks hold stages in a V pattern (e.g., Rank 0 holds Stage 0 and Stage N-1).
+        # We rely on the sorted order of local steps to determine Phase 0 (Forward-going)
+        # and Phase 1 (Backward-coming).
+        stage_to_rank = build_stage_to_host_rank_topology(
+            pp_size=pp_size, num_stages=num_stages, style=ScheduleStyle.v
+        )
+
+        compute_actions: dict[int, list[ActionBase]] = {r: [] for r in range(pp_size)}
+
+        for rank in range(pp_size):
+            compute_actions[rank] = self._build_for_rank(
+                rank=rank,
+                pp_size=pp_size,
+                num_microbatches=num_microbatches,
+                stage_to_rank=stage_to_rank
+            )
+
+        # 4. Inject Communication Operations
+        # This wrapper handles dependency analysis and inserts Send/Recv/Wait ops.
+        return add_communication_ops(
+            compute_actions=compute_actions,
+            stage_to_rank=stage_to_rank,
+            num_stages=num_stages
+        )
+
+    @property
+    def num_stages_per_rank(self) -> int:
+        return 2
+
+    @property
+    def topology_style(self) -> ScheduleStyle:
+        return ScheduleStyle.v

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

@@ -0,0 +1,240 @@
+from collections import defaultdict, deque
+
+from ..component.program import (
+    PipelineProgramBuilder,
+    ScheduleStyle,
+    add_communication_ops,
+    build_stage_to_host_rank_topology,
+)
+from ..component.runtime import (
+    ActionBase,
+    BackwardFullInputComputeAction,
+    BackwardWeightComputeAction,
+    ForwardComputeAction,
+)
+
+
+class Interleaved1F1BPipelineProgramBuilder(PipelineProgramBuilder):
+    """
+    Builder for Interleaved Pipeline Parallelism schedules.
+
+    This builder supports:
+
+    1.  **Standard Interleaved 1F1B**: Assigns multiple stages per rank and prioritizes
+        depth-first execution. (See https://arxiv.org/pdf/2104.04473)
+    2.  **Interleaved Zero Bubble (ZB1P)**: Extends 1F1B by splitting backward passes
+        into Input Gradients and Weight Gradients. Weight gradients are delayed
+        to fill pipeline bubbles. (See https://arxiv.org/pdf/2401.10241)
+    """
+
+    def __init__(self, num_stages_per_rank: int, enable_zero_bubble: bool = False):
+        """
+        Constructs the Interleaved 1F1B builder.
+
+        Args:
+            num_stages_per_rank: Number of stages per rank.
+            enable_zero_bubble: If True, uses the ZB1P schedule variant which
+                splits backward passes to reduce bubble size.
+        """
+        self._num_stages_per_rank = num_stages_per_rank
+        self._enable_zero_bubble = enable_zero_bubble
+
+    def _get_warmup_ops(
+            self,
+            rank: int,
+            microbatches_per_round: int,
+            pp_size: int,
+            n_microbatches: int,
+            multiply_factor: int,
+    ) -> int:
+        """
+        Calculates the number of warmup steps required before entering steady state.
+        """
+        warmups_ops_last_stage = (self._num_stages_per_rank - 1) * microbatches_per_round
+        warmup_ops = warmups_ops_last_stage + multiply_factor * ((pp_size - 1) - rank)
+        return min(warmup_ops, n_microbatches * self._num_stages_per_rank)
+
+    def compose(
+            self, num_microbatches: int, pp_size: int
+    ) -> dict[int, list[ActionBase]]:
+        """
+        Generates the execution program for all ranks.
+
+        Args:
+            num_microbatches: Total microbatches. Must be divisible by the derived
+                number of rounds.
+            pp_size: Number of pipeline ranks.
+
+        Returns:
+            A dictionary mapping rank indices to their list of sequential actions.
+        """
+        num_stages = self.num_stages_per_rank * pp_size
+
+        if num_stages % pp_size != 0:
+            raise ValueError(
+                f"num_stages ({num_stages}) must be divisible by pp_size ({pp_size}) "
+                "for interleaved schedules."
+            )
+
+        # 1. Topology Setup
+        # Use Loop/Round-Robin assignment: Rank 0 gets Stage 0, PP, 2*PP...
+        stage_to_rank = build_stage_to_host_rank_topology(
+            pp_size=pp_size, num_stages=num_stages, style=ScheduleStyle.loop
+        )
+
+        num_rounds = max(1, num_microbatches // pp_size)
+
+        if num_microbatches % num_rounds != 0:
+            raise ValueError(
+                f"microbatches ({num_microbatches}) must be divisible by rounds ({num_rounds})."
+            )
+
+        microbatches_per_round = num_microbatches // num_rounds
+
+        # 2. Schedule Generation
+        actions: dict[int, list[ActionBase]] = {}
+
+        # Zero Bubble 1f1b uses a shorter warmup heuristic (factor 1) than Standard (factor 2)
+        warmup_multiplier = 1 if self._enable_zero_bubble else 2
+
+        for rank in range(pp_size):
+            actions[rank] = self._generate_rank_schedule(
+                rank=rank,
+                pp_size=pp_size,
+                n_microbatches=num_microbatches,
+                microbatches_per_round=microbatches_per_round,
+                multiply_factor=warmup_multiplier,
+            )
+
+        # 3. Communication Injection
+        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,
+            n_microbatches: int,
+            microbatches_per_round: int,
+            multiply_factor: int,
+    ) -> list[ActionBase]:
+        """
+        Generates the sequential list of compute actions for a specific rank.
+        """
+        rank_actions: list[ActionBase] = []
+
+        # -- State Tracking --
+        # Map: stage_idx -> next_microbatch_idx
+        fwd_counters: dict[int, int] = defaultdict(int)
+        bwd_counters: dict[int, int] = defaultdict(int)
+
+        # FIFO Queue for deferred weight gradients in Zero Bubble
+        # Stores: (stage_idx, microbatch_idx)
+        pending_weights: deque[tuple[int, int]] = deque()
+
+        # -- Helpers --
+
+        def get_global_stage(local_idx: int) -> int:
+            """Converts a local virtual stage index (0..N) to global stage ID."""
+            return (local_idx * pp_size) + rank
+
+        def get_fwd_local_idx(op_idx: int) -> int:
+            return (op_idx // microbatches_per_round) % self._num_stages_per_rank
+
+        def get_bwd_local_idx(op_idx: int, warmup_offset: int) -> int:
+            return (self._num_stages_per_rank
+                    - 1
+                    - ((op_idx - warmup_offset) // microbatches_per_round) % self._num_stages_per_rank)
+
+        def emit_forward(op_idx: int):
+            local_idx = get_fwd_local_idx(op_idx)
+            stage = get_global_stage(local_idx)
+            mb = fwd_counters[stage]
+
+            rank_actions.append(ForwardComputeAction(stage_idx=stage, microbatch_idx=mb))
+            fwd_counters[stage] += 1
+
+        def emit_backward(op_idx: int, warmup_offset: int):
+            local_idx = get_bwd_local_idx(op_idx, warmup_offset)
+            stage = get_global_stage(local_idx)
+            mb = bwd_counters[stage]
+
+            # In Zero Bubble, we split: Backward Input (Now) + Backward Weight (Later)
+            # In Standard 1F1B, we do full backward now.
+            is_full = not self._enable_zero_bubble
+
+            rank_actions.append(
+                BackwardFullInputComputeAction(
+                    stage_idx=stage,
+                    microbatch_idx=mb,
+                    full_backward=is_full
+                )
+            )
+
+            if self._enable_zero_bubble:
+                pending_weights.append((stage, mb))
+
+            bwd_counters[stage] += 1
+
+        def try_emit_weight_zb(op_idx: int, warmup_offset: int):
+            if not self._enable_zero_bubble or not pending_weights:
+                return
+
+            steps_into_1f1b = op_idx - warmup_offset
+            # The earliest reasonable time to start weaving in weights is proportional to rank depth
+            if steps_into_1f1b >= rank:
+                w_stage, w_mb = pending_weights.popleft()
+                rank_actions.append(
+                    BackwardWeightComputeAction(stage_idx=w_stage, microbatch_idx=w_mb)
+                )
+
+        # -- Execution Phase Math --
+
+        warmup_ops = self._get_warmup_ops(
+            rank, microbatches_per_round, pp_size, n_microbatches, multiply_factor
+        )
+        total_microbatch_ops = self._num_stages_per_rank * n_microbatches
+        fwd_bwd_ops = total_microbatch_ops - warmup_ops
+        cooldown_ops = total_microbatch_ops - fwd_bwd_ops
+
+        # Combine into one sequence for iteration, but handle logic per phase
+        total_ops = warmup_ops + fwd_bwd_ops + cooldown_ops
+
+        # -- Main Schedule Loop --
+
+        for op in range(total_ops):
+
+            # Phase 1: Warmup (Forward Only)
+            if op < warmup_ops:
+                emit_forward(op)
+
+            # Phase 2: Steady State (1F1B)
+            elif op < warmup_ops + fwd_bwd_ops:
+                emit_forward(op)
+                emit_backward(op, warmup_offset=warmup_ops)
+                try_emit_weight_zb(op, warmup_offset=warmup_ops)
+
+            # Phase 3: Cooldown (Backward Only)
+            else:
+                emit_backward(op, warmup_offset=warmup_ops)
+                try_emit_weight_zb(op, warmup_offset=warmup_ops)
+
+        # -- Post-Loop: Flush Remaining Weights (ZB only) --
+        while pending_weights:
+            w_stage, w_mb = pending_weights.popleft()
+            rank_actions.append(
+                BackwardWeightComputeAction(stage_idx=w_stage, microbatch_idx=w_mb)
+            )
+
+        return rank_actions
+
+    @property
+    def num_stages_per_rank(self) -> int:
+        return self._num_stages_per_rank
+
+    @property
+    def topology_style(self) -> ScheduleStyle:
+        return ScheduleStyle.loop