[Feat] FSDP Fullgraph Overlap — whole-graph weight all-gather / compute overlap

[wtr0504]

🗂️ PR Category

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

📝 Description

Summary

This PR adds an end-to-end, compile-time FSDP weight all-gather / compute overlap feature for the single-fullgraph (non-split) compilation path. When enable_fsdp_fullgraph_overlap=True (requires disable_graph_split=True and cudagraph_mode=NONE), MagiCompiler:

1. Lowers SimpleFSDP weight prim_redistribute nodes into explicit functional all_gather_into_tensor + wait_tensor collectives on the FX graph (passes/fsdp_overlap/redistribute_lowering.py).
2. Optionally buckets the per-parameter gathers into one all_gather_into_tensor_coalesced per (compute-region, process-group, dtype), with an optional per-bucket size cap (passes/fsdp_overlap/bucket_all_gather.py, entry point lower_and_bucket_full_graph).
3. Installs a scheduler-level "latest-safe-launch" reorder pass (passes/fsdp_overlap/reorder.py, replaces Inductor's builtin raise_comms/sink_waits): a single back-to-front two-pointer sweep places each all-gather launch (wait stays put) at the latest position whose upstream compute still hides the collective (compute_window >= comm * contention_factor + slack).
4. Sizes the sweep with a new profiling-based cost model (profiling/runtime_estimator.py): Inductor's analytical roofline is unusable for our nodes (0 for custom ops, ~60x low for deep fused pointwise, ~1500x high for matmul — see scripts/demo/research_estimate_op_runtime_findings.md), so we measure real kernel times instead.

Cost model: ProfilingRuntimeEstimator

A callable snode -> nanoseconds passed directly to the reorder pass as cost_fn (deliberately not installed at the global config.estimate_op_runtime, which Inductor's own scheduling/caching consults in ways that specialize dynamic shapes).

• Op -> time table: keyed by structural identity (target op, input shapes/dtypes) — not the per-node name — so isomorphic ops across repeated layers share one measurement (O(#distinct kernels), not O(#nodes)).
• Fused Triton nodes → scheduler.benchmark_fused_nodes (codegens the same fused kernel production emits; verified 0.995–1.15x of the compiled-graph kernel). Skipped (analytical fallback) while the graph still has free symbols, since that path would specialize dynamic dims.
• Extern nodes (matmul / flash-attn / MoE custom ops) → eager replay on real tensors rebuilt from fx fake-tensor meta via size_hint (no guards added), under @torch._dynamo.disable + no_grad, inside a ShapeEnv sandbox (_shapeenv_sandbox + suppress_guards) so measurement never leaks an Eq(sym, hint) specialization. Safe even on the dynamic base compile.
• Collectives are never benchmarked inside __call__ (per-rank compilation is not co-scheduled; issuing NCCL there desyncs ranks → watchdog hang). They are seeded with Inductor's static analytical estimate and, in profile_sync mode, later overridden by a real rank-lockstep measurement (below).
• Never raises: any measurement failure degrades to the analytical estimate, so compilation is never broken.
• Per-op benchmark-input registry (profiling/benchmark_inputs.py): ops whose replay needs value-consistent metadata (e.g. split sizes feeding an internal CP all_to_all) register a hook via register_benchmark_inputs(op_name, fn, has_internal_collective=...); MagiCompiler holds no model-specific op names.

Multi-rank correctness

The reorder decides how far to hoist each gather from costs; all FSDP gathers run on one process group, so NCCL matches the Nth call positionally — ranks must issue them in identical relative order or they deadlock (observed and root-caused via flight recorder on 8-GPU gaga4). Rank-consistency requires both:

1. Identical input graph per rank — the redistribute lowering emits the pad node unconditionally (zero-width no-op on full-chunk ranks), so uneven Shard(0) params no longer give trailing ranks extra constant_pad_nd nodes.
2. Rank-identical costs — selectable via fsdp_overlap_cost_mode:
   • profile_sync (default): the estimator's warm_and_sync() re-measures every distinct op in rank-lockstep (barrier + fixed iteration counts, so ops with internal collectives issue the same number of NCCL calls on every rank), then MAX-reduces the tables over a dedicated gloo group — real measured costs that are identical on every rank.
   • analytical: Inductor roofline (deterministic by construction, zero overhead, less accurate).
3. Order clamp: targets are additionally clamped to be non-decreasing in original program index, so weight-gather relative order is preserved on all ranks regardless of cost jitter.
4. On any sync failure the pass bails and leaves the graph unchanged (overlap off, no hang).

New config options (CompileConfig)

Option	Default	Purpose
disable_graph_split	False	Hand the whole graph to Inductor as a single piecewise submodule (no FX split at boundary ops).
enable_fsdp_fullgraph_overlap	False	Master switch for the lowering + bucketing + reorder pipeline (requires disable_graph_split=True, cudagraph_mode=NONE).
fsdp_fullgraph_bucket_mode	"none"	"none" (N individual gathers + waits) or "coalesced" (one coalesced gather per region/group/dtype).
fsdp_fullgraph_bucket_size_mib	—	Optional per-bucket size cap (MiB of local shard).
fsdp_overlap_cost_mode	"profile_sync"	profile_sync | analytical | profile (see multi-rank section).
fsdp_overlap_slack_ns	5000.0	Headroom added to each collective's runtime (absorbs estimator error + launch latency).
fsdp_overlap_comm_contention_factor	1.5	Isolated→in-situ scaling of measured comm time when sizing the compute window.

Observability

• INFO: per-compile summary — repositioned M/N weight all-gather launch(es) (cost table: D distinct ops, measured=X reused=Y); rank-synchronized profiling done (N cost entries reconciled).
• DEBUG: one placement line per gather (cur -> target, dep floor, comm/acc_upstream/need, verdict hidden / COMPUTE-LIMITED) answering "why didn't this gather move earlier"; the full op->time table with grep-friendly ESTLINE|kind|label|per_call_us|calls|total_us|measured rows diffable against an nsys trace (scripts/demo/compare_estimate_vs_nsys.py).
• Rebuilt orders are validated against real data deps (WeakDeps excluded — the pass legally crosses Inductor's advisory collective-ordering edges); on validation failure the reorder is dropped with a warning.

Misc

• utils/visualize/visualizer.py: graphviz label sanitization/truncation (backslash escapes, 16384-char label limit) and render failures downgraded to a warning — visualization can no longer abort compilation.

Tests

• tests/feature_tests/test_fsdp_overlap_lowering.py — prim_redistribute → all_gather+wait lowering (incl. unconditional pad / uneven shard determinism).
• tests/feature_tests/test_fsdp_overlap_bucket.py — per-region coalesced bucketing, size caps, group/dtype partitioning.
• tests/feature_tests/test_fsdp_overlap_reorder.py — sweep placement, dep floors, order clamp, validation (via fsdp_overlap_helper/reorder_helper.py).
• tests/feature_tests/test_fsdp_overlap_e2e.py — end-to-end overlap on a real compile, including a world=2 spawn test exercising profile_sync / warm_and_sync rank-lockstep without deadlock.
• tests/feature_tests/test_profiling_estimator.py — table memoization, _static/_realize_arg, extern measurement accuracy vs independent CUDA-event timing, deepcopy safety; fsdp_overlap_helper/estimator_collective_helper.py covers the collective seed → warm_and_sync measured-override path (distributed).
• tests/feature_tests/test_profiling_registry.py — register_benchmark_inputs hook registry & has_internal_collective flagging.

[jiahy0825]

LGTM