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

magi_compiler/config.py

@@ -263,6 +263,85 @@ class CompileConfig(BaseSettings):
             "Each sub-graph between two splitting ops is compiled independently by Inductor."
         ),
     )
+    disable_graph_split: bool = Field(
+        False,
+        description=(
+            "Skip FX-level splitting at the custom subgraph-boundary ops (splitting_ops) and hand the "
+            "WHOLE graph to Inductor as a single piecewise submodule."
+        ),
+    )
+    enable_fsdp_fullgraph_overlap: bool = Field(
+        False,
+        description=(
+            "Whole-graph FSDP weight all-gather / compute overlap (requires disable_graph_split=True, "
+            "cudagraph_mode=NONE). Before Inductor, lower SimpleFSDP weight prim_redistribute to explicit "
+            "collectives and optionally bucket them per compute-region (fsdp_fullgraph_bucket_mode). Then "
+            "install a scheduler-level 'latest-safe-launch' reorder pass (replaces raise_comms/sink_waits): "
+            "each weight all-gather launch is placed at the LATEST position whose downstream compute still "
+            "hides the collective (compute_window >= comm_runtime + slack), instead of PyTorch's earliest. "
+            "Per-snode compute/comm time comes from a profiling estimator installed at config.estimate_op_runtime "
+        ),
+    )
+    fsdp_fullgraph_bucket_mode: str = Field(
+        "none",
+        description=(
+            "Bucketing strategy for enable_fsdp_fullgraph_overlap: 'none' (N individual all_gather + N waits) "
+            "or 'coalesced' (per region: one all_gather_into_tensor_coalesced -> ONE launch, N getitems, N "
+            "waits). Regions are "
+            "delimited by the model's subgraph-boundary ops (splitting_ops) so weights coalesce per compute "
+            "region, not across the whole model."
+        ),
+    )
+    fsdp_fullgraph_bucket_size_mib: int = Field(
+        0,
+        ge=0,
+        description=(
+            "Optional per-bucket size cap (MiB of local shard) for enable_fsdp_fullgraph_overlap bucketing. "
+            "0 (default) = no cap: coalesce every same-(group,dtype) weight in a region into ONE all_gather "
+            "(current behavior). >0 = within a region, walk the weight all-gathers in PROGRAM ORDER and start a "
+            "NEW bucket whenever the dtype changes OR adding the next weight's local-shard size would exceed "
+            "this cap; each bucket becomes one all_gather. Only same-dtype, program-adjacent weights coalesce "
+            "(a different-dtype weight in between breaks the run). Bounds the size of each coalesced collective."
+        ),
+    )
+    fsdp_overlap_cost_mode: Literal["profile_sync", "analytical", "profile"] = Field(
+        "profile_sync",
+        description=(
+            "Cost model the enable_fsdp_fullgraph_overlap reorder pass uses to size each weight "
+            "all-gather's compute window. The costs MUST be rank-identical multi-rank, else gathers "
+            "interleave with the CP all_to_all in rank-divergent order -> NCCL deadlock.\n"
+            "  * 'profile_sync' (default): REAL per-op profiling re-measured in rank-lockstep (barrier + "
+            "fixed iters, MAX-reduced over gloo) -> accurate AND rank-identical. Requires the compiled "
+            "graph to be structurally identical on every rank (guaranteed for gaga4 by replicating "
+            "uneven-Shard(0) params).\n"
+            "  * 'analytical': Inductor roofline (shapes+device only) -> rank-deterministic, zero overhead, "
+            "less accurate. Deadlock-free FALLBACK for a model that still shows per-rank graph divergence.\n"
+            "  * 'profile': plain per-rank profiling, NO cross-rank sync -> rank-nondeterministic, WILL "
+            "deadlock multi-rank; single-rank (world_size==1) only."
+        ),
+    )
+    fsdp_overlap_slack_ns: float = Field(
+        5000.0,
+        ge=0.0,
+        description=(
+            "Extra headroom (ns) added to each collective's runtime when the fullgraph overlap reorder pass "
+            "sizes the compute window, absorbing estimator error + kernel-launch latency so the wait rarely stalls."
+        ),
+    )
+    fsdp_overlap_comm_contention_factor: float = Field(
+        1.5,
+        ge=1.0,
+        description=(
+            "Multiplier on each collective's estimated runtime when the fullgraph overlap reorder pass sizes "
+            "its compute window (need = comm * factor + slack). The profile_sync measurement is taken in "
+            "ISOLATION (rank-lockstep, nothing else running), but in the compiled graph the all-gather runs "
+            "CONCURRENT with the compute it hides and they contend for SMs/bandwidth -- measured in-situ NCCL "
+            "all-gathers run ~1.4-1.5x their isolated time (8xH100 NVSwitch, RING_LL, NCCL_MAX_CTAS=12; nsys "
+            "2026-07-12: ws2 2x[384,256,1280] 4540us vs 3205us isolated, ws8 2x[6144,3072] 2748us vs 1897us). "
+            "Under-reserving exposes the whole tail of the collective (the wait stalls); over-reserving merely "
+            "launches a few hundred us earlier on an otherwise-idle comm stream, so err high."
+        ),
+    )
 
     # ---- torch.compile options keys ----
     post_grad_pass: str = Field(

magi_compiler/magi_backend/magi_backend.py

@@ -542,14 +542,93 @@ def _init_cache(self) -> str:
         self.local_magi_cache_path.mkdir(parents=True, exist_ok=True)
         self.compiler_manager.initialize_cache(self.local_magi_cache_path)
 
+    def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None:
+        """Whole-graph FSDP all-gather / compute overlap (disable_graph_split path).
+
+        1. Lower SimpleFSDP weight prim_redistribute -> explicit collectives and
+           optionally bucket them per compute-region.
+        2. Install the profiling runtime estimator at ``config.estimate_op_runtime``
+           (the analytical roofline is unusable for our sizing decisions).
+        3. Install the latest-safe-launch reorder pass, REPLACING PyTorch's builtin
+           raise_comms/sink_waits, and enable reorder_for_compute_comm_overlap.
+        """
+        if not self.compile_config.disable_graph_split:
+            raise ValueError("enable_fsdp_fullgraph_overlap requires disable_graph_split=True")
+
+        if self.compile_config.cudagraph_mode != CudaGraphMode.NONE:
+            raise ValueError("enable_fsdp_fullgraph_overlap requires cudagraph_mode=NONE")
+
+        from magi_compiler.passes.fsdp_overlap import FsdpOverlapReorder, lower_and_bucket_full_graph
+        from magi_compiler.profiling import ProfilingRuntimeEstimator
+
+        boundary_ops = resolve_defined_ops(
+            [] if self.compile_config.disable_graph_split else [] or self.compile_config.splitting_ops
+        )
+        bucket_size_bytes = int(self.compile_config.fsdp_fullgraph_bucket_size_mib) * 1024 * 1024
+        n_buckets = lower_and_bucket_full_graph(
+            graph,
+            self.compile_config.fsdp_fullgraph_bucket_mode,
+            boundary_ops=boundary_ops,
+            bucket_size_bytes=bucket_size_bytes,
+        )
+        magi_logger.info(
+            "FSDP fullgraph overlap: bucket_mode=%s bucket_size=%d MiB created %d buckets",
+            self.compile_config.fsdp_fullgraph_bucket_mode,
+            self.compile_config.fsdp_fullgraph_bucket_size_mib,
+            n_buckets,
+        )
+
+        _mode = self.compile_config.fsdp_overlap_cost_mode
+        if _mode == "analytical":
+            self._fsdp_overlap_estimator = None
+            cost_fn = None  # reorder pass defaults to Inductor's analytical estimate
+        else:
+            self._fsdp_overlap_estimator = ProfilingRuntimeEstimator()
+            self._fsdp_overlap_estimator._sync_across_ranks = _mode == "profile_sync"
+            cost_fn = self._fsdp_overlap_estimator
+
+        reorder = FsdpOverlapReorder(
+            slack_ns=self.compile_config.fsdp_overlap_slack_ns,
+            cost_fn=cost_fn,
+            comm_contention_factor=self.compile_config.fsdp_overlap_comm_contention_factor,
+        )
+        self.inductor_compile_config["reorder_for_compute_comm_overlap"] = True
+        self.inductor_compile_config["reorder_for_compute_comm_overlap_passes"] = [reorder]
+
     @observe_lifecycle("graph_split")
     def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[SplitItem]]:
-        # Step 1: resolve the splitting ops
-        fx_split_ops = self.compile_config.splitting_ops or []
+        # Step 1: resolve the splitting ops.
+        # When disable_graph_split is set, we deliberately resolve to NO splitting
+        # ops so Step 2 assigns every node to a single subgraph_id -> the whole
+        # graph becomes one piecewise submod handed to Inductor as a unit (the
+        # boundary custom ops remain in the graph as opaque extern calls). The
+        # PIECEWISE cudagraph path assumes >=1 split; forbid the combination.
+        if self.compile_config.disable_graph_split:
+            if self.compile_config.cudagraph_mode == CudaGraphMode.PIECEWISE:
+                raise ValueError("disable_graph_split is incompatible with cudagraph_mode=PIECEWISE")
+            fx_split_ops = []
+            magi_logger.info(
+                "disable_graph_split=True: skipping FX-level graph split; compiling the whole graph as one submod"
+            )
+        else:
+            fx_split_ops = self.compile_config.splitting_ops or []
         resolved_ops: list[torch._ops.OpOverload] = resolve_defined_ops(fx_split_ops)
         magi_logger.info(f"Setting up FX-level graph split with ops: {fx_split_ops=}")
         magi_logger.info(f"Resolved splitting ops for FX-level graph split: {resolved_ops=}")
 
+        # Step 1.4: whole-graph FSDP overlap. Lower SimpleFSDP weight redistribute
+        # to explicit collectives and (optionally) bucket them per compute-region,
+        # then install the latest-safe-launch scheduler reorder pass + profiling
+        # runtime estimator.  Regions are delimited by the model's real
+        # subgraph-boundary ops (the resolved splitting ops) even though
+        # disable_graph_split has emptied fx_split_ops -- the graph itself stays
+        # unsplit (Step 2 assigns every node one sid -> a single submod), so we
+        # pass the model's true boundary ops explicitly for region bucketing.
+        # The piecewise Steps 1.5/2.4/2.5 stay OFF (their config flags default
+        # False) so they never interfere with this path.
+        if self.compile_config.enable_fsdp_fullgraph_overlap:
+            self._apply_fsdp_fullgraph_overlap(graph)
+
         # Step 2: split graph by ops, we split graph based on resolved_ops, which becomes the partitioned single graph.
         subgraph_id = 0
         node_to_subgraph_id = {}

magi_compiler/passes/fsdp_overlap/__init__.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Whole-graph FSDP all-gather / compute overlap (``enable_fsdp_fullgraph_overlap``).
+
+The pipeline, all for the ``disable_graph_split=True`` path (there is no FX-level
+graph split here despite the historical name):
+
+1. :mod:`.redistribute_lowering` -- lower SimpleFSDP weight ``prim_redistribute``
+   into explicit ``all_gather`` + ``wait`` collectives.
+2. :mod:`.bucket_all_gather` -- coalesce the per-region weight gathers into one
+   ``all_gather_into_tensor_coalesced`` per (region, group, dtype), optionally
+   size-capped. :mod:`.lower_and_bucket` is the entry point wrapping steps 1-2.
+3. :mod:`.reorder` -- the scheduler-level ``FsdpOverlapReorder`` pass that hoists
+   each weight all-gather launch into upstream compute so the collective is hidden.
+
+The per-op cost model that feeds the reorder its compute/comm timings lives in the
+general :mod:`magi_compiler.profiling` package (``ProfilingRuntimeEstimator`` +
+``register_benchmark_inputs``), not here -- it is not FSDP-specific.
+"""
+
+from .bucket_all_gather import bucket_weight_all_gather_coalesced_per_region
+from .lower_and_bucket import lower_and_bucket_full_graph
+from .redistribute_lowering import lower_prim_redistribute_to_collectives
+from .reorder import FsdpOverlapReorder
+
+__all__ = [
+    "bucket_weight_all_gather_coalesced_per_region",
+    "lower_prim_redistribute_to_collectives",
+    "lower_and_bucket_full_graph",
+    "FsdpOverlapReorder",
+]