pp-simulator/pp_simulator/dependencies.py

from __future__ import annotations

from collections import defaultdict, deque

from .plan import Plan, PlanOp


class DependencyGraph:
    def __init__(self, plan: Plan):
        self.plan = plan
        self.dependency_reasons = self._build_dependency_reasons()
        self.dependencies = {
            op_id: set(reasons_by_dep)
            for op_id, reasons_by_dep in self.dependency_reasons.items()
        }
        self.topological_order = self._topological_sort()

    def _build_dependency_reasons(self) -> dict[str, dict[str, set[str]]]:
        deps: dict[str, dict[str, set[str]]] = {op.id: {} for op in self.plan.ops}
        by_identity = {
            self._identity(op.rank, op.microbatch_id, op.chunk_id, op.op_type): op
            for op in self.plan.ops
        }
        by_type_and_batch: dict[tuple[str, int], list[PlanOp]] = defaultdict(list)

        for op in self.plan.ops:
            if op.microbatch_id is not None:
                by_type_and_batch[(op.op_type, self._batch_id(op.microbatch_id))].append(op)

        # Local rank order is always a hard execution constraint.
        for rank_ops in self.plan.ops_by_rank.values():
            for previous, current in zip(rank_ops, rank_ops[1:]):
                self._add_dep(current, previous, deps, "rank_order")

        for op in self.plan.ops:
            if op.microbatch_id is None:
                continue

            if op.op_type == "F":
                self._add_llm_forward_deps(op, deps, by_identity, by_type_and_batch)
            elif op.op_type == "B":
                self._add_llm_backward_deps(op, deps, by_identity, by_type_and_batch)
            elif op.op_type == "GF":
                if self._has_disaggregated_layout():
                    self._add_identity_dep(
                        op,
                        deps,
                        by_identity,
                        self._last_llm_rank(),
                        op.microbatch_id,
                        self.plan.vpp_size - 1,
                        "F",
                        "llm_output_ready",
                    )
                else:
                    self._add_identity_dep(
                        op,
                        deps,
                        by_identity,
                        self._last_llm_rank(),
                        op.microbatch_id,
                        self.plan.vpp_size - 1,
                        "F",
                        "llm_output_ready",
                    )
            elif op.op_type == "GB":
                self._add_identity_dep(
                    op,
                    deps,
                    by_identity,
                    op.rank,
                    op.microbatch_id,
                    self.plan.vpp_size - 1,
                    "GF",
                    "generator_backward_consumes_forward",
                )
            elif op.op_type == "VB":
                if self._has_disaggregated_layout():
                    self._add_identity_dep(
                        op,
                        deps,
                        by_identity,
                        self._first_llm_rank(),
                        op.microbatch_id,
                        0,
                        "B",
                        "llm_input_gradient_ready",
                    )
                else:
                    self._add_identity_dep(
                        op,
                        deps,
                        by_identity,
                        self._first_llm_rank(),
                        op.microbatch_id,
                        0,
                        "B",
                        "llm_input_gradient_ready",
                    )

        return deps

    def _add_llm_forward_deps(
        self,
        op: PlanOp,
        deps: dict[str, dict[str, set[str]]],
        by_identity: dict[tuple[int, int | None, int | None, str], PlanOp],
        by_type_and_batch: dict[tuple[str, int], list[PlanOp]],
    ) -> None:
        chunk = self._chunk(op)
        first_llm_rank = self._first_llm_rank()
        last_llm_rank = self._last_llm_rank()

        if op.rank > first_llm_rank:
            self._add_identity_dep(
                op,
                deps,
                by_identity,
                op.rank - 1,
                op.microbatch_id,
                chunk,
                "F",
                "llm_forward_from_previous_rank",
            )
        elif chunk > 0:
            self._add_identity_dep(
                op,
                deps,
                by_identity,
                last_llm_rank,
                op.microbatch_id,
                chunk - 1,
                "F",
                "vpp_forward_from_previous_chunk",
            )

        if op.rank == first_llm_rank and chunk == 0:
            if self._has_disaggregated_layout():
                self._add_identity_dep(
                    op,
                    deps,
                    by_identity,
                    self._encoder_stage(),
                    op.microbatch_id,
                    0,
                    "VF",
                    "encoder_forward_ready",
                )
            else:
                self._add_identity_dep(
                    op,
                    deps,
                    by_identity,
                    self._ve_rank(op.microbatch_id),
                    op.microbatch_id,
                    0,
                    "VF",
                    "ve_forward_ready",
                )

    def _add_llm_backward_deps(
        self,
        op: PlanOp,
        deps: dict[str, dict[str, set[str]]],
        by_identity: dict[tuple[int, int | None, int | None, str], PlanOp],
        by_type_and_batch: dict[tuple[str, int], list[PlanOp]],
    ) -> None:
        chunk = self._chunk(op)
        first_llm_rank = self._first_llm_rank()
        last_llm_rank = self._last_llm_rank()
        self._add_identity_dep(
            op,
            deps,
            by_identity,
            op.rank,
            op.microbatch_id,
            chunk,
            "F",
            "backward_consumes_forward_activation",
        )

        if op.rank < last_llm_rank:
            self._add_identity_dep(
                op,
                deps,
                by_identity,
                op.rank + 1,
                op.microbatch_id,
                chunk,
                "B",
                "llm_backward_from_next_rank",
            )
        elif chunk + 1 < self.plan.vpp_size:
            self._add_identity_dep(
                op,
                deps,
                by_identity,
                first_llm_rank,
                op.microbatch_id,
                chunk + 1,
                "B",
                "vpp_backward_from_next_chunk",
            )

        if op.rank == last_llm_rank and chunk == self.plan.vpp_size - 1:
            if self._has_disaggregated_layout() and self._has_generator():
                self._add_identity_dep(
                    op,
                    deps,
                    by_identity,
                    self._generator_stage(),
                    op.microbatch_id,
                    0,
                    "GB",
                    "generator_backward_ready",
                )
            else:
                self._add_identity_dep(
                    op,
                    deps,
                    by_identity,
                    self._ve_rank(op.microbatch_id),
                    op.microbatch_id,
                    0,
                    "GB",
                    "generator_backward_ready",
                )

    def _add_batch_deps(
        self,
        op: PlanOp,
        deps: dict[str, dict[str, set[str]]],
        by_type_and_batch: dict[tuple[str, int], list[PlanOp]],
        dep_type: str,
        *,
        rank: int | None = None,
        chunk: int | None = None,
        final_forward: bool = False,
        reason: str,
    ) -> None:
        batch_id = self._batch_id(op.microbatch_id)
        for candidate in by_type_and_batch.get((dep_type, batch_id), []):
            if rank is not None and candidate.rank != rank:
                continue
            if chunk is not None and self._chunk(candidate) != chunk:
                continue
            if final_forward and not (
                candidate.rank == self.plan.pp_size - 1
                and self._chunk(candidate) == self.plan.vpp_size - 1
            ):
                continue
            if candidate.id != op.id:
                self._add_dep(op, candidate, deps, reason)

    @staticmethod
    def _add_identity_dep(
        op: PlanOp,
        deps: dict[str, dict[str, set[str]]],
        by_identity: dict[tuple[int, int | None, int | None, str], PlanOp],
        rank: int,
        microbatch_id: int | None,
        chunk_id: int | None,
        op_type: str,
        reason: str,
    ) -> None:
        dependency = by_identity.get((rank, microbatch_id, chunk_id, op_type))
        if dependency is not None and dependency.id != op.id:
            DependencyGraph._add_dep(op, dependency, deps, reason)

    @staticmethod
    def _add_dep(
        op: PlanOp,
        dependency: PlanOp,
        deps: dict[str, dict[str, set[str]]],
        reason: str,
    ) -> None:
        deps[op.id].setdefault(dependency.id, set()).add(reason)

    def _topological_sort(self) -> list[str]:
        reverse_edges: dict[str, set[str]] = defaultdict(set)
        indegree = {op.id: len(self.dependencies[op.id]) for op in self.plan.ops}

        for op_id, dep_ids in self.dependencies.items():
            for dep_id in dep_ids:
                reverse_edges[dep_id].add(op_id)

        queue = deque(op.id for op in self.plan.ops if indegree[op.id] == 0)
        order: list[str] = []

        while queue:
            op_id = queue.popleft()
            order.append(op_id)
            for child_id in sorted(reverse_edges.get(op_id, ())):
                indegree[child_id] -= 1
                if indegree[child_id] == 0:
                    queue.append(child_id)

        if len(order) != len(indegree):
            cyclic_ids = sorted(op_id for op_id, degree in indegree.items() if degree > 0)
            sample = ", ".join(cyclic_ids[:8])
            raise ValueError(f"Dependency graph contains a cycle. Sample blocked ops: {sample}")

        return order

    def _batch_id(self, microbatch_id: int | None) -> int:
        if microbatch_id is None:
            return -1
        return int(microbatch_id) // max(1, self._llm_pp_size())

    @staticmethod
    def _chunk(op: PlanOp) -> int:
        return 0 if op.chunk_id is None else int(op.chunk_id)

    @staticmethod
    def _identity(
        rank: int,
        microbatch_id: int | None,
        chunk_id: int | None,
        op_type: str,
    ) -> tuple[int, int | None, int | None, str]:
        return (rank, microbatch_id, 0 if chunk_id is None else chunk_id, op_type)

    def _first_llm_rank(self) -> int:
        layout = self.plan.pipeline_layout or {}
        return int(layout.get("llm_stage_offset", 0))

    def _llm_pp_size(self) -> int:
        layout = self.plan.pipeline_layout or {}
        return int(layout.get("llm_pp_size", self.plan.pp_size))

    def _last_llm_rank(self) -> int:
        return self._first_llm_rank() + self._llm_pp_size() - 1

    def _has_disaggregated_layout(self) -> bool:
        return bool(self.plan.pipeline_layout)

    def _has_generator(self) -> bool:
        layout = self.plan.pipeline_layout
        if not layout:
            return True
        return bool(layout.get("has_generator", "generator_stage" in layout))

    def _ve_rank(self, microbatch_id: int | None) -> int:
        if microbatch_id is None:
            return 0
        return int(microbatch_id) % self.plan.pp_size

    def _encoder_stage(self) -> int:
        layout = self.plan.pipeline_layout or {}
        return int(layout.get("encoder_stage", 0))

    def _generator_stage(self) -> int:
        layout = self.plan.pipeline_layout or {}
        return int(layout.get("generator_stage", self.plan.pp_size - 1))