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

[wtr0504]

…y by structure, not name

🗂️ PR Category

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

📝 Description

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 (e.g. the local token count, marked dynamic at the same point of the same trace), 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/deduplication history differs per rank. Observed on wan2.2 (2 ranks): the all_to_all output prints s27 + s82 on rank 0 vs s27 + s74 on rank 1 for structurally identical graphs — the primary (local-seq) symbol matches, the second one diverges; a node-by-node fingerprint diff showed only the all_to_all and its downstream attention nodes differing, everything else identical.

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 — nsys showed 82–93% of all-gather time with the compute stream idle.

The two affected checks

1. reorder.py::_graph_fingerprint — the pre-reorder "per-rank graphs are structurally identical" digest hashed raw size reprs. Isomorphic graphs fingerprinted differently → the fail-fast fired → reorder disabled entirely ("overlap OFF for this graph").

Fix: canonicalize s\d+ symbols by first appearance (α-renaming) before hashing. What is structural is the symbols' occurrence pattern, not their names; genuinely different symbolic structure still diverges.

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 (equal counts, unequal reprs) → the whole cost table silently fell back to the analytical estimate. Inductor's roofline overestimates matmul ~400x here (mm: 19µs real vs 7405µs analytical), so the reorder believed one adjacent mm hides any gather and moved every launch by ~2 slots — verdict "hidden", actually exposed.

Fix: key symbolic dims by their size hint (("~", hint)) instead — guard-free (no SymInt specialization), rank-identical (same inputs per rank), and isomorphic kernels keep sharing one measurement. If ranks ever see genuinely different hints, the key-set check degrades safely to analytical as before (no deadlock, no misordering).

Verification

wan2.2, 2×H100, 1280×704, caches off:

• warm_and_sync reconciles — every cost-table entry now tagged measured (previously all analytical);
• placement log: launches hoist hundreds of slots (e.g. cur=2479 → target=2185), 9/10 buckets verdict "hidden";
• exposed all-gather time: 93% → 19% (residual = the graph-head bucket with no upstream compute to hide it — inherent);
• models whose collectives sit inside eager custom-op boundaries never traced the extra symbol and are unaffected (regression suites confirm).

Tests: new α-equivalence unit test for the fingerprint (test_fingerprint_canonicalizes_shape_symbols); estimator key test updated to assert hint-based keys (test_static_symbolic_dim_keyed_by_hint); reorder (3) / e2e (4) / estimator (16) suites pass.