diff --git a/magi_compiler/_magi_register_custom_op.py b/magi_compiler/_magi_register_custom_op.py index 2fa1d7d..490c1ee 100644 --- a/magi_compiler/_magi_register_custom_op.py +++ b/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) diff --git a/magi_compiler/api.py b/magi_compiler/api.py index 3f94989..99602ea 100644 --- a/magi_compiler/api.py +++ b/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, ) diff --git a/magi_compiler/config.py b/magi_compiler/config.py index baf6e1c..f0206ac 100644 --- a/magi_compiler/config.py +++ b/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( diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index 7309264..3c80ab1 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/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,6 +262,7 @@ 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) @@ -264,6 +270,11 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: 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. diff --git a/magi_compiler/profiling/__init__.py b/magi_compiler/profiling/__init__.py index 467b3bf..554bb99 100644 --- a/magi_compiler/profiling/__init__.py +++ b/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", +] diff --git a/magi_compiler/profiling/benchmark_inputs.py b/magi_compiler/profiling/benchmark_inputs.py deleted file mode 100644 index 5f7bcae..0000000 --- a/magi_compiler/profiling/benchmark_inputs.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) 2026 SandAI. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Per-op benchmark-input hooks for the profiling runtime estimator. - -Some custom ops cannot be replayed from generic size-hinted tensors: they carry -VALUE-DEPENDENT metadata that must be self-consistent or they raise (e.g. a CP -attention op whose split-sizes must sum to the sequence length) -- and would fall -back to a 0 cost. The model package that defines such an op registers a hook -here that builds valid replay inputs; MagiCompiler stays free of model-specific -op knowledge. - -Hook: ``fn(fx_node, realize) -> (args, kwargs) | None`` -- ``fx_node`` is the op's -``torch.fx.Node`` (shapes in ``meta['val']``), ``realize`` is the generic arg -realizer to reuse for plain tensor args; return None to fall back to the generic -path. Register at import time of the op-defining module, e.g.:: - - register_benchmark_inputs("mylib::attn_cp", _attn_cp_inputs, has_internal_collective=True) - -Hooks MUST produce rank-identical inputs (derive everything from shapes, no -per-rank state) so the rank-lockstep ``warm_and_sync`` measurement issues any -internal collective in lockstep. -""" - -from __future__ import annotations - -from typing import Callable - -# op name (OpOverload string, e.g. "mylib::attn_cp") -> hook; plus the set of ops -# that issue an internal collective (need fixed-iter lockstep replay). -_BENCHMARK_INPUT_HOOKS: dict[str, Callable] = {} -_INTERNAL_COLLECTIVE_OPS: set[str] = set() - - -def register_benchmark_inputs(op_name: str, fn: Callable, *, has_internal_collective: bool = False) -> None: - """Register a replay-input builder for ``op_name`` (see module docstring). - ``has_internal_collective``: replay with a fixed iteration count under barriers - (an adaptive count would desync the internal NCCL op across ranks).""" - _BENCHMARK_INPUT_HOOKS[op_name] = fn - if has_internal_collective: - _INTERNAL_COLLECTIVE_OPS.add(op_name) - - -def get_benchmark_inputs_hook(op_name: str) -> Callable | None: - return _BENCHMARK_INPUT_HOOKS.get(op_name) - - -def op_has_internal_collective(op_name: str) -> bool: - return op_name in _INTERNAL_COLLECTIVE_OPS diff --git a/magi_compiler/profiling/materialize_inputs.py b/magi_compiler/profiling/materialize_inputs.py new file mode 100644 index 0000000..4214b55 --- /dev/null +++ b/magi_compiler/profiling/materialize_inputs.py @@ -0,0 +1,95 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-op materialize-input hooks for the profiling runtime estimator. + +Some custom ops cannot be replayed from generic size-hinted tensors: they carry +VALUE-DEPENDENT metadata that must be self-consistent or they raise (e.g. a CP +attention op whose split-sizes must sum to the sequence length) -- and would fall +back to a 0 cost. The model package that defines such an op registers a hook +that rebuilds valid replay inputs from the same arguments the custom op sees; +MagiCompiler stays free of model-specific op knowledge. + +Hook signature matches the custom op. MagiCompiler generic-realizes every +argument first, then calls:: + + hook(*realized_args, **realized_kwargs) -> tuple | None + +Return a positional-arg tuple to replay (same slots as the op), or ``None`` to +keep the generic realize. Prefer attaching the hook on the op decorator:: + + def _attn_cp_inputs(q, k, v, cp_split_sizes): + seq = int(q.shape[1] if q.dim() == 4 else q.shape[0]) + return q, k, v, [seq] * len(cp_split_sizes) + + @magi_register_custom_op("mylib::attn_cp", materialize_inputs=_attn_cp_inputs, has_internal_collective=True) + def attn_cp(q, k, v, cp_split_sizes): + ... + +Hooks MUST produce rank-identical inputs (derive everything from shapes, no +per-rank state) so the rank-lockstep ``warm_and_sync`` measurement issues any +internal collective in lockstep. +""" + +from __future__ import annotations + +from typing import Callable + +# TODO: Find a better way to solve the materialize_inputs problem. The current +# per-op hook is a workaround for custom ops whose replay inputs cannot be +# reconstructed from generic size-hinted tensors (value-dependent metadata). +# Look for a more general approach that does not require model-side hooks. + +# op name (OpOverload string, e.g. "mylib::attn_cp") -> hook; plus the set of ops +# that issue an internal collective (need fixed-iter lockstep replay). +_MATERIALIZE_INPUT_HOOKS: dict[str, Callable] = {} +_INTERNAL_COLLECTIVE_OPS: set[str] = set() + + +def register_materialize_inputs(op_name: str, fn: Callable | None = None, *, has_internal_collective: bool = False) -> None: + """Register a same-signature replay-input builder for ``op_name``. + + ``fn`` is optional when only ``has_internal_collective`` is needed (generic + realize is already valid). ``has_internal_collective``: replay with a fixed + iteration count under barriers (an adaptive count would desync the internal + NCCL op across ranks). + """ + if fn is not None: + _MATERIALIZE_INPUT_HOOKS[op_name] = fn + if has_internal_collective: + _INTERNAL_COLLECTIVE_OPS.add(op_name) + + +def get_materialize_inputs_hook(op_name: str) -> Callable | None: + return _MATERIALIZE_INPUT_HOOKS.get(op_name) + + +def op_has_internal_collective(op_name: str) -> bool: + return op_name in _INTERNAL_COLLECTIVE_OPS + + +def apply_materialize_inputs(hook: Callable | None, args: tuple, kwargs: dict) -> tuple[tuple, dict]: + """Run ``hook`` on already-realized custom-op inputs. + + ``hook`` has the same signature as the op. ``None`` keeps ``(args, kwargs)``; + a tuple/list replaces the positional args used for replay. + """ + if hook is None: + return args, kwargs + out = hook(*args, **kwargs) + if out is None: + return args, kwargs + if isinstance(out, (tuple, list)): + return tuple(out), {} + raise TypeError(f"materialize_inputs hook must return None or a tuple of op args, got {type(out).__name__}") diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index 3ba6ca7..c6e3f56 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -27,6 +27,13 @@ NCCL desyncs ranks -> hang); they are seeded with the analytical estimate and re-measured for real in the rank-lockstep ``warm_and_sync``. +``warm_and_sync`` lockstep-measures the INTERSECTION of structural keys present +on every rank (with a stashed snode): full key-set identity is not required, so +graphs that diverge in rank-local compute still get real costs for shared +kernels (weight AG and isomorphic custom ops). Rank-local pure-compute keeps +its warmup measurement; rank-local collective-bearing keys fall back to +analytical. + The op->time table (``self._table``) is keyed by STRUCTURAL identity (op + input shapes/dtypes, ``_structural_key``) -- not the per-node name -- so isomorphic ops across layers share one measurement: O(#distinct kernels), not @@ -48,13 +55,25 @@ from magi_compiler.utils import magi_logger -from .benchmark_inputs import get_benchmark_inputs_hook, op_has_internal_collective +from .materialize_inputs import apply_materialize_inputs, get_materialize_inputs_hook, op_has_internal_collective # Dedicated GLOO (CPU) group for the cost sync, built once -- keeps it off the # NCCL process groups the forward uses (cannot desync weight-gather / CP comms). _COST_SYNC_GROUP = "uninit" +def snode_issues_collective(snode: BaseSchedulerNode) -> bool: + """ + True if replaying / running this snode issues NCCL (collective AG, or a custom + op registered with ``has_internal_collective``). + """ + if contains_collective(snode): + return True + if not isinstance(snode, ExternKernelSchedulerNode): + return False + return _extern_has_internal_collective(snode) + + def _get_cost_sync_group(): global _COST_SYNC_GROUP import torch.distributed as dist @@ -163,6 +182,26 @@ def _concrete_size(s, fallback: int = 1) -> int: return int(s) +# dtypes with no ``randn`` / ``normal_`` kernel (e.g. float8). Replay tensors are +# built as fp32 noise then cast -- values only need to be finite for kernel launch; +# the cost model cares about launch time, not numeric fidelity. +_FLOAT8_DTYPES: frozenset = frozenset( + dt + for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "float8_e8m0fnu") + if (dt := getattr(torch, name, None)) is not None +) + + +def _replay_tensor(shape: tuple, device, dtype) -> torch.Tensor: + """Concrete replay tensor matching ``shape/device/dtype``. float8 has no + ``randn`` kernel -- cast from a float32 draw instead.""" + if dtype in _FLOAT8_DTYPES or "float8" in str(dtype): + return torch.randn(shape, device=device, dtype=torch.float32).to(dtype) + if dtype.is_floating_point: + return torch.randn(shape, device=device, dtype=dtype) + return torch.zeros(shape, device=device, dtype=dtype) + + def _realize_arg(v): """fx arg -> concrete replay input: Node(tensor) -> right-shaped tensor from size-hints; SymInt -> concrete int; containers -> recursively realized PLAIN @@ -172,9 +211,7 @@ def _realize_arg(v): ev = v.meta.get("val") if isinstance(ev, torch.Tensor): shape = tuple(_concrete_size(s) for s in ev.shape) - if ev.is_floating_point(): - return torch.randn(shape, device=ev.device, dtype=ev.dtype) - return torch.zeros(shape, device=ev.device, dtype=ev.dtype) + return _replay_tensor(shape, ev.device, ev.dtype) if _is_symbolic(ev) or isinstance(ev, int): return _concrete_size(ev) # a Node carrying a scalar -> concrete hint return v @@ -196,20 +233,16 @@ def _measure_extern(snode: ExternKernelSchedulerNode, fixed_iters: bool = False) (CP all_to_all inside attention/MoE): adaptive iteration counts differ per rank -> NCCL count mismatch -> deadlock. - Ops whose replay needs value-consistent metadata register a hook via - ``benchmark_inputs.register_benchmark_inputs``; otherwise ``_realize_arg``.""" + Replay inputs: generic ``_realize_arg``, then an optional same-signature + hook (``materialize_inputs``) that rebuilds value-consistent metadata.""" fx_node = snode.node.get_origin_node() if fx_node is None: return 0.0 target = fx_node.target - hook = get_benchmark_inputs_hook(_op_name(target)) - built = hook(fx_node, _realize_arg) if hook is not None else None - if built is not None: - args, kwargs = built - else: - args = tuple(_realize_arg(a) for a in fx_node.args) - kwargs = {k: _realize_arg(v) for k, v in fx_node.kwargs.items()} + args = tuple(_realize_arg(a) for a in fx_node.args) + kwargs = {k: _realize_arg(v) for k, v in fx_node.kwargs.items()} + args, kwargs = apply_materialize_inputs(get_materialize_inputs_hook(_op_name(target)), args, kwargs) # Replay eagerly, decoupled from the enclosing compile: # * dynamo.disable: an op whose impl contains torch.compile'd regions would @@ -243,7 +276,7 @@ def fn(): def _op_name(target) -> str: - """Overload-qualified op name (e.g. 'athena::gaga4_fa3_with_sink_cp'), or '' for + """Overload-qualified op name (e.g. 'mylib::attn_cp'), or '' for a non-op target. Used to look ops up in the benchmark-input registry.""" name = getattr(target, "name", None) if callable(name): @@ -256,7 +289,7 @@ def _op_name(target) -> str: def _extern_has_internal_collective(snode: BaseSchedulerNode) -> bool: """Ops that issue collectives internally (declared via - ``register_benchmark_inputs(..., has_internal_collective=True)``) must be + ``register_materialize_inputs(..., has_internal_collective=True)``) must be measured with fixed iterations under a barrier.""" node = getattr(snode, "node", None) origin = node.get_origin_node() if (node is not None and hasattr(node, "get_origin_node")) else None @@ -287,6 +320,7 @@ def _leaf_collective(snode: BaseSchedulerNode): def _collective_spec(node): """(op_overload, group_name, group_size, [(shape, dtype, device), ...]) for a collective IR node, or None if it isn't a benchmarkable all-gather.""" + # TODO: support more graph collectives than AG / AG_COALESCED. op = getattr(node, "op_overload", None) if op not in (_AG, _AG_COALESCED): return None @@ -381,18 +415,19 @@ def table(self) -> "dict[tuple, ProfileEntry]": return self._table def warm_and_sync(self) -> int: - """Rank-lockstep re-measurement of every table entry, giving REAL costs - that are also rank-identical (required for a rank-identical reorder - schedule). Steps: verify key sets match across ranks; per sorted key, - barrier -> measure with FIXED iters -> barrier; finally all_gather the - {key: ns} maps and take the max. Returns #entries whose cost changed. - - The key-set check is a hard precondition: with divergent key sequences the - barrier loop deadlocks (barrier/all_gather count mismatch on gloo, or two - ranks lockstep-measuring DIFFERENT internal-collective ops -> NCCL - mismatch). On mismatch we warn and degrade this compile's entries to the - analytical estimate -- rank-deterministic, same guarantee as - fsdp_config.cost_mode=analytical.""" + """Rank-lockstep re-measurement of table entries shared across ranks. + + Steps: + 1. all_gather key sets; lockstep-measure the INTERSECTION (keys present on + every rank, in a rank-identical sorted order); + 2. max-reduce the measured ns maps; + 3. for keys unique to this rank: keep the local measurement if the op has + no (internal) collective; otherwise fall back to analytical (solo NCCL + replay would hang). + + Full key-set identity is NOT required: per-rank tables may diverge on + rank-local compute, but shared kernels remain isomorphic and are still + lockstep-measured. Returns #entries whose cost changed.""" import torch.distributed as dist if not (dist.is_available() and dist.is_initialized()): @@ -403,55 +438,46 @@ def warm_and_sync(self) -> int: group = _get_cost_sync_group() keys = sorted(self._table.keys(), key=repr) - - # Fail-fast key-set check (see docstring). all_gather_object is a single - # symmetric collective and every rank sees the same result, so all ranks - # take the same branch. key_reprs = [repr(k) for k in keys] all_key_reprs: list = [None] * world dist.all_gather_object(all_key_reprs, key_reprs, group=group) + + sets = [set(kr or []) for kr in all_key_reprs] + shared_reprs = sets[0].intersection(*sets[1:]) if sets else set() + # Only measure keys that every rank both owns AND has a stashed snode for: + # measuring a collective when one rank has no snode (keeps prior ns without + # issuing NCCL) desyncs the rest. + has_snode_reprs = {repr(k) for k in self._key_snode} + all_has: list = [None] * world + dist.all_gather_object(all_has, list(has_snode_reprs), group=group) + measurable_reprs = shared_reprs.intersection(*(set(h or []) for h in all_has)) + if any(kr != all_key_reprs[0] for kr in all_key_reprs[1:]): ref = set(all_key_reprs[0] or []) mine = set(key_reprs) - missing = sorted(ref - mine)[:3] - extra = sorted(mine - ref)[:3] magi_logger.warning( "warm_and_sync: cross-rank profiling key sets DIFFER (counts per rank: %s; " - "this rank vs rank0 -- missing %d e.g. %s, extra %d e.g. %s). The per-rank " - "graphs are not structurally identical, so rank-lockstep measurement would " - "deadlock. Falling back to the ANALYTICAL cost estimate for this graph " - "(rank-deterministic, less accurate). Consider fsdp_config.cost_mode=analytical.", + "this rank vs rank0 -- missing %d, extra %d). Lockstep-measuring the " + "intersection (%d shared keys with snodes on every rank); rank-local " + "keys stay local-measured (or analytical if they issue a collective).", [len(kr or []) for kr in all_key_reprs], len(ref - mine), - missing, len(mine - ref), - extra, + len(measurable_reprs), ) - n = 0 - for k, snode in self._key_snode.items(): - e = self._table.get(k) - if e is None: - continue - ns = _safe_analytical(snode) - if ns != e.ns: - n += 1 - e.ns = ns - e.measured = False - self._key_snode.clear() - return n + + # Rank-identical iteration order (repr sort) over the measurable intersection. + shared_keys = sorted((k for k in keys if repr(k) in measurable_reprs), key=repr) local_ns: dict = {} - for k in keys: + measured_here: set = set() + for k in shared_keys: snode = self._key_snode.get(k) dist.barrier(group=group) - if snode is not None: - local_ns[k] = self._measure_one(snode) - else: - local_ns[k] = self._table[k].ns # no cached snode -> keep prior measurement + # snode is non-None on every rank by measurable_reprs construction. + local_ns[k] = self._measure_one(snode) + measured_here.add(k) dist.barrier(group=group) - # Union of measured keys across ranks -> flag entries measured=True. - measured_here = set(self._key_snode.keys()) - gathered: list = [None] * world dist.all_gather_object(gathered, local_ns, group=group) gathered_measured: list = [None] * world @@ -464,14 +490,26 @@ def warm_and_sync(self) -> int: measured_keys = set() for mk in gathered_measured: measured_keys.update(mk or []) + n = 0 for k, e in self._table.items(): if k in measured_keys: - e.measured = True # reconciled from a real rank-lockstep measurement - m = merged.get(k) - if m is not None and m != e.ns: - e.ns = m - n += 1 + e.measured = True + m = merged.get(k) + if m is not None and m != e.ns: + e.ns = m + n += 1 + continue + # Rank-local key (not in the all-rank intersection): cannot lockstep + # measure. Keep the warmup measurement for pure compute; degrade + # collective / internal-collective ops to analytical (solo replay hangs). + snode = self._key_snode.get(k) + if snode is not None and _needs_lockstep_measure(snode): + ns = _safe_analytical(snode) + if ns != e.ns: + n += 1 + e.ns = ns + e.measured = False self._key_snode.clear() # drop snode refs (unpicklable) once sync is done return n @@ -542,6 +580,11 @@ def __call__(self, snode: BaseSchedulerNode) -> float: # Extern replay is ShapeEnv-isolated -> safe with free symbols; fused # Triton (benchmark_fused_nodes) would specialize the dynamic dim, so it # stays analytical while the graph is dynamic. + # TODO: measure fused Triton on dynamic graphs the way externs are + # replayed -- run the generated kernel on size-hint tensors inside a + # sandbox that cannot leak Eq(sym, hint) into the live ShapeEnv. + # benchmark_fused_nodes is unsafe here; eager-replay of the original + # aten.sin/add/... would time the unfused launches, not the fused kernel. if not is_extern and _graph_has_free_symbols(): return _safe_analytical(snode) @@ -555,7 +598,7 @@ def __call__(self, snode: BaseSchedulerNode) -> float: return entry.ns # Extern with an INTERNAL collective (CP attention / MoE): in sync mode, - # never measure it here -- the warm-up runs per-rank WITHOUT barriers, and + # never measure it here -- the warm-up runs per-rank without barriers, and # the adaptive benchmarker would issue rank-dependent numbers of the # internal NCCL op -> count mismatch -> hang. Seed analytical + stash the # snode; warm_and_sync re-measures it in rank-lockstep (fixed iters). @@ -615,6 +658,15 @@ def _measure_inner(self, snode: BaseSchedulerNode) -> float: return _safe_analytical(snode) +def _needs_lockstep_measure(snode: BaseSchedulerNode) -> bool: + """True if replaying this snode alone would issue a collective (hang without + peer ranks). Shared keys are measured under barriers; rank-local ones must + fall back to analytical instead.""" + if contains_collective(snode): + return True + return isinstance(snode, ExternKernelSchedulerNode) and _extern_has_internal_collective(snode) + + def _safe_analytical(snode: BaseSchedulerNode) -> float: try: return snode.get_estimated_runtime() diff --git a/tests/api_tests/test_register_custom_op.py b/tests/api_tests/test_register_custom_op.py index 071cee0..91722c9 100644 --- a/tests/api_tests/test_register_custom_op.py +++ b/tests/api_tests/test_register_custom_op.py @@ -68,6 +68,49 @@ def _multi_input_op(a: torch.Tensor, b: torch.Tensor, scale: float) -> torch.Ten assert_close(output, expected) + def test_materialize_inputs_same_signature_as_op(self): + """Decorator hook is called with the custom op's realized arguments.""" + from magi_compiler.profiling import apply_materialize_inputs, get_materialize_inputs_hook, op_has_internal_collective + from magi_compiler.profiling.materialize_inputs import _INTERNAL_COLLECTIVE_OPS, _MATERIALIZE_INPUT_HOOKS + + def _materialize(q, k, v, cp_split_sizes): + seq = int(q.shape[0]) + return q, k, v, [seq] * len(cp_split_sizes) + + @magi_register_custom_op( + name="test::materialize_same_sig_op", + infer_output_meta_fn=["q"], + materialize_inputs=_materialize, + has_internal_collective=True, + ) + def _op(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cp_split_sizes: list[int]) -> torch.Tensor: + return q + + try: + assert op_has_internal_collective("test::materialize_same_sig_op") + hook = get_materialize_inputs_hook("test::materialize_same_sig_op") + q = torch.zeros(8, 4) + args, kwargs = apply_materialize_inputs(hook, (q, q, q, [0, 0]), {}) + assert kwargs == {} + assert args[3] == [8, 8] + finally: + _MATERIALIZE_INPUT_HOOKS.pop("test::materialize_same_sig_op", None) + _INTERNAL_COLLECTIVE_OPS.discard("test::materialize_same_sig_op") + + def test_has_internal_collective_without_hook(self): + from magi_compiler.profiling import get_materialize_inputs_hook, op_has_internal_collective + from magi_compiler.profiling.materialize_inputs import _INTERNAL_COLLECTIVE_OPS + + @magi_register_custom_op(name="test::moe_flag_only_op", has_internal_collective=True) + def _moe(x: torch.Tensor) -> torch.Tensor: + return x + + try: + assert op_has_internal_collective("test::moe_flag_only_op") + assert get_materialize_inputs_hook("test::moe_flag_only_op") is None + finally: + _INTERNAL_COLLECTIVE_OPS.discard("test::moe_flag_only_op") + class TestInferOutputMeta: """Tests for custom op with infer_output_meta_fn.""" diff --git a/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py b/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py index 45e2a85..e124dae 100644 --- a/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py +++ b/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py @@ -24,10 +24,10 @@ 4. assert the estimate is within a tolerance band of the independent measurement. Also exercises ``ProfilingRuntimeEstimator.warm_and_sync`` (the profile_sync entry) to confirm it runs rank-lockstep without deadlock and reconciles a collective entry, and -the KEY-SET MISMATCH fail-fast: when one rank's table has an extra key (simulating a -per-rank structural divergence), warm_and_sync must NOT enter the barrier loop (which -would deadlock) -- it must detect the mismatch on every rank, warn, and degrade every -stashed entry to the analytical estimate (measured=False). +the KEY-SET INTERSECTION path: when one rank's table has an extra key (simulating a +per-rank structural divergence), warm_and_sync must still lockstep-measure the SHARED +collective key (not abandon the whole table) and degrade only the rank-local +collective-bearing keys to analytical -- completing at all proves no hang. Run: torchrun --nproc_per_node=2 .../estimator_collective_helper.py @@ -147,20 +147,23 @@ def main() -> None: coll_entries = [e for e in est.table.values() if e.kind == "collective"] warmsync_ok = len(coll_entries) >= 1 and all(e.measured for e in coll_entries) - # 5. key-set MISMATCH fail-fast: rank 1 injects an extra table entry so the - # cross-rank key sets differ. warm_and_sync must detect this on EVERY rank - # (symmetric all_gather_object), skip the per-key barrier loop entirely (which - # would deadlock on the count mismatch), and degrade this compile's entries to - # the analytical estimate (measured=False). Completing at all proves no hang. + # 5. key-set INTERSECTION: rank 1 injects an extra table entry so the + # cross-rank key sets differ. warm_and_sync must still lockstep-measure the + # shared collective key (measured=True), and degrade only the rank1-only + # collective-bearing fake key to analytical. Completing at all proves no hang. est2 = ProfilingRuntimeEstimator() est2._sync_across_ranks = True est2(coll_snode) # both ranks: seed the shared collective entry if rank == 1: fake_key = ("mismatch_only_on_rank1",) est2._table[fake_key] = ProfileEntry(ns=1.0, kind="extern", label="fake", measured=True) - est2._key_snode[fake_key] = coll_snode + est2._key_snode[fake_key] = coll_snode # coll-bearing -> analytical on local-only path est2.warm_and_sync() - mismatch_ok = all(not e.measured for e in est2.table.values()) and not est2._key_snode + shared_entries = [e for k, e in est2.table.items() if k != ("mismatch_only_on_rank1",)] + mismatch_ok = len(shared_entries) >= 1 and all(e.measured for e in shared_entries) and not est2._key_snode + if rank == 1: + fake = est2.table.get(("mismatch_only_on_rank1",)) + mismatch_ok = mismatch_ok and fake is not None and not fake.measured dist.barrier() # gather agreement across ranks diff --git a/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py b/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py index 3cabadb..7de0f86 100644 --- a/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py +++ b/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py @@ -19,24 +19,37 @@ The reorder pass is an Inductor ``reorder_for_compute_comm_overlap_passes`` callback; it needs a process group (its multi-rank-determinism warmup calls dist.get_rank()). We build a fn: y = (x @ w0).relu() # upstream compute + y = all_reduce(y) ; wait # a NON-weight collective in between g = all_gather(shard) ; wait # a weight gather to hoist - out = y + gathered_use # consumer after the compute -and wrap the pass so we can assert it RAN and returned a valid schedule. + out = y @ gathered_use # consumer after the compute +and wrap the pass so we can assert it RAN and returned a valid schedule. The +all_reduce stands in for a CP / EP kernel: the gather's compute window reaches past +it, so hoisting requires hopping another collective. With ``--mismatch`` (needs >=2 ranks): rank 1's fn gets an EXTRA compute op so the per-rank graphs are structurally DIFFERENT. The reorder pass's cross-rank -graph-fingerprint fail-fast must fire on EVERY rank (symmetric all_gather), warn, -and leave the schedule unchanged -- completing at all proves the check itself does -not desync. +graph-fingerprint check must fire on EVERY rank (symmetric all_gather), warn, and +continue in SLOT-consensus mode (the collective skeleton still matches) -- the +gather may hop the all_reduce as long as EVERY rank hops it. The invariant that +actually matters is checked directly: the collective sequence of the FINAL schedule +is all_gathered and compared across ranks. + +With ``--modes-only`` (>=2 ranks, gloo, no CUDA / no compile): drive +``_negotiate_mode`` directly with synthetic per-rank inputs and assert it returns +the expected mode for each rung of the ladder (identical / slot / pinned / abort). Run: torchrun --nproc_per_node=1 tests/feature_tests/fsdp_overlap_helper/reorder_helper.py torchrun --nproc_per_node=2 ... reorder_helper.py --mismatch + torchrun --nproc_per_node=2 ... reorder_helper.py --modes-only Markers (rank 0): REORDER_CALLED gathers= REORDER_OK moved= (pass returned; N launches repositioned) REORDER_FINITE ok= (compiled output finite + matches eager) - REORDER_MISMATCH unchanged= (--mismatch only: schedule left untouched) + REORDER_SKELETON ok= (final collective sequence identical on all ranks) + REORDER_MISMATCH local= (--mismatch only: divergent-graph path taken) + REORDER_SLOT rank= (--mismatch only: SLOT-consensus mode chosen) + REORDER_MODES ok= (--modes-only: the mode ladder returned as expected) REORDER_PASS / REORDER_FAIL """ @@ -53,10 +66,65 @@ from magi_compiler.passes.fsdp_overlap import reorder as _ro +class _FakeIR: + op_overload = "fake.op" + origins = None + + def get_size(self): + return [8, 8] + + +class _FakeSnode: + """Enough of a snode for ``_graph_fingerprint`` (the only thing the mode + negotiation reads out of the schedule).""" + + snodes = None + + def __init__(self) -> None: + self.node = _FakeIR() + + +def _mode_ladder_selfcheck(rank: int) -> bool: + """Assert every rung of ``_negotiate_mode``'s ladder, with rank 1 feeding the + divergent input. All ranks walk the cases in the same order, so the symmetric + all_gather inside each call stays lockstep.""" + negotiate = FsdpOverlapReorder._negotiate_mode + odd = rank == 1 + ag, other = (True, "ag", (8, 8)), (False, "cp", (4,)) + cases = { + # (n_snodes, weight-AG count, skeleton kinds) -> expected mode + "identical": (4, 2, [ag, other, ag]), + "slot": (5 if odd else 4, 2, [ag, other, ag]), # graphs differ, skeleton does not + "pinned": (5 if odd else 4, 2, [ag, other, ag] if odd else [ag, ag, other]), + "abort": (5 if odd else 4, 3 if odd else 2, [ag, other, ag]), + } + ok = True + for expected, (n_snodes, n_ag, kinds) in cases.items(): + got = negotiate([_FakeSnode() for _ in range(n_snodes)], [None] * n_ag, kinds)[0] + ok = ok and got == expected + print(f"REORDER_MODE_CASE rank={rank} expected={expected} got={got}", flush=True) + return ok + + def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--mismatch", action="store_true", help="rank1 compiles a structurally different graph") + ap.add_argument("--modes-only", action="store_true", help="only self-check the mode ladder (gloo, no compile)") args = ap.parse_args() + + if args.modes_only: + dist.init_process_group("gloo") + my_rank = dist.get_rank() + t = torch.tensor([1 if _mode_ladder_selfcheck(my_rank) else 0]) + dist.all_reduce(t, op=dist.ReduceOp.MIN) # every rank sees the same verdict + all_ok = bool(t.item()) + if my_rank == 0: + print(f"REORDER_MODES ok={all_ok}", flush=True) + print("REORDER_PASS" if all_ok else "REORDER_FAIL", flush=True) + dist.barrier() + dist.destroy_process_group() + raise SystemExit(0 if all_ok else 1) + dist.init_process_group("cpu:gloo,cuda:nccl") rank = dist.get_rank() world = dist.get_world_size() @@ -66,6 +134,7 @@ def main() -> None: torch.manual_seed(0) _AG = torch.ops._c10d_functional.all_gather_into_tensor.default + _AR = torch.ops._c10d_functional.all_reduce.default _WAIT = torch.ops._c10d_functional.wait_tensor.default H = 512 @@ -76,6 +145,7 @@ def main() -> None: def fn(x, w0, shard): y = (x @ w0).relu() # upstream compute the gather can hide behind + y = _WAIT(_AR(y, "sum", grp)) # non-weight collective the gather must hop if extra_op: y = y.sin() # rank1-only node -> graphs differ across ranks g = _WAIT(_AG(shard, world, grp)) # weight all-gather + wait @@ -83,18 +153,22 @@ def fn(x, w0, shard): return y @ gathered # instrument the pass: count how many times it runs, how many launches move, - # and whether the returned schedule is identical to the input (fail-fast path). - calls = {"n": 0, "gathers": 0, "moved": 0, "unchanged": True, "warned_mismatch": False} + # and whether the returned schedule is identical to the input (LOCAL path). + calls = {"n": 0, "gathers": 0, "moved": 0, "unchanged": True, "warned_mismatch": False, "slot_mode": False} + skeletons: list = [] # collective sequence of every schedule the pass returned orig_call = FsdpOverlapReorder.__call__ # magi_logger output from inside an Inductor compile does not reliably reach the - # subprocess streams; intercept the warning call itself to detect the fail-fast. + # subprocess streams; intercept the warning call itself to detect the mode taken. orig_warning = _ro.magi_logger.warning def spy_warning(msg, *a, **kw): if "NOT structurally identical" in str(msg): calls["warned_mismatch"] = True print(f"REORDER_WARNED rank={rank}", flush=True) + if "SLOT-consensus" in str(msg): + calls["slot_mode"] = True + print(f"REORDER_SLOT rank={rank}", flush=True) return orig_warning(msg, *a, **kw) _ro.magi_logger.warning = spy_warning @@ -105,6 +179,7 @@ def spy(self, snodes): before = list(snodes) out = orig_call(self, snodes) calls["unchanged"] = len(out) == len(before) and all(a is b for a, b in zip(out, before)) + skeletons.append(_ro._collective_skeleton(out)[1]) return out FsdpOverlapReorder.__call__ = spy @@ -134,11 +209,18 @@ def spy(self, snodes): rel = ((out.float() - eager.float()).norm() / (eager.float().norm() + 1e-6)).item() numeric_ok = finite and rel < 5e-2 - # In --mismatch mode the fail-fast must leave the schedule untouched on EVERY - # rank; agree across ranks before printing. - ok_local = calls["n"] > 0 and calls["gathers"] >= 1 and numeric_ok + # THE invariant: whatever each rank decided, the collective sequence of the + # schedules it emitted must be identical on every rank -- that (not identical + # absolute placement) is what keeps NCCL's positional matching intact. + peer_skeletons: list = [None] * world + dist.all_gather_object(peer_skeletons, skeletons) + skeleton_ok = all(s == peer_skeletons[0] for s in peer_skeletons[1:]) + + # In --mismatch mode the divergent-graph path must warn on EVERY rank; agree + # across ranks before printing. Schedule may change under SLOT reordering. + ok_local = calls["n"] > 0 and calls["gathers"] >= 1 and numeric_ok and skeleton_ok if args.mismatch: - ok_local = ok_local and calls["unchanged"] and calls["warned_mismatch"] + ok_local = ok_local and calls["warned_mismatch"] and calls["slot_mode"] t = torch.tensor([1 if ok_local else 0], device=dev) dist.all_reduce(t) all_ok = int(t.item()) == world @@ -147,8 +229,9 @@ def spy(self, snodes): print(f"REORDER_CALLED gathers={calls['gathers']}", flush=True) print(f"REORDER_OK ran={calls['n'] > 0}", flush=True) print(f"REORDER_FINITE ok={numeric_ok} rel={rel:.5f}", flush=True) + print(f"REORDER_SKELETON ok={skeleton_ok}", flush=True) if args.mismatch: - print(f"REORDER_MISMATCH unchanged={calls['unchanged']}", flush=True) + print(f"REORDER_MISMATCH local={calls['warned_mismatch']} unchanged={calls['unchanged']}", flush=True) print("REORDER_PASS" if all_ok else "REORDER_FAIL", flush=True) rc = 0 if all_ok else 1 else: diff --git a/tests/feature_tests/test_fsdp_overlap_reorder.py b/tests/feature_tests/test_fsdp_overlap_reorder.py index fa6c65a..39388d4 100644 --- a/tests/feature_tests/test_fsdp_overlap_reorder.py +++ b/tests/feature_tests/test_fsdp_overlap_reorder.py @@ -67,6 +67,7 @@ def test_reorder_multi_rank(): p = _run(2) out = p.stdout + p.stderr assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" + assert "REORDER_SKELETON ok=True" in p.stdout, out[-3000:] assert "REORDER_PASS" in p.stdout, out[-3000:] @@ -144,16 +145,33 @@ def syms(*names): assert _graph_fingerprint(rank0) != _graph_fingerprint(doubled) +@requires_torchrun +def test_reorder_mode_ladder(): + """The cross-rank safety ladder, driven directly on synthetic inputs (gloo, no + CUDA): identical graphs, graphs differing only in compute (-> slot consensus), + differing collective skeletons (-> pinned) and differing weight-AG counts + (-> abort).""" + p = _run(2, "--modes-only", port="29633") + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" + assert "REORDER_MODES ok=True" in p.stdout, out[-3000:] + + @requires_cuda @requires_torchrun @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") -def test_reorder_graph_mismatch_fail_fast(): +def test_reorder_graph_mismatch_slot_mode(): """world=2 with rank1 compiling a structurally DIFFERENT graph: the cross-rank - graph-fingerprint check must fire on both ranks (warning), leave the schedule - unchanged, and complete without deadlock.""" + graph-fingerprint check must fire on both ranks (warning) and continue in + SLOT-consensus mode -- the graphs differ only in compute, so the collective + skeleton still matches and the gather may hop the graph's other collective, as + long as every rank hops it. The emitted collective sequence must come out + identical on both ranks (REORDER_SKELETON), which is what rules out deadlock.""" p = _run(2, "--mismatch", port="29632") out = p.stdout + p.stderr assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" - assert "REORDER_MISMATCH unchanged=True" in p.stdout, out[-3000:] - assert "REORDER_WARNED" in p.stdout, out[-3000:] # the fail-fast warning fired + assert "REORDER_MISMATCH local=True" in p.stdout, out[-3000:] + assert "REORDER_WARNED" in p.stdout, out[-3000:] # the divergent-graph warning fired + assert "REORDER_SLOT" in p.stdout, out[-3000:] # ... and it chose SLOT consensus + assert "REORDER_SKELETON ok=True" in p.stdout, out[-3000:] assert "REORDER_PASS" in p.stdout, out[-3000:] diff --git a/tests/feature_tests/test_profiling_estimator.py b/tests/feature_tests/test_profiling_estimator.py index e30697d..79317d5 100644 --- a/tests/feature_tests/test_profiling_estimator.py +++ b/tests/feature_tests/test_profiling_estimator.py @@ -54,6 +54,22 @@ def test_realize_tensor_node_builds_matching_tensor(): assert out.device.type == "cuda" +@requires_cuda +def test_realize_float8_tensor_does_not_raise(): + """float8 has no ``randn`` kernel; replay must cast from float32 noise so fp8 + custom-op / SmoothQuant GEMM measurement does not silently fall back to 0.""" + if not hasattr(torch, "float8_e4m3fn"): + pytest.skip("float8_e4m3fn not available") + dev = torch.cuda.current_device() + node = _node_with_val(torch.empty(4, 8, device=dev, dtype=torch.float8_e4m3fn)) + out = _realize_arg(node) + assert isinstance(out, torch.Tensor) + assert out.shape == (4, 8) + assert out.dtype == torch.float8_e4m3fn + assert out.device.type == "cuda" + assert torch.isfinite(out.float()).all() + + def test_realize_scalar_int_node(): node = _node_with_val(16) # a Node carrying a plain int scalar assert _realize_arg(node) == 16 @@ -248,9 +264,9 @@ def test_internal_collective_extern_not_measured_in_sync_warmup(monkeypatch, syn adaptive benchmarker would issue rank-dependent numbers of the internal NCCL op -> hang). It is seeded analytical + stashed for warm_and_sync. In non-sync mode the normal measurement path still runs.""" - from magi_compiler.profiling import register_benchmark_inputs + from magi_compiler.profiling import register_materialize_inputs from magi_compiler.profiling import runtime_estimator as re_mod - from magi_compiler.profiling.benchmark_inputs import _BENCHMARK_INPUT_HOOKS, _INTERNAL_COLLECTIVE_OPS + from magi_compiler.profiling.materialize_inputs import _INTERNAL_COLLECTIVE_OPS, _MATERIALIZE_INPUT_HOOKS measured = {"called": False} @@ -258,7 +274,7 @@ def _boom(*a, **k): measured["called"] = True raise AssertionError("must not be measured in sync warm-up") - register_benchmark_inputs("aten::mm", lambda fx_node, realize: None, has_internal_collective=True) + register_materialize_inputs("aten::mm", has_internal_collective=True) monkeypatch.setattr(re_mod, "_measure_extern", _boom) try: est = ProfilingRuntimeEstimator() @@ -273,7 +289,7 @@ def _boom(*a, **k): else: assert measured["called"] is True # non-sync: measurement path still taken finally: - _BENCHMARK_INPUT_HOOKS.pop("aten::mm", None) + _MATERIALIZE_INPUT_HOOKS.pop("aten::mm", None) _INTERNAL_COLLECTIVE_OPS.discard("aten::mm") @@ -490,6 +506,6 @@ def test_collective_profile_accuracy_multi_rank(): assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" assert "COLL_ACCURATE ok=True" in p.stdout, out[-3000:] assert "COLL_WARMSYNC" in p.stdout and "ok=True" in p.stdout, out[-3000:] - # key-set mismatch fail-fast: no deadlock + degraded to analytical on all ranks + # key-set intersection: shared keys still measured; no hang on mismatch assert "COLL_MISMATCH ok=True" in p.stdout, out[-3000:] assert "COLL_PASS" in p.stdout, out[-3000:] diff --git a/tests/feature_tests/test_profiling_registry.py b/tests/feature_tests/test_profiling_registry.py index 42587b2..c88a1f0 100644 --- a/tests/feature_tests/test_profiling_registry.py +++ b/tests/feature_tests/test_profiling_registry.py @@ -12,71 +12,86 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the profiling benchmark-input registry -(``magi_compiler.profiling.benchmark_inputs``). +"""Unit tests for the profiling materialize-input registry +(``magi_compiler.profiling.materialize_inputs``). Pure-CPU: the registry is a plain module-level dict/set, no torch/GPU/distributed. """ import pytest -from magi_compiler.profiling import benchmark_inputs as bi -from magi_compiler.profiling import get_benchmark_inputs_hook, op_has_internal_collective, register_benchmark_inputs +from magi_compiler.profiling import apply_materialize_inputs, get_materialize_inputs_hook +from magi_compiler.profiling import materialize_inputs as mi +from magi_compiler.profiling import op_has_internal_collective, register_materialize_inputs @pytest.fixture def clean_registry(): """Snapshot + restore the global registry so tests don't leak into each other.""" - hooks = dict(bi._BENCHMARK_INPUT_HOOKS) - coll = set(bi._INTERNAL_COLLECTIVE_OPS) + hooks = dict(mi._MATERIALIZE_INPUT_HOOKS) + coll = set(mi._INTERNAL_COLLECTIVE_OPS) yield - bi._BENCHMARK_INPUT_HOOKS.clear() - bi._BENCHMARK_INPUT_HOOKS.update(hooks) - bi._INTERNAL_COLLECTIVE_OPS.clear() - bi._INTERNAL_COLLECTIVE_OPS.update(coll) + mi._MATERIALIZE_INPUT_HOOKS.clear() + mi._MATERIALIZE_INPUT_HOOKS.update(hooks) + mi._INTERNAL_COLLECTIVE_OPS.clear() + mi._INTERNAL_COLLECTIVE_OPS.update(coll) def test_unregistered_returns_none(clean_registry): - assert get_benchmark_inputs_hook("ns::never_registered") is None + assert get_materialize_inputs_hook("ns::never_registered") is None assert op_has_internal_collective("ns::never_registered") is False def test_register_and_get(clean_registry): - def hook(fx_node, realize): + def hook(q, k, v, cp_split_sizes): return None - register_benchmark_inputs("ns::op_a", hook) - assert get_benchmark_inputs_hook("ns::op_a") is hook + register_materialize_inputs("ns::op_a", hook) + assert get_materialize_inputs_hook("ns::op_a") is hook # not flagged as internal-collective by default assert op_has_internal_collective("ns::op_a") is False def test_internal_collective_flag(clean_registry): - register_benchmark_inputs("ns::coll_op", lambda n, r: None, has_internal_collective=True) + register_materialize_inputs("ns::coll_op", has_internal_collective=True) assert op_has_internal_collective("ns::coll_op") is True - assert get_benchmark_inputs_hook("ns::coll_op") is not None + assert get_materialize_inputs_hook("ns::coll_op") is None # a hook registered WITHOUT the flag must not be marked - register_benchmark_inputs("ns::plain_op", lambda n, r: None) + register_materialize_inputs("ns::plain_op", lambda *a, **k: None) assert op_has_internal_collective("ns::plain_op") is False def test_register_overrides_previous(clean_registry): - def hook1(n, r): - return "1" + def hook1(q): + return (q,) - def hook2(n, r): - return "2" + def hook2(q): + return (q,) - register_benchmark_inputs("ns::op_b", hook1) - assert get_benchmark_inputs_hook("ns::op_b") is hook1 - register_benchmark_inputs("ns::op_b", hook2) - assert get_benchmark_inputs_hook("ns::op_b") is hook2 + register_materialize_inputs("ns::op_b", hook1) + assert get_materialize_inputs_hook("ns::op_b") is hook1 + register_materialize_inputs("ns::op_b", hook2) + assert get_materialize_inputs_hook("ns::op_b") is hook2 def test_hook_is_callable_returning_none(clean_registry): - """A no-op hook (returns None -> fall back to generic realize) is valid.""" - register_benchmark_inputs("ns::noop", lambda fx_node, realize: None, has_internal_collective=True) - hook = get_benchmark_inputs_hook("ns::noop") - assert hook(object(), lambda x: x) is None + """A no-op hook (returns None -> keep generic realize) is valid.""" + register_materialize_inputs("ns::noop", lambda *a, **k: None, has_internal_collective=True) + hook = get_materialize_inputs_hook("ns::noop") + assert hook(object()) is None assert op_has_internal_collective("ns::noop") is True + + +def test_apply_materialize_inputs_same_signature_as_op(clean_registry): + """Hook sees realized custom-op args and returns a rewritten positional tuple.""" + + def hook(q, k, v, cp_split_sizes): + seq = q[0] + return q, k, v, [seq] * len(cp_split_sizes) + + args, kwargs = apply_materialize_inputs(hook, ((16,), (16,), (16,), [0, 0]), {}) + assert kwargs == {} + assert args[3] == [16, 16] + assert apply_materialize_inputs(None, (1,), {}) == ((1,), {}) + assert apply_materialize_inputs(lambda *a, **k: None, (1, 2), {"x": 3}) == ((1, 2), {"x": 3})