[Feat] Keep FSDP overlap when per-rank graphs diverge in compute

magi_compiler/_magi_register_custom_op.py

@@ -886,6 +886,14 @@ class _DataclassRuntimeAdapter:
 # ==============================================================================
 
 
+def _maybe_register_op_profiling(op_name: str, has_internal_collective: bool, materialize_inputs: Callable | None) -> None:
+    if not has_internal_collective and materialize_inputs is None:
+        return
+    from magi_compiler.profiling import register_materialize_inputs
+
+    register_materialize_inputs(op_name, materialize_inputs, has_internal_collective=has_internal_collective)
+
+
 def _magi_register_custom_op_impl(
     name: str | None = None,
     mutates_args: tuple[str, ...] = (),
@@ -894,6 +902,8 @@ def _magi_register_custom_op_impl(
     backward_fn: Callable | None = None,
     is_compute_sensitive: bool = False,
     is_subgraph_boundary: bool = False,
+    has_internal_collective: bool = False,
+    materialize_inputs: Callable | None = None,
 ):
     def decorator(fn: Callable) -> Callable:
         # A 4-slot pipeline.
@@ -903,6 +913,7 @@ def decorator(fn: Callable) -> Callable:
             get_compile_config().recompute_config.custom_compute_sensitive_ops.append(op_name)
         if is_subgraph_boundary:
             get_compile_config().splitting_ops.append(op_name)
+        _maybe_register_op_profiling(op_name, has_internal_collective, materialize_inputs)
 
         _validate_op_signature_constraints(fn)
         original_sig, lowered_sig, param_mapping_tree = _lower_op_signature(fn)

magi_compiler/api.py

@@ -202,6 +202,8 @@ def magi_register_custom_op(
     backward_fn: Callable | None = None,
     is_compute_sensitive: bool = False,
     is_subgraph_boundary: bool = False,
+    has_internal_collective: bool = False,
+    materialize_inputs: Callable | None = None,
 ):
     """
     A unified decorator to register a custom operator with PyTorch's library.
@@ -233,6 +235,14 @@ def magi_register_custom_op(
             ops are prioritised for saving rather than recomputing.
         is_subgraph_boundary: Split the FX graph at this op during compilation.
             Each sub-graph between boundary ops is compiled independently.
+        has_internal_collective: The op issues NCCL internally (CP all-to-all,
+            EP dispatch, ...).  The profiler then lockstep-replays
+            the whole custom op with a fixed iteration count.
+        materialize_inputs: Optional hook with the **same signature as the op**.
+            MagiCompiler generic-realizes every argument, then calls
+            ``fn(*args, **kwargs)``. Use this to rebuild value-dependent metadata
+            (``cp_split_sizes``, ``cu_seqlens``, ...) from the tensors the
+            op actually receives.
 
     Returns:
         A callable with the user's original signature.
@@ -313,4 +323,6 @@ def magi_register_custom_op(
         backward_fn=backward_fn,
         is_compute_sensitive=is_compute_sensitive,
         is_subgraph_boundary=is_subgraph_boundary,
+        has_internal_collective=has_internal_collective,
+        materialize_inputs=materialize_inputs,
     )

magi_compiler/config.py

@@ -214,11 +214,10 @@ class FSDPConfig(BaseModel):
     cost_mode: Literal["profile_sync", "analytical"] = Field(
         "profile_sync",
         description=(
-            "Cost model for the reorder placement; must be rank-identical multi-rank (else NCCL deadlock). "
-            "'profile_sync' (default): real per-op profiling, re-measured in rank-lockstep and max-reduced; "
-            "requires structurally identical per-rank graphs (verified at runtime, degrades safely on "
-            "mismatch). 'analytical': Inductor roofline -- rank-deterministic, less accurate, deadlock-free "
-            "fallback."
+            "Cost model for the reorder placement. "
+            "'profile_sync' (default): real per-op profiling; entries shared across ranks "
+            "are re-measured in rank-lockstep and max-reduced (works even when per-rank graphs diverge). "
+            "'analytical': Inductor roofline -- rank-deterministic, less accurate, deadlock-free fallback."
         ),
     )
     comm_overlap_window_margin_ns: float = Field(

magi_compiler/passes/fsdp_overlap/reorder.py

@@ -36,10 +36,12 @@
 (1 packed launch + N MultiOutput members moved together as one block + N waits).
 """
 
+import bisect
 import hashlib
 from collections import defaultdict
 
 import torch
+import torch.distributed as dist
 from torch._inductor.comms import _is_fake_dep
 from torch._inductor.ir import MultiOutput
 from torch._inductor.scheduler import BaseSchedulerNode
@@ -91,6 +93,29 @@ def _size_hint_of(sym) -> int:
         return 0
 
 
+def _collective_kind_key(snode: BaseSchedulerNode) -> tuple:
+    """Coarse, rank-comparable identity of one NCCL-issuing snode."""
+    node = _leaf_collective_node(snode) or getattr(snode, "node", None)
+    op = getattr(node, "op_overload", None) or getattr(node, "python_kernel_name", None) or type(node).__name__
+    dims: tuple = ()
+    try:
+        dims = tuple("?" if getattr(d, "free_symbols", None) else int(d) for d in node.get_size())
+    except Exception:  # noqa: BLE001
+        pass
+    return (_is_weight_gather(snode), str(op), dims)
+
+
+def _collective_skeleton(order: list[BaseSchedulerNode]) -> tuple[list[int], list[tuple]]:
+    """The graph's collective skeleton: indices (ascending) and rank-comparable
+    kinds of every snode that ISSUES NCCL -- functional collectives plus custom ops
+    with an internal collective . This sequence is what must stay rank-identical;
+    the compute between two consecutive entries is rank-private."""
+    from magi_compiler.profiling.runtime_estimator import snode_issues_collective
+
+    idx = [i for i, s in enumerate(order) if snode_issues_collective(s)]
+    return idx, [_collective_kind_key(order[i]) for i in idx]
+
+
 def _graph_fingerprint(order: list[BaseSchedulerNode]) -> str:
     """Rank-comparable digest of the snode sequence: type + op identity + output
     sizes + sorted origin fx TARGETS.  Origins are required -- a fused pointwise
@@ -211,34 +236,14 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]:
 
         index_of = {s: i for i, s in enumerate(order)}
 
-        # Fail-fast: the index-based sweep (and the lockstep profiling below) both
-        # require structurally IDENTICAL per-rank graphs, else the weight gathers
-        # interleave with other collectives in rank-divergent order -> NCCL
-        # deadlock.  Verify with one symmetric all_gather of a graph digest; every
-        # rank sees the same result, so all ranks take the same branch.
-        import torch.distributed as dist
-
-        if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
-            from magi_compiler.profiling.runtime_estimator import _get_cost_sync_group
-
-            fp = (_graph_fingerprint(order), len(order), len(launches))
-            world = dist.get_world_size()
-            all_fp: list = [None] * world
-            dist.all_gather_object(all_fp, fp, group=_get_cost_sync_group())
-            if any(f != all_fp[0] for f in all_fp[1:]):
-                magi_logger.warning(
-                    "FSDP overlap reorder: per-rank graphs are NOT structurally identical "
-                    "((digest, n_snodes, n_weight_gathers) per rank: %s). Reordering would "
-                    "produce rank-divergent collective order -> NCCL deadlock; leaving the "
-                    "graph unchanged (overlap OFF for this graph). Likely cause: uneven "
-                    "Shard(0) params -- replicate them or use chunk-padded uniform shards.",
-                    [(f[0][:12], f[1], f[2]) if f else None for f in all_fp],
-                )
-                return order
+        skel_idx, skel_kinds = _collective_skeleton(order)
+        mode, sync_group, world = self._negotiate_mode(order, launches, skel_kinds)
+        if mode == "abort":
+            return order
 
         # profile_sync: warm the estimator table on every node, then re-measure in
-        # rank-lockstep (warm_and_sync) so costs are rank-identical.  On failure,
-        # leave the graph unchanged (overlap off, no hang).
+        # rank-lockstep (warm_and_sync) so shared keys get real, max-reduced costs.
+        # On failure, leave the graph unchanged (overlap off, no hang).
         if hasattr(self._cost_fn, "warm_and_sync") and getattr(self._cost_fn, "_sync_across_ranks", False):
             try:
                 for s in order:
@@ -257,13 +262,19 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]:
         launches_in_order = sorted(launches, key=lambda s: index_of[s])  # original program order
 
         plans = []  # (launch, group, fc_idx, comm_runtime, lower)
+        lowers: dict = {}  # launch -> earliest legal index (real-dep floor)
         for launch in launches_in_order:
             group = self._launch_group(launch, order, buf_to_snode, users)
             fc_idx = self._first_consumer_index(launch, group, order, users)
             if fc_idx is None:
                 continue
             comm_runtime = self._cost(launch)
             lower = self._earliest_legal_index(group, order, index_of, buf_to_snode, op_to_snode)
+            if mode == "pinned":
+                # No skeleton to negotiate against: keep the AG between the same two
+                # NCCL-issuing snodes it already sat between.
+                lower = self._raise_lower_for_nccl_barriers(lower, index_of[launch], order)
+            lowers[launch] = lower
             plans.append((launch, group, fc_idx, comm_runtime, lower))
 
         targets: dict = {}  # launch -> target index (in original order space)
@@ -312,12 +323,11 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]:
                 "hidden" if acc >= need else "COMPUTE-LIMITED",
             )
 
-        # Clamp targets NON-DECREASING in original program order.  NCCL matches the
-        # Nth call on a PG positionally across ranks, so the gathers' relative order
-        # must be rank-identical; `max(lower, t)` can invert two gathers and whether
-        # the inversion happens depends on per-rank cost jitter -> deadlock.  The
-        # clamp pins the original subsequence at the cost of occasionally placing a
-        # launch later than its compute window would allow.
+        if world > 1 and mode != "pinned":
+            self._consensus_slot_targets(targets, lowers, launches_in_order, skel_idx, index_of, sync_group, world)
+
+        # Keep gather order: clamp targets non-decreasing so cost jitter or
+        # same-slot per-rank indices cannot swap two launches.
         running = -1
         for launch in launches_in_order:
             if launch not in targets:
@@ -345,10 +355,19 @@ def _key(s):
 
         new_order = sorted(order, key=_key)
         # Validate the rebuilt order is a valid topological order; only commit if so.
-        if self._validate_full(new_order, op_to_snode, buf_to_snode, users):
+        # The verdict is reduced across ranks: committing on some ranks and not on
+        # others is itself a divergent NCCL sequence.
+        ok = self._validate_full(new_order, op_to_snode, buf_to_snode, users)
+        if not ok:
+            magi_logger.warning("FSDP overlap reorder: rebuilt order failed validation; leaving graph unchanged")
+        if self._agree(ok, sync_group, world):
             order[:] = new_order
         else:
-            magi_logger.warning("FSDP overlap reorder: rebuilt order failed validation; leaving graph unchanged")
+            if ok:
+                magi_logger.warning(
+                    "FSDP overlap reorder: another rank did not commit its rebuilt order; "
+                    "leaving this rank's graph unchanged too"
+                )
             moved = 0
 
         measured = getattr(self._cost_fn, "n_measured", None)
@@ -369,6 +388,108 @@ def _key(s):
             magi_logger.debug("FSDP overlap %s", self._cost_fn.summary())
         return order
 
+    # -- multi-rank agreement ---------------------------------------------
+    @staticmethod
+    def _negotiate_mode(order, launches, skel_kinds) -> tuple[str, object, int]:
+        """Rank-identical placement mode: (mode, group, world).
+
+        ``identical`` / ``slot``: skeletons match → consensus slots (in-slot index is per-rank).
+        ``pinned``: skeletons differ → keep each AG between its neighboring NCCL snodes.
+        ``abort``: weight-AG counts differ → leave the graph unchanged.
+        """
+        from magi_compiler.profiling.runtime_estimator import _get_cost_sync_group
+
+        group = _get_cost_sync_group()
+        world = dist.get_world_size()
+        mine = ((_graph_fingerprint(order), len(order), len(launches)), tuple(skel_kinds))
+        peers: list = [None] * world
+        dist.all_gather_object(peers, mine, group=group)
+        if all(p == peers[0] for p in peers[1:]):
+            return "identical", group, world
+
+        desc = [(p[0][0][:12], p[0][1], p[0][2], len(p[1])) for p in peers]
+        n_ag = [p[0][2] for p in peers]
+        if any(g != n_ag[0] for g in n_ag[1:]):
+            magi_logger.warning(
+                "FSDP overlap reorder: per-rank graphs differ AND weight-AG counts diverge "
+                "((digest, n_snodes, n_weight_gathers, n_collectives) per rank: %s). No rank "
+                "correspondence to reconcile; leaving the graph unchanged (overlap OFF).",
+                desc,
+            )
+            return "abort", group, world
+        if all(p[1] == peers[0][1] for p in peers[1:]):
+            magi_logger.warning(
+                "FSDP overlap reorder: per-rank graphs are NOT structurally identical "
+                "((digest, n_snodes, n_weight_gathers, n_collectives) per rank: %s), but the "
+                "collective skeleton matches. Continuing in SLOT-consensus mode: gathers are "
+                "placed in a rank-negotiated skeleton slot (they MAY hop CP / EP kernels, as "
+                "long as every rank hops the same one).",
+                desc,
+            )
+            return "slot", group, world
+        magi_logger.warning(
+            "FSDP overlap reorder: per-rank graphs are NOT structurally identical AND their "
+            "collective skeletons differ ((digest, n_snodes, n_weight_gathers, n_collectives) "
+            "per rank: %s). Continuing in PINNED mode: gathers keep their position relative to "
+            "every NCCL-issuing snode (no hop over CP / EP kernels).",
+            desc,
+        )
+        return "pinned", group, world
+
+    @staticmethod
+    def _agree(ok: bool, sync_group, world: int) -> bool:
+        """Reduce a local yes/no into a rank-identical one (AND over ranks)."""
+        if world <= 1:
+            return ok
+
+        try:
+            t = torch.tensor([1 if ok else 0], dtype=torch.int32)
+            dist.all_reduce(t, op=dist.ReduceOp.MIN, group=sync_group)
+            return bool(t.item())
+        except Exception as exc:  # noqa: BLE001
+            magi_logger.warning("FSDP overlap reorder: cross-rank agreement failed (%s); leaving graph unchanged", exc)
+            return False
+
+    @staticmethod
+    def _consensus_slot_targets(targets, lowers, launches_in_order, skel_idx, index_of, sync_group, world) -> None:
+        """Put each gather in the same skeleton slot on every rank (max of desired
+        slot and dep floor, then non-decreasing).  Index inside the slot stays local.
+        """
+
+        def slot_of(idx: int) -> int:
+            return bisect.bisect_left(skel_idx, idx)
+
+        # One entry per launch in program order so skipped gathers stay aligned.
+        mine = []
+        for launch in launches_in_order:
+            own = slot_of(index_of[launch])
+            mine.append((slot_of(targets[launch][0]), slot_of(lowers[launch])) if launch in targets else (own, own))
+        peers: list = [None] * world
+        dist.all_gather_object(peers, mine, group=sync_group)
+
+        running = 0
+        for j, launch in enumerate(launches_in_order):
+            q = max(max(p[j][0] for p in peers), max(p[j][1] for p in peers))
+            q = running = max(q, running)
+            if launch not in targets:
+                continue
+            target, group = targets[launch]
+            slot_lo = max(lowers[launch], skel_idx[q - 1] + 1 if q > 0 else 0)
+            slot_hi = skel_idx[q] if q < len(skel_idx) else index_of[launch]
+            new_target = min(max(target, slot_lo), slot_hi)
+            targets[launch] = (new_target, group)
+            magi_logger.debug(
+                "FSDP overlap slot consensus: launch cur=%d slot=%d/%d (mine=%s) target %d -> %d [%d, %d]",
+                index_of[launch],
+                q,
+                slot_of(index_of[launch]),
+                mine[j],
+                target,
+                new_target,
+                slot_lo,
+                slot_hi,
+            )
+
     # -- group detection --------------------------------------------------
     def _launch_group(self, launch, order, buf_to_snode, users) -> list[BaseSchedulerNode]:
         """The snodes that must move together with the launch.
@@ -444,6 +565,18 @@ def _is_transparent(self, snode: BaseSchedulerNode) -> bool:
         return self._cost(snode) <= 1.0
 
     # -- repositioning ----------------------------------------------------
+    @staticmethod
+    def _raise_lower_for_nccl_barriers(lower: int, launch_idx: int, order: list) -> int:
+        """Raise ``lower`` so a weight AG cannot hop any NCCL-issuing snode that
+        originally precedes it."""
+        from magi_compiler.profiling.runtime_estimator import snode_issues_collective
+
+        barrier = lower
+        for i in range(lower, launch_idx):
+            if snode_issues_collective(order[i]):
+                barrier = i + 1
+        return barrier
+
     def _earliest_legal_index(self, group, order, index_of, buf_to_snode, op_to_snode) -> int:
         """1 + max index of any REAL (non-fake buffer) producer the group needs.
 

magi_compiler/profiling/__init__.py

@@ -13,7 +13,18 @@
 # limitations under the License.
 
 
-from .benchmark_inputs import get_benchmark_inputs_hook, op_has_internal_collective, register_benchmark_inputs
+from .materialize_inputs import (
+    apply_materialize_inputs,
+    get_materialize_inputs_hook,
+    op_has_internal_collective,
+    register_materialize_inputs,
+)
 from .runtime_estimator import ProfilingRuntimeEstimator
 
-__all__ = ["ProfilingRuntimeEstimator", "register_benchmark_inputs", "get_benchmark_inputs_hook", "op_has_internal_collective"]
+__all__ = [
+    "ProfilingRuntimeEstimator",
+    "register_materialize_inputs",
+    "get_materialize_inputs_hook",
+    "op_has_internal_collective",
+    "apply_materialize_inputs",
+]