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))