diff --git a/magi_compiler/magi_backend/compile_artifacts.py b/magi_compiler/magi_backend/compile_artifacts.py index 0af9438..59309c4 100644 --- a/magi_compiler/magi_backend/compile_artifacts.py +++ b/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( diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 8ef38bd..239b224 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/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, diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index 2e78e44..7309264 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/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) diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index 9b5bb4e..3ba6ca7 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/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) diff --git a/tests/feature_tests/test_compile_artifacts.py b/tests/feature_tests/test_compile_artifacts.py index 2a072fe..967c565 100644 --- a/tests/feature_tests/test_compile_artifacts.py +++ b/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 diff --git a/tests/feature_tests/test_fsdp_overlap_reorder.py b/tests/feature_tests/test_fsdp_overlap_reorder.py index 5e92142..fa6c65a 100644 --- a/tests/feature_tests/test_fsdp_overlap_reorder.py +++ b/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") diff --git a/tests/feature_tests/test_profiling_estimator.py b/tests/feature_tests/test_profiling_estimator.py index 3f2ae5b..3ec5edd 100644 --- a/tests/feature_tests/test_profiling_estimator.py +++ b/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: ("~", ) 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 # ---------------------------------------------------------------------------