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

[wtr0504]

🗂️ PR Category

• ✨ New Feature
• 🚀 Optimization (performance, memory, etc.)
• 💥 Breaking Change
• 🐛 Bug Fix
• 🛠️ Development / Refactoring
• 📚 Documentation
• 🧹 Chore (Dependencies, CI/CD, Configuration, etc.)
• 🧪 Testing

📝 Description

This PR contains two independent fixes in the compile/runtime layer, both found while profiling wan2.2 SimpleFSDP inference.

---

Fix 1 — fsdp_overlap: dynamic-shape symbol names are per-rank noise; key by structure, not name

Root cause. Dynamo allocates shape symbols from a process-local ShapeEnv counter with no cross-process coordination: the name a dynamic dim gets (s27, s82, …) depends on the order and history of symbol allocation during that rank's trace — not on the graph's structure. As long as every symbol in the graph derives from the ranks' common inputs, the names happen to match across ranks and the problem stays invisible.

A traced CP all_to_all breaks exactly this condition. Its output size is sum(output_split_sizes) — the local chunk plus the other ranks' chunk sizes. That is the one quantity in the graph not derived from local inputs, so dynamo materializes it as an additional symbol mid-trace, and where that symbol lands in each rank's private allocation history differs per rank. The all_to_all output prints s27 + s82 on rank 0 vs s27 + s74 on rank 1 for structurally identical graphs; a node-by-node fingerprint diff showed only the all_to_all and its downstream attention nodes differing.

Two cross-rank consistency checks treated these symbol names as part of the graph structure. Both are fail-safe, so instead of crashing they silently degraded, and the weight all-gathers ran fully exposed.

Changes.

1. reorder.py::_graph_fingerprint — the pre-reorder "per-rank graphs are structurally identical" digest hashed raw size reprs; isomorphic graphs fingerprinted differently → fail-fast fired → reorder disabled entirely.
   Now canonicalizes shape symbols at the sympy-expression level, before printing. Renaming the printed string is not sound: sympy's StrPrinter orders commutative terms by default_sort_key ≈ symbol name (s27 + s82 prints as-is, but s27 + s174 prints as s174 + s27), so "first appearance in the string" is itself a function of the per-rank names being erased — digit-count crossings and same-digit relative-order flips would still diverge. Instead, a symbol → c{i:04d} map is built in a name-free order (first appearance across the node walk; within one expression, fresh symbols ordered by size hint, then occurrence count; name only as a last-resort tie-break for the fully-symmetric case, which is unobserved in practice and fails safe — overlap off, never a divergent schedule), then xreplace is applied before repr: sympy still name-sorts commutative terms, but now by canonical names that are rank-identical, and zero-padding keeps their own sort order immune to digit-count effects. What is structural is the symbols' occurrence pattern, not their names; genuinely different symbolic structure (different linkage, 2*s vs s + s') still diverges. The fingerprint test goes through real sympy printing (the previous version fed hand-written strings, bypassing the layer the bug lives in) and covers digit-crossing and order-flip cases.

2. runtime_estimator.py::_static — the profiling table's structural key stringified symbolic dims, so warm_and_sync's cross-rank key-set check failed → the whole cost table silently fell back to the analytical estimate, which overestimates matmul ~~400x here (mm: 19µs real vs 7405µs analytical) → the reorder believed one adjacent mm hides any gather and moved every launch by 2 slots, verdict "hidden", actually exposed. Now keys symbolic dims by their size hint (("~", hint)) — guard-free (no SymInt specialization), rank-identical, and isomorphic kernels keep sharing one measurement. If ranks ever see genuinely different hints, the key-set check degrades safely to analytical as before.

---

Fix 2 — Remove example_inputs from MagiSerializableFunction (multi-GB memory pin)

Root cause. MagiSerializableFunction stored the example_inputs dynamo hands the backend — which are the real first-call input tensors, not fakes. The object lives in the dynamo code cache for the process lifetime, so the entire first-call input set of every compiled region stayed pinned in GPU memory.

The tensor values were never actually consumed anywhere:

• the live call path only invokes optimized_call;
• serialize_compile_artifacts already nulled every tensor leaf before pickling;
• rebuild_backend (deserialize path) only ever saw those nulls and re-filled them from the graph placeholders' example_value metadata.

Changes.

• Drop the example_inputs constructor parameter and attribute entirely (call site in magi_backend.py updated).
• rebuild_backend now derives compile inputs purely from placeholder example_value metadata (FakeTensors for tensor inputs, SymInts/scalars for the rest), and raises a clear RuntimeError naming the offending placeholders if any lacks metadata — missing metadata would otherwise silently miscompile.
• deserialize_compile_artifacts stays backward compatible: the legacy example_inputs key in old artifacts is popped unread (same treatment as triton_kernel_info).

Verification. test_compile_artifacts.py 38/38 (constructors updated to the new signature); AOT cache end-to-end suites (test_transformer_cache_reuse, test_restart_analysis_cache, test_autograd_function_cache_flag) 4/4 — these exercise the real serialize → deserialize → rebuild_backend path across processes; weakref check confirms first-call inputs are released after compile (pinned without the fix).

---

Tests

• New: test_fingerprint_canonicalizes_shape_symbols (α-equivalence of the fingerprint).
• Updated: test_static_symbolic_dim_keyed_by_hint (hint-based keys); test_compile_artifacts.py constructors.
• Passing: reorder (3) / fsdp_overlap e2e (4) / estimator (16) / compile_artifacts (38) / AOT cache reuse (4).

---

[jiahy0825]

LGTM