[Fix] fsdp_overlap symbol-name rank divergence + drop pinned example_inputs in MagiSerializableFunction

magi_compiler/magi_backend/compile_artifacts.py

@@ -24,7 +24,6 @@
 from unittest.mock import patch
 
 import torch
-from torch.utils._pytree import tree_map_only
 
 from magi_compiler.utils import magi_logger
 
@@ -48,13 +47,18 @@ class MagiSerializableFunction(SerializableCallable):
     disk. There's no need to wrap around the compiled function if we don't want
     to serialize them in particular cases.
     Right now serialization for the custom backend is done via
-    serializing the Dynamo fx graph plus example inputs.
+    serializing the Dynamo fx graph; compile inputs are re-derived from the
+    graph placeholders' ``example_value`` metadata on rebuild.
+
+    Deliberately does NOT hold the example inputs dynamo hands the backend:
+    those are the REAL first-call input tensors, and this object lives in the
+    dynamo code cache for the process lifetime, so storing them would pin the
+    whole first-call input set (multi-GB of activations for a large region).
     """
 
     def __init__(
         self,
         graph_module,
-        example_inputs,
         model_tag,
         optimized_call,
         model_idx: int = 0,
@@ -63,7 +67,6 @@ def __init__(
     ):
         assert isinstance(graph_module, torch.fx.GraphModule)
         self.graph_module = graph_module
-        self.example_inputs = example_inputs
         self.model_idx = model_idx
         self.model_tag = model_tag
         self.traced_files = traced_files or []
@@ -92,14 +95,12 @@ def serialize_compile_artifacts(cls, compiled_fn: MagiSerializableFunction) -> b
         patched_op_pickle = GraphNodeOpPatchUtils.make_patch_for_pickle()
 
         # Pickle under all patches
-        state["example_inputs"] = tree_map_only(torch.Tensor, lambda _: None, state["example_inputs"])
         with (
             patch.object(GraphPickler, "reducer_override", patched_reducer),
             patch.object(_NodePickleData, "__init__", patched_node_init),
             patch.object(_OpPickleData, "pickle", patched_op_pickle),
         ):
             state["graph_module"] = GraphPickler.dumps(state["graph_module"], Options(ops_filter=None))
-            state["example_inputs"] = GraphPickler.dumps(state["example_inputs"])
 
         return pickle.dumps(state)
 
@@ -125,16 +126,16 @@ def deserialize_compile_artifacts(cls, data: bytes) -> MagiSerializableFunction:
 
         state = pickle.loads(data)
 
-        # Backward compat: pop triton_kernel_info from old serialized artifacts.
+        # Backward compat: pop keys that old serialized artifacts carried.
         state.pop("triton_kernel_info", None)
+        state.pop("example_inputs", None)  # discarded unread; rebuild uses placeholder metadata
 
         fake_mode = FakeTensorMode(shape_env=ShapeEnv())
 
-        # Unpickle graph & inputs under node-level patches
+        # Unpickle the graph under node-level patches
         patched_unpickle = GraphNodePicklePatchUtils.make_patch_for_unpickle()
         with patch.object(_NodePickleData, "unpickle", patched_unpickle):
             state["graph_module"] = GraphPickler.loads(state["graph_module"], fake_mode)
-            state["example_inputs"] = GraphPickler.loads(state["example_inputs"], fake_mode)
 
         # Reconstruct CompileConfig from the serialized artifact (self-contained).
         compile_config_data = state.get("compile_config")
@@ -164,11 +165,17 @@ def rebuild_backend(self) -> None:
         from magi_compiler.magi_backend import MagiBackend
         from magi_compiler.utils import OrderedSet
 
-        # Fill None placeholders in example_inputs with FakeTensors from graph metadata.
-        placeholder_fake_values = [
-            node.meta.get("example_value") for node in self.graph_module.graph.nodes if node.op == "placeholder"
-        ]
-        compile_inputs = [inp if inp is not None else placeholder_fake_values[i] for i, inp in enumerate(self.example_inputs)]
+        # Compile inputs come entirely from the graph placeholders' metadata:
+        # FakeTensors for tensor inputs, SymInts/scalars for the rest. Missing
+        # metadata would silently miscompile downstream, so fail loudly instead.
+        placeholders = [node for node in self.graph_module.graph.nodes if node.op == "placeholder"]
+        compile_inputs = [node.meta.get("example_value") for node in placeholders]
+        missing = [node.name for node, val in zip(placeholders, compile_inputs) if val is None]
+        if missing:
+            raise RuntimeError(
+                f"AOT rebuild: {len(missing)} graph placeholder(s) lack example_value metadata "
+                f"(e.g. {missing[:5]}); cannot reconstruct compile inputs for model_tag={self.model_tag}"
+            )
 
         fake_mode = detect_fake_mode(compile_inputs)
         magi_backend = MagiBackend(

magi_compiler/magi_backend/magi_backend.py

@@ -711,7 +711,6 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> MagiSerializableFun
 
         return MagiSerializableFunction(
             graph,
-            example_inputs,
             self.model_tag,
             runnable_gm,
             model_idx=self.model_idx,

magi_compiler/passes/fsdp_overlap/reorder.py

@@ -80,13 +80,46 @@ def _is_multi_output(snode: BaseSchedulerNode) -> bool:
     return type(node) is MultiOutput
 
 
+def _size_hint_of(sym) -> int:
+    """Rank-identical size hint for a sympy symbol (0 if unavailable, e.g. in
+    unit tests without a live Inductor graph)."""
+    try:
+        from torch._inductor.virtualized import V
+
+        return int(V.graph.sizevars.size_hint(sym, fallback=0))
+    except Exception:  # noqa: BLE001
+        return 0
+
+
 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
     kernel is one ComputedBuffer whose class/size hide its contents (relu vs
     relu+sin look identical without them).  Targets only, not node names: names
-    carry per-rank numbering noise."""
+    carry per-rank numbering noise.
+
+    """
+    import sympy
+
     h = hashlib.sha256()
+    sym_canon: dict = {}  # sympy.Symbol -> canonical sympy.Symbol
+
+    def _canon_size(size) -> str:
+        dims = []
+        for d in size:
+            free = getattr(d, "free_symbols", None)
+            if not free:
+                dims.append(repr(d))
+                continue
+            fresh = [sym for sym in free if sym not in sym_canon]
+            # Name-free assignment order; symbol name only as the last-resort
+            # tie-break (see docstring: that case fails safe).
+            fresh.sort(key=lambda sym: (_size_hint_of(sym), d.count(sym), sym.name))
+            for sym in fresh:
+                sym_canon[sym] = sympy.Symbol(f"c{len(sym_canon):04d}")
+            dims.append(repr(d.xreplace(sym_canon)))
+        return "[" + ", ".join(dims) + "]"
+
     for s in order:
         h.update(type(s).__name__.encode())
         for sub in getattr(s, "snodes", None) or (s,):
@@ -96,7 +129,7 @@ def _graph_fingerprint(order: list[BaseSchedulerNode]) -> str:
             op = getattr(n, "op_overload", None) or getattr(n, "python_kernel_name", None) or type(n).__name__
             h.update(str(op).encode())
             try:
-                h.update(repr(n.get_size()).encode())
+                h.update(_canon_size(n.get_size()).encode())
             except Exception:  # noqa: BLE001
                 pass
             origins = getattr(n, "origins", None)

magi_compiler/profiling/runtime_estimator.py

@@ -135,11 +135,18 @@ def _is_symbolic(s) -> bool:
 def _static(shape) -> tuple:
     """Cache-key shape.  MUST NOT call ``int()`` on a SymInt -- that adds an
     ``Eq(sym, value)`` guard and specializes the dynamic dim, breaking dynamic
-    shape compilation.  Symbolic dims are stringified (stable within a compile)."""
+    shape compilation.
+
+    Symbolic dims are keyed by their size hint (guard-free, see
+    ``_concrete_size``), tagged "~".  Not by ``str(s)``: dynamo names shape
+    symbols differently per rank, so symbol-name keys broke
+    ``warm_and_sync``'s cross-rank check and silently fell back to the
+    inaccurate analytical estimate.  Hints are rank-identical and stable
+    within a compile, so isomorphic kernels share one measurement."""
     out = []
     for s in shape:
         if _is_symbolic(s):
-            out.append(str(s))
+            out.append(("~", _concrete_size(s)))
         else:
             out.append(int(s))
     return tuple(out)

tests/feature_tests/test_compile_artifacts.py

@@ -748,7 +748,7 @@ def test_basic_roundtrip(self):
         nodes[1].meta["example_value"] = ft_mul
         nodes[2].meta["example_value"] = ft_add
 
-        fn = MagiSerializableFunction(gm, [ft], "test_basic", lambda *a: None)
+        fn = MagiSerializableFunction(gm, "test_basic", lambda *a: None)
         data = MagiSerializableFunction.serialize_compile_artifacts(fn)
         assert isinstance(data, bytes) and len(data) > 0
 
@@ -800,7 +800,7 @@ def test_einops_roundtrip(self):
         nodes[0].meta["example_value"] = ft
         nodes[1].meta["example_value"] = ft_out
 
-        fn = MagiSerializableFunction(gm, [ft], "test_einops", lambda *a: None)
+        fn = MagiSerializableFunction(gm, "test_einops", lambda *a: None)
         data = MagiSerializableFunction.serialize_compile_artifacts(fn)
         assert isinstance(data, bytes) and len(data) > 0
 
@@ -846,7 +846,7 @@ def test_triton_roundtrip(self):
 
         list(gm.graph.nodes)[0].meta["example_value"] = ft
 
-        fn = MagiSerializableFunction(gm, [ft], "test_triton", lambda *a: None)
+        fn = MagiSerializableFunction(gm, "test_triton", lambda *a: None)
         data = MagiSerializableFunction.serialize_compile_artifacts(fn)
         assert isinstance(data, bytes) and len(data) > 0
 
@@ -890,7 +890,7 @@ def test_slice_roundtrip(self):
         nodes[1].meta["example_value"] = 5
         nodes[2].meta["example_value"] = ft_sliced
 
-        fn = MagiSerializableFunction(gm, [ft, None], "test_slice", lambda *a: None)
+        fn = MagiSerializableFunction(gm, "test_slice", lambda *a: None)
         data = MagiSerializableFunction.serialize_compile_artifacts(fn)
         assert isinstance(data, bytes) and len(data) > 0
 

tests/feature_tests/test_fsdp_overlap_reorder.py

@@ -70,6 +70,80 @@ def test_reorder_multi_rank():
     assert "REORDER_PASS" in p.stdout, out[-3000:]
 
 
+class _FakeIR:
+    """Sizes are REAL sympy expressions, as ``node.get_size()`` returns: the
+    canonicalization must survive sympy's StrPrinter, which orders commutative
+    terms by symbol NAME -- a hand-written repr string would bypass exactly the
+    layer the bug lives in."""
+
+    def __init__(self, size):
+        self.op_overload = "fake.op"
+        self._size = size
+        self.origins = None
+
+    def get_size(self):
+        return self._size
+
+
+class _FakeSnode:
+    snodes = None
+
+    def __init__(self, size):
+        self.node = _FakeIR(size)
+
+
+def _graph(*sizes):
+    """Build a fake snode list; each size is a list of sympy exprs / ints."""
+    return [_FakeSnode(size) for size in sizes]
+
+
+def test_fingerprint_canonicalizes_shape_symbols():
+    """Dynamic-shape symbol NAMES are per-rank numbering noise (the same
+    logical dim is s82 on rank 0 and s74 on rank 1 for the same graph): the
+    fingerprint must be invariant under symbol renaming, while still
+    distinguishing genuinely different symbolic structure.
+
+    Modeled like the real wan graph: the local-seq symbol first appears alone
+    in a placeholder-like dim, then the CP all_to_all output sums it with the
+    fresh cross-rank symbol, then downstream nodes use the fresh symbol alone.
+    """
+    import sympy
+
+    from magi_compiler.passes.fsdp_overlap.reorder import _graph_fingerprint
+
+    def syms(*names):
+        return [sympy.Symbol(n, positive=True, integer=True) for n in names]
+
+    # Same digit count (the case the string-level rename happened to handle).
+    (a0, b0), (a1, b1) = syms("s27", "s82"), syms("s27", "s74")
+    rank0 = _graph([a0, 3072], [a0 + b0, 6, 64], [b0, 64])
+    rank1 = _graph([a1, 3072], [a1 + b1, 6, 64], [b1, 64])
+    assert _graph_fingerprint(rank0) == _graph_fingerprint(rank1)
+
+    # Digit-count crossing: sympy prints ``s27 + s174`` as ``s174 + s27``
+    # (StrPrinter sorts terms by name), so any rename applied AFTER printing
+    # sees a different first-appearance order and diverges.
+    (a2, b2) = syms("s27", "s174")
+    rank2 = _graph([a2, 3072], [a2 + b2, 6, 64], [b2, 64])
+    assert _graph_fingerprint(rank0) == _graph_fingerprint(rank2)
+
+    # Same digit count but flipped relative order: fresh symbol sorts BEFORE
+    # the shared one on one rank (s34 < s50) and AFTER it on the other.
+    (a3, b3), (a4, b4) = syms("s50", "s82"), syms("s50", "s34")
+    rank3 = _graph([a3, 3072], [a3 + b3, 6, 64], [b3, 64])
+    rank4 = _graph([a4, 3072], [a4 + b4, 6, 64], [b4, 64])
+    assert _graph_fingerprint(rank3) == _graph_fingerprint(rank4)
+
+    # Genuinely different LINKAGE must still differ: downstream node reuses the
+    # local-seq symbol instead of the cross-rank one.
+    linked_other = _graph([a0, 3072], [a0 + b0, 6, 64], [a0, 64])
+    assert _graph_fingerprint(rank0) != _graph_fingerprint(linked_other)
+
+    # Genuinely different EXPRESSION structure must differ: 2*s vs s + s'.
+    doubled = _graph([a0, 3072], [2 * a0, 6, 64], [b0, 64])
+    assert _graph_fingerprint(rank0) != _graph_fingerprint(doubled)
+
+
 @requires_cuda
 @requires_torchrun
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs")

tests/feature_tests/test_profiling_estimator.py

@@ -83,16 +83,66 @@ def test_static_concrete_dims_are_ints():
     assert _static((4, 8, 16)) == (4, 8, 16)
 
 
-def test_static_symbolic_dim_stringified():
-    class _FakeSym:
-        # mimic a SymInt: _is_symbolic() returns True for objects with a `.node`
-        node = object()
-
-        def __str__(self):
-            return "s7"
-
-    out = _static((_FakeSym(), 8))
-    assert out == ("s7", 8)  # symbolic -> str, static -> int (no specializing int())
+class _FakeSym:
+    """Mimics a torch.SymInt: ``_is_symbolic()`` keys on having a ``.node``."""
+
+    node = object()
+
+    def __init__(self, name="s7", hint=None):
+        self._name = name
+        self.hint = hint
+
+    def __str__(self):
+        return self._name
+
+
+def test_static_symbolic_dim_keyed_by_hint():
+    """The PRODUCTION path (live V.graph): symbolic dims key by their
+    rank-identical size hint, NOT the symbol name -- dynamo numbers symbols
+    per rank (s82 vs s74 for the same dim), and symbol-name keys broke
+    warm_and_sync's cross-rank key-set match on any dynamic graph ->
+    analytical-cost fallback -> exposed all-gathers.
+
+    Two symbols sharing a hint sharing one table entry is exact, not
+    approximate: ``_measure_extern`` realizes replay inputs at these same
+    hints (``_realize_arg`` -> ``_concrete_size``), so the measured ns is a
+    pure function of (op, dtype, hint shape) -- re-measuring under a separate
+    key would produce the same value."""
+    from torch._inductor.virtualized import V
+
+    class _FakeSizevars:
+        def size_hint(self, sym, fallback=0):
+            return sym.hint
+
+    class _FakeGraph:
+        sizevars = _FakeSizevars()
+
+    with V.set_graph_handler(_FakeGraph()):
+        seq, tok = _FakeSym("s27", hint=4096), _FakeSym("s82", hint=2048)
+        seq_other_rank = _FakeSym("s74", hint=4096)  # same dim, renamed by rank 1
+
+        # hint reaches the key (not the fallback), tagged "~" for dynamic
+        assert _static((seq, 3072)) == (("~", 4096), 3072)
+        # different hints -> different keys: no false sharing across dims
+        assert _static((seq, 3072)) != _static((tok, 3072))
+        # same hint, different symbol NAME -> same key (the fix's purpose)
+        assert _static((seq, 3072)) == _static((seq_other_rank, 3072))
+        # a dynamic dim never collides with a static dim of the same value
+        assert _static((seq,)) != _static((4096,))
+
+
+def test_static_symbolic_dim_fallback_without_graph():
+    """DEGRADED path only (no live V.graph, e.g. bare unit tests):
+    ``_concrete_size`` falls back to 1 for every symbolic dim.  In production
+    ``_structural_key`` always runs inside Inductor scheduling where V.graph
+    is set, so this collapse never happens there -- and if it somehow did, the
+    measurement itself realizes inputs through the same fallback, so keys and
+    values degrade together."""
+    out = _static((_FakeSym("s7"), 8))
+    # the key SHAPE is what matters: ("~", <hint>) marks the dim dynamic,
+    # static dims stay plain ints (never int()'d from a SymInt -> no guard)
+    assert out == (("~", 1), 8)
+    assert _static((_FakeSym("s99"), 8)) == out
 
 
 # ---------------------------------------------------------------------------