From cc8f966dca52a5ae2d9e7eca097774bbe7af6317 Mon Sep 17 00:00:00 2001 From: wtr Date: Tue, 30 Jun 2026 17:59:50 +0800 Subject: [PATCH 1/7] [Feat] Add fsdp weight prefetch in infer --- magi_compiler/config.py | 77 ++++ magi_compiler/magi_backend/magi_backend.py | 104 ++++- magi_compiler/passes/graph_split/__init__.py | 13 + .../graph_split/fsdp_bucket_all_gather.py | 342 ++++++++++++++++ .../graph_split/fsdp_collective_prefetch.py | 367 ++++++++++++++++++ .../graph_split/fsdp_redistribute_lowering.py | 170 ++++++++ magi_compiler/utils/visualize/visualizer.py | 35 +- .../test_fsdp_collective_prefetch.py | 229 +++++++++++ 8 files changed, 1332 insertions(+), 5 deletions(-) create mode 100644 magi_compiler/passes/graph_split/__init__.py create mode 100644 magi_compiler/passes/graph_split/fsdp_bucket_all_gather.py create mode 100644 magi_compiler/passes/graph_split/fsdp_collective_prefetch.py create mode 100644 magi_compiler/passes/graph_split/fsdp_redistribute_lowering.py create mode 100644 tests/feature_tests/test_fsdp_collective_prefetch.py diff --git a/magi_compiler/config.py b/magi_compiler/config.py index 44fbd03..725f5cb 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -263,6 +263,83 @@ class CompileConfig(BaseSettings): "Each sub-graph between two splitting ops is compiled independently by Inductor." ), ) + enable_fsdp_all_gather_prefetch: bool = Field( + False, + description=( + "Enable one-submod-ahead prefetch for SimpleFSDP/DTensor weight all-gather. " + "Before split_module, MagiCompiler moves prim_redistribute to the beginning of the previous submod " + "and leaves prim_to_local at the original weight-use point." + ), + ) + fsdp_all_gather_prefetch_distance: int = Field( + 1, + ge=0, + description="How many submods ahead to launch SimpleFSDP weight all-gather; the initial supported value is 1.", + ) + fsdp_prefetch_mode: str = Field( + "collective", + description=( + "Which FSDP all-gather prefetch pass to use when enable_fsdp_all_gather_prefetch is set. " + "'collective' (default): operate on explicit _c10d_functional.all_gather_into_tensor/wait_tensor " + "nodes, moving ONLY the launch to the previous submod and leaving the wait at the use site " + "(this is the one that actually overlaps). 'redistribute': legacy pass that moves the DTensor " + "prim_redistribute node (kept for reference; does not separate launch from wait). " + "NOTE: the 'collective' FX pass only fires when the pre-split graph already contains explicit " + "all_gather_into_tensor nodes. The real gaga4 SimpleFSDP graph instead contains opaque DTensor " + "prim_redistribute, so that pass is a no-op there; for real overlap on gaga4 use " + "enable_inductor_comm_reorder below (the within-submod Inductor reorder works regardless of " + "how the collectives originate)." + ), + ) + fsdp_lower_redistribute_to_collectives: bool = Field( + False, + description=( + "Before the collective prefetch pass, decompose each SimpleFSDP weight DTensor " + "prim_redistribute/prim_to_local into explicit _c10d_functional.all_gather_into_tensor + " + "wait_tensor FX nodes. The real gaga4 graph only has opaque prim_redistribute (no explicit " + "collectives), so this is required for the collective prefetch pass to fire and to let the " + "all_gather launch be moved across the MoE boundary while the wait stays at the use site. " + "Shard(0) only; requires cudagraph_mode=NONE." + ), + ) + fsdp_bucket_weight_all_gather: bool = Field( + False, + description=( + "After lowering SimpleFSDP weight prim_redistribute to explicit collectives, coalesce each " + "submod's weight all-gathers into one all_gather_into_tensor_coalesced per (submod, group, " + "dtype) -- fewer launched NCCL kernels. Waits are kept separate (one per weight) so the " + "collective prefetch can still hoist the coalesced launch across the MoE boundary while each " + "wait stays at its use site. Requires fsdp_lower_redistribute_to_collectives=True and cudagraph NONE." + ), + ) + fsdp_bucket_mode: str = Field( + "concat", + description=( + "Which bucketing implementation to use when fsdp_bucket_weight_all_gather is set. " + "'concat' (default): flatten each weight's local shard, cat them, do ONE " + "all_gather_into_tensor on the concatenated buffer, then reshape + split_with_sizes back to " + "each weight. The gathered buffer is rank-major while weights are weight-major, so the split " + "is lowered as a clone into a fresh buffer per weight (transient ~2x memory + a multi-GB cat " + "on the compute stream). 'coalesced': fuse the N launches into a single " + "all_gather_into_tensor_coalesced (one NCCL group, one weight-major output buffer per input), " + "then unpack each member with operator.getitem (zero-copy, no cat/clone/2x spike) and keep a " + "separate wait per member. The coalesced op returns a Tensor[] that cannot cross a piecewise " + "submod boundary, so the launch and its getitems are kept together in the launch submod; only " + "the unpacked single tensors cross. Requires fsdp_bucket_weight_all_gather=True and cudagraph NONE." + ), + ) + enable_inductor_comm_reorder: bool = Field( + False, + description=( + "Enable Inductor's comm/compute overlap reorder for every piecewise submod " + "(reorder_for_compute_comm_overlap=True with passes raise_comms + sink_waits). This hoists each " + "submod's weight all-gather launches to the top of the generated code and sinks each wait to just " + "before its use, so a gather overlaps the compute that precedes its consumer. Works on lowered " + "_c10d_functional collectives inside each submod, so it applies to gaga4 SimpleFSDP regardless of " + "the prefetch pass. The collective prefetch pass also turns this on automatically when it moves a " + "launch across a submod boundary. Requires cudagraph_mode=NONE." + ), + ) # ---- torch.compile options keys ---- post_grad_pass: str = Field( diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 9569bb2..4b72d20 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -34,6 +34,12 @@ from magi_compiler.magi_depyf.timeline import observe_lifecycle, observe_lifecycle_context from magi_compiler.offload.offload_warpper import OffloadWrapper from magi_compiler.passes import CustomJointGraphPartitionFn, FullGraphPassManager, PostGradPassManager, pass_context +from magi_compiler.passes.graph_split import ( + apply_fsdp_collective_prefetch, + bucket_weight_all_gather_coalesced_per_submod, + bucket_weight_all_gather_per_submod, + lower_prim_redistribute_to_collectives, +) from magi_compiler.utils import compilation_counter, compute_code_hash, compute_hash, magi_logger from magi_compiler.utils.visualize import save_fx_graph_visualization @@ -542,6 +548,22 @@ 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 _enable_inductor_comm_reorder(self) -> None: + """Turn on Inductor's comm/compute overlap reorder for piecewise submods. + + ``raise_comms`` hoists each submod's collective launches to the top of + the generated code (Inductor otherwise sinks a collective whose result + is only a graph output to just before ``return``, defeating prefetch); + ``sink_waits`` pushes each ``wait_tensor`` down to just before its first + use. Together they let a weight all-gather overlap the compute that runs + before its consumer. Inductor's default is + ``reorder_for_compute_comm_overlap=False``. + """ + self.inductor_compile_config.setdefault("reorder_for_compute_comm_overlap", True) + self.inductor_compile_config.setdefault( + "reorder_for_compute_comm_overlap_passes", ["raise_comms", "sink_waits"] + ) + @observe_lifecycle("graph_split") def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[SplitItem]]: # Step 1: resolve the splitting ops @@ -550,6 +572,17 @@ def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[Spli 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.5: optionally lower SimpleFSDP weight DTensor prim_redistribute/ + # prim_to_local into explicit all_gather_into_tensor + wait_tensor nodes. + # Must run BEFORE Step 2 builds node_to_subgraph_id (it mutates the graph). + # This is what makes the 'collective' prefetch pass fire on the real gaga4 + # graph (which otherwise only contains opaque prim_redistribute) and lets + # the launch be moved across the MoE boundary while the wait stays put. + if self.compile_config.fsdp_lower_redistribute_to_collectives: + if self.compile_config.cudagraph_mode != CudaGraphMode.NONE: + raise ValueError("FSDP redistribute lowering currently requires cudagraph_mode=NONE") + lower_prim_redistribute_to_collectives(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 = {} @@ -569,6 +602,75 @@ def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[Spli else: node_to_subgraph_id[node] = subgraph_id + # Step 2.4: coalesce each submod's lowered weight all-gathers into one + # all_gather_into_tensor_coalesced per (submod, group, dtype) -- fewer + # launched NCCL kernels. Must run after node_to_subgraph_id is built and + # BEFORE the prefetch pass (so prefetch sees one coalesced launch/submod). + if self.compile_config.fsdp_bucket_weight_all_gather: + if self.compile_config.cudagraph_mode != CudaGraphMode.NONE: + raise ValueError("FSDP all-gather bucketing currently requires cudagraph_mode=NONE") + bucket_mode = self.compile_config.fsdp_bucket_mode + if bucket_mode == "concat": + n_buckets = bucket_weight_all_gather_per_submod(graph, node_to_subgraph_id) + elif bucket_mode == "coalesced": + n_buckets = bucket_weight_all_gather_coalesced_per_submod(graph, node_to_subgraph_id) + else: + raise ValueError( + f"Unknown fsdp_bucket_mode={bucket_mode!r}; expected 'concat' or 'coalesced'" + ) + magi_logger.info("FSDP all-gather bucketing (%s) created %d buckets", bucket_mode, n_buckets) + + # Step 2.5: SimpleFSDP weight all-gather prefetch at submod granularity. + # Goal: launch a submod's weight all-gather during the *previous* submod's + # compute and only wait on it right before the weight is used, so the NCCL + # all-gather overlaps compute. + # + # 'collective' mode (default) operates on the explicit + # _c10d_functional.all_gather_into_tensor / wait_tensor nodes and moves + # ONLY the launch to the previous submod, leaving the wait at the use + # site. This is what actually produces overlap. 'redistribute' mode is + # the legacy pass that moves the opaque DTensor prim_redistribute node; + # it is kept for reference but does NOT separate launch from wait (Inductor + # lowers a single prim_redistribute into both the all_gather and its wait). + if self.compile_config.enable_fsdp_all_gather_prefetch: + if self.compile_config.cudagraph_mode != CudaGraphMode.NONE: + raise ValueError("FSDP all-gather prefetch currently requires cudagraph_mode=NONE") + distance = self.compile_config.fsdp_all_gather_prefetch_distance + if self.compile_config.fsdp_prefetch_mode == "collective": + moved = apply_fsdp_collective_prefetch(graph, node_to_subgraph_id, distance=distance) + magi_logger.info( + "FSDP collective prefetch moved %d all_gather_into_tensor launches by distance=%d", + moved, + distance, + ) + # The prefetch pass puts each hoisted all_gather at the FRONT of + # the previous submod's FX graph, but Inductor's per-submod + # scheduler re-sinks a collective that has no consumer inside the + # submod (its result is only a graph output) back to just before + # `return` -- so at runtime the launch is issued AFTER the + # submod's compute and cannot overlap. Enabling Inductor's + # comm/compute reorder with `raise_comms` hoists the launch back + # to the top of the generated code so it actually runs + # concurrently with this submod's compute on the NCCL stream. + if moved: + self._enable_inductor_comm_reorder() + else: + raise ValueError( + f"Unknown fsdp_prefetch_mode={self.compile_config.fsdp_prefetch_mode!r}; " + "expected 'collective' or 'redistribute'" + ) + + # Standalone Inductor comm/compute overlap reorder. Unlike the FX + # prefetch pass (which only fires when the pre-split graph has explicit + # all_gather_into_tensor nodes), this operates on the lowered collectives + # inside every submod, so it produces overlap on the real gaga4 SimpleFSDP + # graph (whose pre-split form is opaque DTensor prim_redistribute). + if self.compile_config.enable_inductor_comm_reorder: + if self.compile_config.cudagraph_mode != CudaGraphMode.NONE: + raise ValueError("Inductor comm reorder currently requires cudagraph_mode=NONE") + self._enable_inductor_comm_reorder() + magi_logger.info("Enabled Inductor comm/compute overlap reorder (raise_comms + sink_waits)") + # Step 3: split the graph based on node_to_subgraph_id # pytorch might reorder the nodes and the semantics of the graph will change when we have mutations in the graph, if we don't set keep_original_order=True split_gm = torch.fx.passes.split_module.split_module( @@ -593,7 +695,7 @@ def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[Spli # Step 5: visualize the split graph if envs.MAGI_ENABLE_FX_GRAPH_VIZ: - save_fx_graph_visualization(split_gm.graph, sub_dir="after_split", filename="split_gm_root") + # save_fx_graph_visualization(split_gm.graph, sub_dir="after_split", filename="split_gm_root") for item in piecewise_graphs: save_fx_graph_visualization(item.graph.graph, sub_dir="after_split", filename=item.submod_name) diff --git a/magi_compiler/passes/graph_split/__init__.py b/magi_compiler/passes/graph_split/__init__.py new file mode 100644 index 0000000..e5f74d2 --- /dev/null +++ b/magi_compiler/passes/graph_split/__init__.py @@ -0,0 +1,13 @@ +from .fsdp_bucket_all_gather import ( + bucket_weight_all_gather_coalesced_per_submod, + bucket_weight_all_gather_per_submod, +) +from .fsdp_collective_prefetch import apply_fsdp_collective_prefetch +from .fsdp_redistribute_lowering import lower_prim_redistribute_to_collectives + +__all__ = [ + "apply_fsdp_collective_prefetch", + "bucket_weight_all_gather_per_submod", + "bucket_weight_all_gather_coalesced_per_submod", + "lower_prim_redistribute_to_collectives", +] \ No newline at end of file diff --git a/magi_compiler/passes/graph_split/fsdp_bucket_all_gather.py b/magi_compiler/passes/graph_split/fsdp_bucket_all_gather.py new file mode 100644 index 0000000..cbc8d50 --- /dev/null +++ b/magi_compiler/passes/graph_split/fsdp_bucket_all_gather.py @@ -0,0 +1,342 @@ +from __future__ import annotations + +import math +import operator +from collections import defaultdict, deque + +import torch +import torch.fx as fx + +from magi_compiler.utils import magi_logger + +_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default +_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default +_WAIT = torch.ops._c10d_functional.wait_tensor.default +_CAT = torch.ops.aten.cat.default +_RESHAPE = torch.ops.aten.reshape.default +_SPLIT = torch.ops.aten.split_with_sizes.default + + +def _is_param_like(node: fx.Node) -> bool: + """True for a placeholder/get_attr that names a SimpleFSDP weight shard.""" + if node.op not in ("placeholder", "get_attr"): + return False + name = f"{node.name} {node.target}".lower() + return any(t in name for t in ("parameter", "parameters", "weight", "bias", "shard")) + + +def _gathers_a_weight(node: fx.Node) -> bool: + """Walk back from an all_gather through cheap producers (to_local / cast / + pad / view / reshape); return True if the gathered source is a weight shard. + + Used so the bucketing pass also fires on graphs whose weight gathers are NOT + tagged ``magi_fsdp_weight_ag`` (e.g. the demo, which emits explicit + ``all_gather_into_tensor`` directly off a ``*_shard`` parameter).""" + q: deque[fx.Node] = deque(node.all_input_nodes) + seen: set[fx.Node] = set() + while q: + dep = q.popleft() + if dep in seen: + continue + seen.add(dep) + if _is_param_like(dep): + return True + if dep.op == "call_method" and str(dep.target) in {"to_local", "contiguous", "to", "view", "reshape"}: + q.extend(dep.all_input_nodes) + elif dep.op == "call_function": + nm = getattr(dep.target, "__name__", "") or str(dep.target) + if any(t in nm for t in ("constant_pad_nd", "_to_copy", "convert_element_type", "view", "reshape", "clone")): + q.extend(dep.all_input_nodes) + return False + + +def _is_weight_all_gather(node: fx.Node) -> bool: + """A SimpleFSDP weight ``all_gather_into_tensor`` launch. + + Recognized either by the ``magi_fsdp_weight_ag`` tag (set by the redistribute + lowering pass on the real gaga4 graph) OR structurally, when the gathered + source traces back to a weight/param placeholder (covers untagged graphs such + as the demo).""" + if node.op != "call_function" or node.target is not _ALL_GATHER: + return False + if node.meta.get("magi_fsdp_weight_ag"): + return True + return _gathers_a_weight(node) + + +def _producer_chain(node: fx.Node) -> list[fx.Node]: + """The local-shard prep nodes feeding ``node`` (to_local / _to_copy / pad). + These depend only on the weight placeholder, so they may be hoisted. Returns + the movable producers (excludes placeholders/get_attr).""" + chain: list[fx.Node] = [] + seen: set[fx.Node] = set() + stack = [node] + while stack: + n = stack.pop() + for dep in n.all_input_nodes: + if dep in seen: + continue + if dep.op in ("call_function", "call_method"): + seen.add(dep) + chain.append(dep) + stack.append(dep) + return chain + + +def bucket_weight_all_gather_per_submod( + graph: fx.GraphModule, + node_to_subgraph_id: dict[fx.Node, int], +) -> int: + """Coalesce, per submod, the SimpleFSDP weight all-gathers into ONE + ``all_gather_into_tensor`` per ``(subgraph_id, group_name, dtype)`` using the + torch.compile-style flatten/cat/gather/split merge (single collective, single + wait), instead of N individual gathers. + + For a group of N weight gathers (each a single ``all_gather_into_tensor`` from + the redistribute-lowering pass, tagged ``magi_fsdp_weight_ag``, input = padded + local shard of shape ``(chunk_i, *rest_i)``), this builds:: + + # launch side (movable by the prefetch pass): + flat_i = reshape(local_i, [-1]) for each member + cat_in = cat([flat_0, ..., flat_{N-1}]) # (sum_numel,) + ag = all_gather_into_tensor(cat_in, W, group) # (W*sum_numel,) + # use side (wait stays before the consuming compute): + waited = wait_tensor(ag) # ONE wait + resh = reshape(waited, [W, sum_numel]) + splits = split_with_sizes(resh, [numel_0, ...], dim=1) + out_i = reshape(splits[i], [W*chunk_i, *rest_i]) + + ``out_i`` has exactly the shape of the member's old ``all_gather`` output, so + each member's existing downstream (the optional ``slice`` back to the true + full size, then the real use) is simply re-pointed from its old + ``wait_tensor`` to ``out_i``. The old per-member ``all_gather`` and + ``wait_tensor`` are erased. + + Only the single ``ag`` tensor crosses submod boundaries (the ``split_with_sizes`` + list and its getitems live entirely on the use side), so this avoids the AOT + output-spec mismatch that the list-returning ``all_gather_into_tensor_coalesced`` + op triggers under piecewise split. + + Runs AFTER redistribute lowering and BEFORE the collective prefetch pass. + Returns the number of coalesced buckets created. + """ + node_index = {n: i for i, n in enumerate(graph.graph.nodes)} + + groups: dict[tuple[int, str, torch.dtype], list[fx.Node]] = defaultdict(list) + for node in graph.graph.nodes: + if not _is_weight_all_gather(node): + continue + sid = node_to_subgraph_id.get(node) + if sid is None: + continue + _, _world, group_name = node.args + dtype = node.meta["example_value"].dtype + groups[(sid, group_name, dtype)].append(node) + + buckets = 0 + for (sid, group_name, dtype), ag_nodes in groups.items(): + if len(ag_nodes) < 2: + continue # nothing to coalesce + + world = int(ag_nodes[0].args[1]) + ag_nodes.sort(key=lambda n: node_index[n]) + + # Per-member: the padded local shard (ag input) and its full gathered meta. + locals_ = [ag.args[0] for ag in ag_nodes] + ag_metas = [ag.meta["example_value"] for ag in ag_nodes] # (W*chunk_i, *rest_i) + local_metas = [loc.meta["example_value"] for loc in locals_] # (chunk_i, *rest_i) + numels = [int(math.prod(m.shape)) for m in local_metas] + sum_numel = sum(numels) + dev = local_metas[0].device + + first_ag = ag_nodes[0] + # Sole user of each member's all_gather is its wait_tensor. + waits = [next(iter(ag.users)) for ag in ag_nodes] + + # Hoist every member's local-shard prep (to_local / _to_copy / pad chain) + # above the FIRST member's all_gather, so all `locals_` are defined before + # the coalesced launch we insert there. These chains depend only on the + # weight placeholder, so they can always move up. + for loc in locals_: + chain = [loc, *_producer_chain(loc)] if loc.op in ("call_function", "call_method") else _producer_chain(loc) + for prod in sorted(chain, key=lambda n: node_index[n]): + first_ag.prepend(prod) + node_to_subgraph_id[prod] = sid + + # ---- launch side: flatten each local, concat, single all_gather ---- + # Insert just before the FIRST member's all_gather (all locals now precede it). + with graph.graph.inserting_before(first_ag): + flats = [] + for loc, lm, n in zip(locals_, local_metas, numels): + fl = graph.graph.call_function(_RESHAPE, (loc, [-1])) + fl.meta["example_value"] = lm.reshape(-1) + node_to_subgraph_id[fl] = sid + flats.append(fl) + cat_in = graph.graph.call_function(_CAT, (flats, 0)) + cat_in.meta["example_value"] = local_metas[0].new_empty((sum_numel,)) + node_to_subgraph_id[cat_in] = sid + + ag = graph.graph.call_function(_ALL_GATHER, (cat_in, world, group_name)) + ag.meta["example_value"] = local_metas[0].new_empty((world * sum_numel,)) + ag.meta["magi_fsdp_weight_ag"] = True # so the prefetch pass moves it + node_to_subgraph_id[ag] = sid + + # ---- use side: one wait, reshape, split back to per-member tensors ---- + # Insert before the EARLIEST member's wait (its use site within this submod). + first_wait = min(waits, key=lambda n: node_index[n]) + with graph.graph.inserting_before(first_wait): + waited = graph.graph.call_function(_WAIT, (ag,)) + waited.meta["example_value"] = ag.meta["example_value"] + node_to_subgraph_id[waited] = sid + + resh = graph.graph.call_function(_RESHAPE, (waited, [world, sum_numel])) + resh.meta["example_value"] = local_metas[0].new_empty((world, sum_numel)) + node_to_subgraph_id[resh] = sid + + split = graph.graph.call_function(_SPLIT, (resh, numels, 1)) + split.meta["example_value"] = [local_metas[0].new_empty((world, n)) for n in numels] + node_to_subgraph_id[split] = sid + + for i, (am, lm) in enumerate(zip(ag_metas, local_metas)): + gi = graph.graph.call_function(operator.getitem, (split, i)) + gi.meta["example_value"] = local_metas[0].new_empty((world, numels[i])) + node_to_subgraph_id[gi] = sid + out_i = graph.graph.call_function(_RESHAPE, (gi, list(am.shape))) + out_i.meta["example_value"] = am # (W*chunk_i, *rest_i): same as old ag out + node_to_subgraph_id[out_i] = sid + + # Re-point the member's downstream (slice / real use) from its old + # wait to out_i, then drop the old wait + all_gather. + waits[i].replace_all_uses_with(out_i) + + for ag_old, wait_old in zip(ag_nodes, waits): + node_to_subgraph_id.pop(wait_old, None) + node_to_subgraph_id.pop(ag_old, None) + graph.graph.erase_node(wait_old) + graph.graph.erase_node(ag_old) + + buckets += 1 + + if buckets: + graph.graph.lint() + graph.recompile() + magi_logger.info( + "FSDP weight all-gather bucketing (concat): created %d coalesced buckets across submods", buckets + ) + return buckets + + +def bucket_weight_all_gather_coalesced_per_submod( + graph: fx.GraphModule, + node_to_subgraph_id: dict[fx.Node, int], +) -> int: + """Coalesce, per submod, the SimpleFSDP weight all-gathers into ONE + ``all_gather_into_tensor_coalesced`` per ``(subgraph_id, group_name, dtype)``. + + Unlike :func:`bucket_weight_all_gather_per_submod` (the ``concat`` strategy), + this does NOT cat the shards into one buffer. ``all_gather_into_tensor_coalesced`` + fuses the N launches into a single NCCL group while returning one + *weight-major* output buffer per input, so each member is recovered with a + zero-copy ``operator.getitem`` -- no ``cat`` on the compute stream, no + ``split_with_sizes`` clone, and no transient ~2x memory spike that the concat + path incurs. + + For a group of N weight gathers (each a single ``all_gather_into_tensor`` whose + input is the padded local shard ``(chunk_i, *rest_i)``) this builds:: + + # launch side (kept together in ONE submod -- the list must not cross a + # piecewise split boundary): + coalesced = all_gather_into_tensor_coalesced([local_0, ..., local_{N-1}], W, group) + out_i = getitem(coalesced, i) # (W*chunk_i, *rest_i), weight-major + # use side (one wait per member, left at its consumer): + wait_i = wait_tensor(out_i) + + ``out_i`` has exactly the shape of the member's old ``all_gather`` output, so + each member's existing downstream (slice / real use) is simply re-pointed from + its old ``wait_tensor`` to ``wait_i``. The ``coalesced`` launch and all its + ``getitem`` nodes stay in the launch submod; only the per-member ``out_i`` + single tensors cross to the consumer submod. The prefetch pass moves the + launch together with its getitems (see ``apply_fsdp_collective_prefetch``). + + Runs AFTER redistribute lowering and BEFORE the collective prefetch pass. + Returns the number of coalesced buckets created. + """ + node_index = {n: i for i, n in enumerate(graph.graph.nodes)} + + groups: dict[tuple[int, str, torch.dtype], list[fx.Node]] = defaultdict(list) + for node in graph.graph.nodes: + if not _is_weight_all_gather(node): + continue + sid = node_to_subgraph_id.get(node) + if sid is None: + continue + _, _world, group_name = node.args + dtype = node.meta["example_value"].dtype + groups[(sid, group_name, dtype)].append(node) + + buckets = 0 + for (sid, group_name, dtype), ag_nodes in groups.items(): + if len(ag_nodes) < 2: + continue # nothing to coalesce + + world = int(ag_nodes[0].args[1]) + ag_nodes.sort(key=lambda n: node_index[n]) + + locals_ = [ag.args[0] for ag in ag_nodes] + ag_metas = [ag.meta["example_value"] for ag in ag_nodes] # (W*chunk_i, *rest_i) + + first_ag = ag_nodes[0] + # Sole user of each member's all_gather is its wait_tensor. + waits = [next(iter(ag.users)) for ag in ag_nodes] + + # Hoist every member's local-shard prep (to_local / _to_copy / pad chain) + # above the FIRST member's all_gather so all `locals_` precede the + # coalesced launch we insert there. These depend only on the weight + # placeholder, so they can always move up. + for loc in locals_: + chain = [loc, *_producer_chain(loc)] if loc.op in ("call_function", "call_method") else _producer_chain(loc) + for prod in sorted(chain, key=lambda n: node_index[n]): + first_ag.prepend(prod) + node_to_subgraph_id[prod] = sid + + # ---- launch side: ONE coalesced all_gather + per-member getitem unpack ---- + # Insert before the FIRST member's all_gather (all locals now precede it). + with graph.graph.inserting_before(first_ag): + coalesced = graph.graph.call_function(_ALL_GATHER_COALESCED, (list(locals_), world, group_name)) + # The coalesced op returns a list[Tensor]; meta mirrors the per-member + # gathered shapes (weight-major, same as each old all_gather output). + coalesced.meta["example_value"] = list(ag_metas) + coalesced.meta["magi_fsdp_weight_ag"] = True # so prefetch moves it + coalesced.meta["magi_fsdp_weight_ag_coalesced"] = True # move getitems too + node_to_subgraph_id[coalesced] = sid + + outs = [] + for i, am in enumerate(ag_metas): + gi = graph.graph.call_function(operator.getitem, (coalesced, i)) + gi.meta["example_value"] = am # (W*chunk_i, *rest_i) + node_to_subgraph_id[gi] = sid + outs.append(gi) + + # ---- use side: one wait per member, left before its consumer ---- + for i, (out_i, old_wait) in enumerate(zip(outs, waits)): + with graph.graph.inserting_before(old_wait): + wait_i = graph.graph.call_function(_WAIT, (out_i,)) + wait_i.meta["example_value"] = ag_metas[i] + node_to_subgraph_id[wait_i] = node_to_subgraph_id.get(old_wait, sid) + old_wait.replace_all_uses_with(wait_i) + + for ag_old, wait_old in zip(ag_nodes, waits): + node_to_subgraph_id.pop(wait_old, None) + node_to_subgraph_id.pop(ag_old, None) + graph.graph.erase_node(wait_old) + graph.graph.erase_node(ag_old) + + buckets += 1 + + if buckets: + graph.graph.lint() + graph.recompile() + magi_logger.info( + "FSDP weight all-gather bucketing (coalesced): created %d coalesced buckets across submods", buckets + ) + return buckets diff --git a/magi_compiler/passes/graph_split/fsdp_collective_prefetch.py b/magi_compiler/passes/graph_split/fsdp_collective_prefetch.py new file mode 100644 index 0000000..fe0d838 --- /dev/null +++ b/magi_compiler/passes/graph_split/fsdp_collective_prefetch.py @@ -0,0 +1,367 @@ +from __future__ import annotations + +from collections import deque + +import torch +import torch.fx as fx + +from magi_compiler.utils import magi_logger + +# The two functional-collective ops a SimpleFSDP weight gather lowers to once +# the DTensor ``redistribute``/``to_local`` prims have been turned into explicit +# collectives. This pass operates on *these* nodes (not on the opaque DTensor +# prims), because only here are the launch (``all_gather_into_tensor``) and the +# wait (``wait_tensor``) separate, individually movable FX nodes. +_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default +_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default +_WAIT_TENSOR = torch.ops._c10d_functional.wait_tensor.default + + +def _target_name(target: object) -> str: + if hasattr(target, "name") and callable(getattr(target, "name")): + try: + return str(target.name()) + except Exception: + pass + name = getattr(target, "__name__", None) + if name is not None: + return str(name) + return str(target) + + +def _is_all_gather(node: fx.Node) -> bool: + # Both the single launch and the per-submod coalesced launch are movable. + return node.op == "call_function" and node.target in (_ALL_GATHER, _ALL_GATHER_COALESCED) + + +def _is_wait_tensor(node: fx.Node) -> bool: + return node.op == "call_function" and node.target is _WAIT_TENSOR + + +def _is_param_like(node: fx.Node) -> bool: + """True for a placeholder/get_attr that names a SimpleFSDP weight shard.""" + if node.op not in ("placeholder", "get_attr"): + return False + name = f"{node.name} {node.target}".lower() + return any(token in name for token in ("parameter", "parameters", "weight", "bias")) + + +def _is_transparent_producer(node: fx.Node) -> bool: + """Cheap, structure-only ops that may sit *between* a weight placeholder and + the ``all_gather`` (dtype cast / pad / view). Used when walking backwards + from the launch to (a) decide it gathers a weight and (b) collect the + input-producing chain that must move together with the launch.""" + # The redistribute-lowering pass extracts the local shard with + # ``call_method("to_local", (weight_placeholder,))`` before the gather, so + # the chain from all_gather back to the weight passes through it. + if node.op == "call_method" and str(node.target) in {"to_local", "contiguous", "to", "view", "reshape"}: + return True + if node.op != "call_function": + return False + name = _target_name(node.target) + return any( + token in name + for token in ( + "constant_pad_nd", + "_to_copy", + "aten.to", + "aten::to", + "convert_element_type", + # bucketing inserts cat(reshape(local),...) before the coalesced launch + "aten.cat", + "aten::cat", + "aten.view", + "aten::view", + "aten.reshape", + "aten::reshape", + "_unsafe_view", + "aten.clone", + "aten::clone", + "aten.detach", + "aten::detach", + ) + ) + + +def _is_transparent_consumer(node: fx.Node) -> bool: + """Nodes that forward the gathered weight to its real use without being the + real use themselves: the ``wait_tensor`` plus cheap view/cast/slice ops.""" + if _is_wait_tensor(node): + return True + if node.op == "call_method" and str(node.target) in { + "view", + "reshape", + "flatten", + "squeeze", + "unsqueeze", + "permute", + "transpose", + "t", + "contiguous", + "to", + "detach", + "clone", + }: + return True + if node.op != "call_function": + return False + name = _target_name(node.target) + return any( + token in name + for token in ( + "aten.view", + "aten::view", + "aten.reshape", + "aten::reshape", + "_unsafe_view", + "aten.squeeze", + "aten::squeeze", + "aten.unsqueeze", + "aten::unsqueeze", + "aten.permute", + "aten::permute", + "aten.transpose", + "aten::transpose", + "aten.t", + "aten::t", + "aten.slice", + "aten::slice", + # bucketing's use-side chain: wait -> reshape -> split_with_sizes -> getitem -> reshape + "split_with_sizes", + "aten.detach", + "aten::detach", + "aten.clone", + "aten::clone", + "aten.to", + "aten::to", + "_to_copy", + "convert_element_type", + "operator.getitem", + "getitem", + ) + ) + + +def _gathers_a_weight(all_gather: fx.Node) -> bool: + """Walk backwards from the launch through transparent producers; return True + if the gathered source is a weight/bias placeholder.""" + q: deque[fx.Node] = deque(all_gather.all_input_nodes) + seen: set[fx.Node] = set() + while q: + dep = q.popleft() + if dep in seen: + continue + seen.add(dep) + if _is_param_like(dep): + return True + if _is_transparent_producer(dep): + q.extend(dep.all_input_nodes) + return False + + +def _find_consumer_subgraph_id( + all_gather: fx.Node, node_to_subgraph_id: dict[fx.Node, int] +) -> int | None: + """Forward BFS through wait/transparent nodes to the first real consumer; + return the minimum subgraph id among real consumers.""" + q: deque[fx.Node] = deque(all_gather.users) + seen: set[fx.Node] = set() + consumer_ids: list[int] = [] + while q: + user = q.popleft() + if user in seen: + continue + seen.add(user) + if _is_transparent_consumer(user): + q.extend(user.users) + continue + sid = node_to_subgraph_id.get(user) + if sid is not None: + consumer_ids.append(sid) + if not consumer_ids: + return None + return min(consumer_ids) + + +def _collect_launch_input_chain(all_gather: fx.Node) -> list[fx.Node]: + """Backwards from the launch, collect transparent producer nodes that feed + *only* into this launch (their results are not used elsewhere). These must + move with the launch so it is self-contained in the previous submod. + + The weight placeholder/get_attr itself is NOT included (placeholders are not + movable graph nodes); only intermediate compute that prepares the shard. + """ + chain: list[fx.Node] = [] + q: deque[fx.Node] = deque(n for n in all_gather.all_input_nodes if _is_transparent_producer(n)) + seen: set[fx.Node] = set() + frontier = {all_gather} + while q: + node = q.popleft() + if node in seen: + continue + # Movable only if every consumer is within the chain we are moving. + if not all(u in frontier or u in seen or u is all_gather for u in node.users): + continue + seen.add(node) + chain.append(node) + frontier.add(node) + for dep in node.all_input_nodes: + if _is_transparent_producer(dep): + q.append(dep) + return chain + + +def _build_first_compute_node_by_subgraph_id( + graph: fx.GraphModule, node_to_subgraph_id: dict[fx.Node, int] +) -> dict[int, fx.Node]: + """First *compute* node of each submod, used as the prefetch anchor. + + The anchor MUST be a node that genuinely STAYS in the target submod — i.e. a + real compute node (matmul, custom op, norm, ...), NOT any weight-gather-chain + node. After redistribute lowering a gather is a chain + ``to_local -> [_to_copy] -> [pad] -> all_gather -> wait -> [slice]``; every one + of those is itself relocatable by this pass. If we anchored on, say, the + first ``to_local`` of the target submod, that node could be hoisted away to an + even earlier submod, and prepending another launch before the (now-moved) + anchor would drag it to the very front of the graph — corrupting submod + ordering and making ``split_module`` raise KeyError on partition inputs. + Skipping the whole gather-chain vocabulary keeps the anchor stable. + """ + first: dict[int, fx.Node] = {} + for node in graph.graph.nodes: + if node.op in ("placeholder", "output"): + continue + # Skip every node that is part of a (relocatable) weight-gather chain. + if ( + _is_all_gather(node) + or _is_wait_tensor(node) + or _is_transparent_producer(node) + or _is_transparent_consumer(node) + ): + continue + sid = node_to_subgraph_id.get(node) + if sid is not None and sid not in first: + first[sid] = node + return first + + +def _deps_available_before_anchor( + nodes: list[fx.Node], anchor: fx.Node, node_index: dict[fx.Node, int] +) -> bool: + """Every external dependency of the to-be-moved nodes must already be + defined before the anchor (so the move does not read an undefined value).""" + moved = set(nodes) + anchor_index = node_index[anchor] + for node in nodes: + for dep in node.all_input_nodes: + if dep in moved: + continue + if dep.op in ("placeholder", "get_attr"): + continue + if node_index.get(dep, anchor_index + 1) > anchor_index: + return False + return True + + +def apply_fsdp_collective_prefetch( + graph: fx.GraphModule, + node_to_subgraph_id: dict[fx.Node, int], + *, + distance: int = 1, +) -> int: + """Prefetch SimpleFSDP weight all-gather across a submod boundary. + + For each ``all_gather_into_tensor`` whose gathered result feeds a weight use + in submod ``N``, move ONLY the launch node (and the transparent producer + chain that feeds exclusively into it) to the beginning of submod + ``N - distance``. The ``wait_tensor`` and everything after it are left + untouched at the original use site in submod ``N``. + + Because ``split_module(keep_original_order=True)`` runs afterwards, the + gathered tensor automatically becomes an output of submod ``N - distance`` + and an input of submod ``N`` -- so the collective is launched during the + previous submod's compute and only waited on right before use, giving real + compute/communication overlap. + + Unlike :func:`apply_fsdp_all_gather_prefetch`, this pass never moves the + wait: at the explicit-collective level the launch and wait are distinct FX + nodes, so leaving the wait behind actually keeps it at the consumer. + + IMPORTANT -- this FX move is necessary but NOT sufficient for overlap. The + hoisted ``all_gather`` has no consumer inside submod ``N - distance`` (its + result is only a graph output), so Inductor's per-submod scheduler re-sinks + it to just before ``return`` -- at runtime the launch is then issued AFTER + that submod's compute and cannot overlap. The caller must therefore also + enable Inductor's comm/compute reorder with the ``raise_comms`` pass + (``reorder_for_compute_comm_overlap=True``, + ``reorder_for_compute_comm_overlap_passes=["raise_comms", "sink_waits"]``), + which hoists the launch back to the top of the generated code. MagiBackend + wires this automatically when this pass moves any launch. + """ + if distance <= 0: + return 0 + + first_node_by_subgraph_id = _build_first_compute_node_by_subgraph_id(graph, node_to_subgraph_id) + node_index = {node: idx for idx, node in enumerate(graph.graph.nodes)} + moved = 0 + changed = False + + for node in list(graph.graph.nodes): + if not _is_all_gather(node): + continue + if not _gathers_a_weight(node): + continue + + current_id = node_to_subgraph_id.get(node) + if current_id is None: + continue + + consumer_id = _find_consumer_subgraph_id(node, node_to_subgraph_id) + if consumer_id is None: + continue + + target_id = max(0, consumer_id - distance) + if current_id <= target_id: + # Already launched at or before the desired submod. + continue + + anchor = first_node_by_subgraph_id.get(target_id) + if anchor is None: + continue + + # Move the launch plus the producer chain that feeds only into it. + chain = _collect_launch_input_chain(node) + # For a coalesced launch the op returns a list[Tensor] that must NOT cross + # a submod boundary; its operator.getitem unpackers (whose only input is + # the launch) move together with it so the list stays in the target submod + # and only the per-member single tensors cross to the consumer. + getitems = [] + if node.target is _ALL_GATHER_COALESCED or node.meta.get("magi_fsdp_weight_ag_coalesced"): + getitems = [ + u for u in node.users + if u.op == "call_function" and getattr(u.target, "__name__", "") == "getitem" + ] + # Order from earliest to latest so prepend keeps dependency order. + move_group = sorted([node, *chain, *getitems], key=lambda n: node_index[n]) + + if not _deps_available_before_anchor(move_group, anchor, node_index): + magi_logger.debug( + "Skip collective prefetch for %s: deps not available before %s", node.name, anchor.name + ) + continue + + for moved_node in move_group: + node_to_subgraph_id[moved_node] = target_id + anchor.prepend(moved_node) + + node.meta["magi_fsdp_prefetch_from_subgraph"] = current_id + node.meta["magi_fsdp_prefetch_to_subgraph"] = target_id + node.meta["magi_fsdp_prefetch_for_consumer_subgraph"] = consumer_id + + moved += 1 + changed = True + + if changed: + graph.graph.lint() + graph.recompile() + return moved diff --git a/magi_compiler/passes/graph_split/fsdp_redistribute_lowering.py b/magi_compiler/passes/graph_split/fsdp_redistribute_lowering.py new file mode 100644 index 0000000..4e28ff7 --- /dev/null +++ b/magi_compiler/passes/graph_split/fsdp_redistribute_lowering.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import math + +import torch +import torch.fx as fx + +from magi_compiler.utils import magi_logger + +# The functional collectives we lower a SimpleFSDP weight gather into. +_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default +_WAIT = torch.ops._c10d_functional.wait_tensor.default +_PAD = torch.ops.aten.constant_pad_nd.default +_SLICE = torch.ops.aten.slice.Tensor +_TO_COPY = torch.ops.aten._to_copy.default + + +def _is_prim(node: fx.Node, name: str) -> bool: + return node.op == "call_function" and getattr(node.target, "__name__", None) == name + + +def _input_is_weight(node: fx.Node) -> bool: + """The redistribute input is a SimpleFSDP weight/bias param placeholder.""" + src = node.args[0] if node.args else None + if not isinstance(src, fx.Node) or src.op not in ("placeholder", "get_attr"): + return False + name = f"{src.name} {src.target}".lower() + return any(t in name for t in ("parameter", "parameters", "weight", "bias")) + + +def _dtensor_meta(node: fx.Node): + """Return the DTensor example_value carried on a node, or None.""" + ev = node.meta.get("example_value") + return ev if getattr(ev, "_spec", None) is not None else None + + +def lower_prim_redistribute_to_collectives(graph: fx.GraphModule) -> int: + """Rewrite SimpleFSDP weight ``prim_redistribute`` + ``prim_to_local`` pairs + into explicit functional collectives, so the launch and the wait become two + distinct FX nodes that a later graph-split pass can place in *different* + submods (enabling cross-boundary overlap with the MoE op). + + For one ``Shard(0)`` weight (full dim0 ``F``, world ``W``, local shard the + input placeholder's ``_local_tensor``), this emits the exact sequence Inductor + itself produces for a ``redistribute(Replicate()).to_local()``: + + [ _to_copy(local, fwd_dtype) ]? # only if forward_dtype set + [ constant_pad_nd(local, dim0 -> chunk) ]? # only if local < chunk (uneven) + all_gather_into_tensor(local, W, group_name) # -> (W*chunk, ...) + wait_tensor(ag) + [ slice(wait, 0, 0, F) ]? # only if W*chunk != F (uneven) + + where ``chunk = ceil(F / W)``. ``all_gather_into_tensor`` and ``wait_tensor`` + are kept as separate nodes. Only ``Shard(0)`` is handled; anything else is + left as the original prim path and logged. + + Returns the number of redistribute pairs lowered. + """ + lowered = 0 + skipped = 0 + + for node in list(graph.graph.nodes): + if not _is_prim(node, "prim_redistribute"): + continue + if not _input_is_weight(node): + continue + + # prim_to_local consumer (the node whose output the rest of the graph uses). + to_local = next((u for u in node.users if _is_prim(u, "prim_to_local")), None) + if to_local is None: + skipped += 1 + continue + + src = node.args[0] + out_dt = _dtensor_meta(node) + in_dt = _dtensor_meta(src) + if out_dt is None or in_dt is None: + skipped += 1 + continue + + spec = in_dt._spec + placement = spec.placements[0] if len(spec.placements) == 1 else None + if placement is None or not getattr(placement, "is_shard", lambda: False)() or placement.dim != 0: + # Only single-mesh Shard(0) is supported; leave others on the prim path. + skipped += 1 + continue + + mesh = spec.mesh + try: + group_name = mesh._dim_group_names[0] + world = int(mesh.size(0)) + except Exception as exc: # pragma: no cover - defensive + magi_logger.warning("FSDP lowering skip %s: cannot resolve group/world (%s)", node.name, exc) + skipped += 1 + continue + + full = out_dt._local_tensor # FakeTensor with the global (replicated) shape + local = in_dt._local_tensor # FakeTensor with this rank's shard shape + F = int(full.shape[0]) + L = int(local.shape[0]) + chunk = math.ceil(F / world) + + # forward_dtype: cast the local shard before the gather (matches torchtitan). + fwd_dtype = None + try: + import inspect + + fwd_dtype = inspect.getclosurevars(node.target).nonlocals.get("kwargs_as_value", {}).get("forward_dtype") + except Exception: + fwd_dtype = None + + with graph.graph.inserting_before(node): + # The weight placeholder is still a Shard(0) DTensor; the functional + # collective needs a PLAIN local tensor (DTensor has no sharding + # strategy for all_gather_into_tensor). Extract the local shard first + # -- the same `param.to_local()` step the eager _materialize_param uses. + cur = graph.graph.call_method("to_local", (src,)) + cur.meta["example_value"] = local + cur_dtype = local.dtype + + if fwd_dtype is not None and fwd_dtype != cur_dtype: + cur = graph.graph.call_function(_TO_COPY, (cur,), {"dtype": fwd_dtype}) + cur.meta["example_value"] = local.to(fwd_dtype) + cur_dtype = fwd_dtype + + # Pad the local shard up to `chunk` rows when this rank owns fewer + # (uneven Shard(0): trailing ranks get remainder/empty). + # + # NOTE: meta example_values MUST be created via FakeTensor ops + # (``local.new_empty``), NOT ``torch.empty(..., device=cuda)``. A + # FakeTensor's ``.device`` is a real ``cuda:N``, so ``torch.empty`` + # allocates a REAL full-size buffer per weight; across all weights + # that leaks ~the whole unsharded model (tens of GB) and OOMs. + if L < chunk: + pad = [0, 0] * (local.dim() - 1) + [0, chunk - L] + padded = graph.graph.call_function(_PAD, (cur, pad, 0.0)) + padded.meta["example_value"] = local.new_empty((chunk, *local.shape[1:]), dtype=cur_dtype) + cur = padded + + ag = graph.graph.call_function(_ALL_GATHER, (cur, world, group_name)) + ag.meta["example_value"] = local.new_empty((world * chunk, *local.shape[1:]), dtype=cur_dtype) + # Mark as a SimpleFSDP weight gather so the per-submod bucketing pass + # can coalesce these into a single all_gather_into_tensor_coalesced. + ag.meta["magi_fsdp_weight_ag"] = True + + wait = graph.graph.call_function(_WAIT, (ag,)) + wait.meta["example_value"] = local.new_empty((world * chunk, *local.shape[1:]), dtype=cur_dtype) + + result = wait + if world * chunk != F: + sliced = graph.graph.call_function(_SLICE, (wait, 0, 0, F)) + sliced.meta["example_value"] = local.new_empty((F, *local.shape[1:]), dtype=cur_dtype) + result = sliced + + # Re-point every user of prim_to_local at the explicit-collective result, + # then drop the two prim nodes. + to_local.replace_all_uses_with(result) + graph.graph.erase_node(to_local) + graph.graph.erase_node(node) + lowered += 1 + + if lowered: + graph.graph.lint() + graph.recompile() + magi_logger.info( + "FSDP redistribute lowering: lowered %d weight prim_redistribute -> explicit collectives (skipped %d)", + lowered, + skipped, + ) + return lowered diff --git a/magi_compiler/utils/visualize/visualizer.py b/magi_compiler/utils/visualize/visualizer.py index 3bd62e8..a6adfa0 100644 --- a/magi_compiler/utils/visualize/visualizer.py +++ b/magi_compiler/utils/visualize/visualizer.py @@ -70,6 +70,27 @@ def build_node_to_code_map(graph: torch.fx.Graph) -> Dict[torch.fx.Node, str]: return node_to_code +# Graphviz aborts with a "syntax error ... missing endquote? longer than 16384?" +# when any single quoted string (i.e. a node label) exceeds 16384 characters. +# Keep labels comfortably below that ceiling. +_MAX_LABEL_LEN = 8000 + + +def _sanitize_label(text: str) -> str: + # Graphviz treats backslashes as escapes inside quoted strings. A literal + # backslash in a target/meta repr (or one left at the end of a truncated + # string) can escape the closing quote and produce + # "syntax error in line N". Replace backslashes so the label is inert. + return text.replace("\\", "/") + + +def _truncate_label(text: str) -> str: + if len(text) <= _MAX_LABEL_LEN: + return text + omitted = len(text) - _MAX_LABEL_LEN + return text[:_MAX_LABEL_LEN] + f"\n...[truncated {omitted} chars]" + + def extract_fx_graph_structure(graph: torch.fx.Graph, simple_desc: bool = False) -> Tuple[List[Dict], List[Dict]]: import textwrap @@ -107,6 +128,7 @@ def wrap_str_to_multi_lines(text, width=30): node_label = f"Op: {node.op}\nTarget: {target_str}\nName: {name_str}\nMeta: {meta_str}\nCode: {node_code}" if simple_desc: node_label = f"Op: {node.op}\nTarget: {target_str}\nName: {name_str}" + node_label = _truncate_label(_sanitize_label(node_label)) nodes.append({"id": node.name, "style": style_info, "node_label": node_label}) @@ -223,7 +245,12 @@ def save_fx_graph_visualization(graph: torch.fx.Graph, sub_dir: str = "", filena file_path = get_fx_graph_path(sub_dir=sub_dir, filename=filename) - nodes, edges = extract_fx_graph_structure(graph, simple_desc=simple_desc) - dot = create_fx_graph_dot(nodes, edges) - dot.render(filename=file_path, view=False, cleanup=True) - magi_logger.info("FX graph visualization saved to: %s.pdf", file_path) + # Visualization is purely diagnostic; never let a render failure abort + # compilation/inference. + try: + nodes, edges = extract_fx_graph_structure(graph, simple_desc=simple_desc) + dot = create_fx_graph_dot(nodes, edges) + dot.render(filename=file_path, view=False, cleanup=True) + magi_logger.info("FX graph visualization saved to: %s.pdf", file_path) + except Exception as e: + magi_logger.warning("Skipping FX graph visualization (%s): %s", file_path, e) diff --git a/tests/feature_tests/test_fsdp_collective_prefetch.py b/tests/feature_tests/test_fsdp_collective_prefetch.py new file mode 100644 index 0000000..139b00e --- /dev/null +++ b/tests/feature_tests/test_fsdp_collective_prefetch.py @@ -0,0 +1,229 @@ +# 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. + +""" +Unit test for apply_fsdp_collective_prefetch. + +Builds a tiny FX graph that mirrors the gaga4 split structure: + + submod 0 (compute): a = x @ w0 + submod 1 (boundary): b = relu(a) # stands in for a custom-op boundary + submod 2 (compute): g = all_gather(w2) + wt = wait_tensor(g) + out = b @ wt # weight `w2` is used here + +The weight all-gather for `w2` lives in submod 2 alongside its use. With +distance=2 the pass should move ONLY the `all_gather_into_tensor` launch into +submod 0 (the previous *compute* submod, skipping the odd boundary submod), +while leaving `wait_tensor` and the matmul in submod 2. After split_module that +makes the collective launch during submod 0's compute and wait right before use. +""" + +import operator + +import torch +import torch.fx as fx + +from magi_compiler.passes.graph_split import ( + apply_fsdp_collective_prefetch, + bucket_weight_all_gather_coalesced_per_submod, +) + +_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default +_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default +_WAIT = torch.ops._c10d_functional.wait_tensor.default +_CAT = torch.ops.aten.cat.default +_SPLIT = torch.ops.aten.split_with_sizes.default + + +def _build_graph(): + g = fx.Graph() + x = g.placeholder("x") + w0 = g.placeholder("l_self_modules_proj0_parameters_weight_") + w2 = g.placeholder("l_self_modules_proj2_parameters_weight_") + + # submod 0: compute + a = g.call_function(torch.ops.aten.mm.default, (x, w0)) + # submod 1: boundary op (semantics irrelevant to the pass) + b = g.call_function(torch.ops.aten.relu.default, (a,)) + # submod 2: weight gather + wait + use + ag = g.call_function(_ALL_GATHER, (w2, 8, "0")) + wt = g.call_function(_WAIT, (ag,)) + out = g.call_function(torch.ops.aten.mm.default, (b, wt)) + g.output((out,)) + + gm = fx.GraphModule(torch.nn.Module(), g) + node_to_subgraph_id = {a: 0, b: 1, ag: 2, wt: 2, out: 2} + return gm, node_to_subgraph_id, {"a": a, "b": b, "ag": ag, "wt": wt, "out": out} + + +def test_launch_moves_wait_stays(): + gm, mapping, nodes = _build_graph() + + moved = apply_fsdp_collective_prefetch(gm, mapping, distance=2) + + assert moved == 1, f"expected exactly one launch moved, got {moved}" + # Launch relocated to the previous compute submod ... + assert mapping[nodes["ag"]] == 0, mapping[nodes["ag"]] + # ... but the wait and the real use stay at the consumer site. + assert mapping[nodes["wt"]] == 2, mapping[nodes["wt"]] + assert mapping[nodes["out"]] == 2, mapping[nodes["out"]] + + # Launch is tagged with provenance metadata. + assert nodes["ag"].meta.get("magi_fsdp_prefetch_to_subgraph") == 0 + assert nodes["ag"].meta.get("magi_fsdp_prefetch_for_consumer_subgraph") == 2 + + # The all_gather now precedes the first node of submod 0 in graph order. + order = list(gm.graph.nodes) + assert order.index(nodes["ag"]) < order.index(nodes["a"]), "launch must come before submod-0 compute" + assert order.index(nodes["wt"]) > order.index(nodes["b"]), "wait must remain after the boundary" + + # Graph stays valid. + gm.graph.lint() + + +def test_no_move_when_already_early(): + """If the consumer is in submod 0 there is no previous submod to move to.""" + gm, mapping, nodes = _build_graph() + # Pretend everything is in submod 0. + for n in nodes.values(): + mapping[n] = 0 + moved = apply_fsdp_collective_prefetch(gm, mapping, distance=2) + assert moved == 0 + assert mapping[nodes["ag"]] == 0 + + +def test_non_weight_all_gather_is_ignored(): + """An all_gather of a non-weight (activation) tensor must not be moved.""" + g = fx.Graph() + x = g.placeholder("x") # activation, not a weight + w0 = g.placeholder("l_self_modules_proj0_parameters_weight_") + a = g.call_function(torch.ops.aten.mm.default, (x, w0)) + b = g.call_function(torch.ops.aten.relu.default, (a,)) + ag = g.call_function(_ALL_GATHER, (b, 8, "0")) # gathers an activation + wt = g.call_function(_WAIT, (ag,)) + out = g.call_function(torch.ops.aten.relu.default, (wt,)) + g.output((out,)) + gm = fx.GraphModule(torch.nn.Module(), g) + mapping = {a: 0, b: 1, ag: 2, wt: 2, out: 2} + + moved = apply_fsdp_collective_prefetch(gm, mapping, distance=2) + assert moved == 0 + assert mapping[ag] == 2 + + +# --------------------------------------------------------------------------- # +# Coalesced bucketing: N same-submod weight gathers -> ONE +# all_gather_into_tensor_coalesced unpacked with operator.getitem (no cat, no +# split_with_sizes clone), one wait per member. +# --------------------------------------------------------------------------- # +def _build_coalesced_graph(world: int = 4): + """Two weight gathers in submod 2 (a compute submod after a boundary), + each off its own Shard(0) weight placeholder, plus a use of each.""" + from torch._subclasses.fake_tensor import FakeTensorMode + + fake = FakeTensorMode() + g = fx.Graph() + x = g.placeholder("x") + w0 = g.placeholder("l_self_modules_proj0_parameters_weight_") + wa = g.placeholder("l_self_modules_mlp_parameters_weight_") + wb = g.placeholder("l_self_modules_mlp_o_parameters_weight_") + + a = g.call_function(torch.ops.aten.mm.default, (x, w0)) # submod 0 compute + b = g.call_function(torch.ops.aten.relu.default, (a,)) # submod 1 boundary + # submod 2: two weight gathers + waits + uses + ag_a = g.call_function(_ALL_GATHER, (wa, world, "0")) + wt_a = g.call_function(_WAIT, (ag_a,)) + ag_b = g.call_function(_ALL_GATHER, (wb, world, "0")) + wt_b = g.call_function(_WAIT, (ag_b,)) + use_a = g.call_function(torch.ops.aten.mm.default, (b, wt_a)) + out = g.call_function(torch.ops.aten.mm.default, (use_a, wt_b)) + g.output((out,)) + + # Attach meta example_values (the bucket pass reads dtype + shapes). + with fake: + # local shards (chunk, rest) and gathered (world*chunk, rest) + loc_a = torch.empty((8, 16), dtype=torch.bfloat16) + loc_b = torch.empty((16, 8), dtype=torch.bfloat16) + wa.meta["example_value"] = loc_a + wb.meta["example_value"] = loc_b + ag_a.meta["example_value"] = loc_a.new_empty((world * 8, 16)) + ag_b.meta["example_value"] = loc_b.new_empty((world * 16, 8)) + wt_a.meta["example_value"] = ag_a.meta["example_value"] + wt_b.meta["example_value"] = ag_b.meta["example_value"] + + gm = fx.GraphModule(torch.nn.Module(), g) + node_to_subgraph_id = {a: 0, b: 1, ag_a: 2, wt_a: 2, ag_b: 2, wt_b: 2, use_a: 2, out: 2} + nodes = {"a": a, "b": b, "ag_a": ag_a, "ag_b": ag_b, "wt_a": wt_a, "wt_b": wt_b, "use_a": use_a, "out": out} + return gm, node_to_subgraph_id, nodes + + +def _count(gm, target): + return sum(1 for n in gm.graph.nodes if n.op == "call_function" and n.target is target) + + +def test_coalesced_bucket_unpacks_with_getitem(): + gm, mapping, _ = _build_coalesced_graph() + n = bucket_weight_all_gather_coalesced_per_submod(gm, mapping) + assert n == 1, f"expected one coalesced bucket, got {n}" + # exactly one coalesced launch, two getitem, two waits ... + assert _count(gm, _ALL_GATHER_COALESCED) == 1 + assert _count(gm, _WAIT) == 2 + assert sum(1 for x in gm.graph.nodes if x.op == "call_function" and x.target is operator.getitem) == 2 + # ... and NONE of the concat-path machinery (no cat, no split_with_sizes, no + # leftover per-member single all_gather). + assert _count(gm, _CAT) == 0 + assert _count(gm, _SPLIT) == 0 + assert _count(gm, _ALL_GATHER) == 0 + # launch + getitems all tagged / in submod 2; waits in submod 2 too. + coal = next(x for x in gm.graph.nodes if x.target is _ALL_GATHER_COALESCED) + assert coal.meta.get("magi_fsdp_weight_ag_coalesced") is True + assert mapping[coal] == 2 + gm.graph.lint() + + +def test_coalesced_launch_and_getitems_move_together(): + gm, mapping, _ = _build_coalesced_graph() + bucket_weight_all_gather_coalesced_per_submod(gm, mapping) + moved = apply_fsdp_collective_prefetch(gm, mapping, distance=2) + assert moved == 1, f"expected the coalesced launch moved once, got {moved}" + + coal = next(x for x in gm.graph.nodes if x.target is _ALL_GATHER_COALESCED) + gis = [x for x in gm.graph.nodes if x.op == "call_function" and x.target is operator.getitem] + waits = [x for x in gm.graph.nodes if x.op == "call_function" and x.target is _WAIT] + + # launch + both getitems hoisted into submod 0 (the previous compute submod) ... + assert mapping[coal] == 0 + for gi in gis: + assert mapping[gi] == 0, f"getitem must move with the coalesced launch, got {mapping[gi]}" + # ... but the waits stay at the use site in submod 2. + for wt in waits: + assert mapping[wt] == 2, f"wait must stay at consumer, got {mapping[wt]}" + + # The list (coalesced launch) and its getitems never become a submod boundary + # edge: only the getitem outputs cross. Order check: launch precedes submod-0 + # compute. + order = list(gm.graph.nodes) + a = next(x for x in gm.graph.nodes if x.op == "call_function" and x.target is torch.ops.aten.mm.default) + assert order.index(coal) < order.index(a) + gm.graph.lint() + + +if __name__ == "__main__": + test_launch_moves_wait_stays() + test_no_move_when_already_early() + test_non_weight_all_gather_is_ignored() + test_coalesced_bucket_unpacks_with_getitem() + test_coalesced_launch_and_getitems_move_together() + print("PASS: fsdp_collective_prefetch unit tests") From 36b8cdcdc34d88f3b50ed72a4a9b7f13cddff089 Mon Sep 17 00:00:00 2001 From: wtr Date: Thu, 9 Jul 2026 17:22:31 +0800 Subject: [PATCH 2/7] [Feat] Update FSDP Overlap V2 --- magi_compiler/config.py | 98 +-- magi_compiler/magi_backend/magi_backend.py | 239 +++--- magi_compiler/passes/fsdp_overlap/__init__.py | 43 + .../passes/fsdp_overlap/bucket_all_gather.py | 269 +++++++ .../passes/fsdp_overlap/lower_and_bucket.py | 116 +++ .../redistribute_lowering.py} | 36 +- magi_compiler/passes/fsdp_overlap/reorder.py | 527 ++++++++++++ magi_compiler/passes/graph_split/__init__.py | 13 - .../graph_split/fsdp_bucket_all_gather.py | 342 -------- .../graph_split/fsdp_collective_prefetch.py | 367 --------- magi_compiler/profiling/__init__.py | 31 + magi_compiler/profiling/benchmark_inputs.py | 71 ++ magi_compiler/profiling/runtime_estimator.py | 758 ++++++++++++++++++ .../fsdp_overlap_helper/__init__.py | 13 + .../fsdp_overlap_helper/e2e_helper.py | 172 ++++ .../estimator_collective_helper.py | 166 ++++ .../fsdp_overlap_helper/reorder_helper.py | 121 +++ .../test_fsdp_collective_prefetch.py | 229 ------ .../feature_tests/test_fsdp_overlap_bucket.py | 204 +++++ tests/feature_tests/test_fsdp_overlap_e2e.py | 95 +++ .../test_fsdp_overlap_lowering.py | 144 ++++ .../test_fsdp_overlap_reorder.py | 70 ++ .../feature_tests/test_profiling_estimator.py | 373 +++++++++ .../feature_tests/test_profiling_registry.py | 82 ++ 24 files changed, 3464 insertions(+), 1115 deletions(-) create mode 100644 magi_compiler/passes/fsdp_overlap/__init__.py create mode 100644 magi_compiler/passes/fsdp_overlap/bucket_all_gather.py create mode 100644 magi_compiler/passes/fsdp_overlap/lower_and_bucket.py rename magi_compiler/passes/{graph_split/fsdp_redistribute_lowering.py => fsdp_overlap/redistribute_lowering.py} (77%) create mode 100644 magi_compiler/passes/fsdp_overlap/reorder.py delete mode 100644 magi_compiler/passes/graph_split/__init__.py delete mode 100644 magi_compiler/passes/graph_split/fsdp_bucket_all_gather.py delete mode 100644 magi_compiler/passes/graph_split/fsdp_collective_prefetch.py create mode 100644 magi_compiler/profiling/__init__.py create mode 100644 magi_compiler/profiling/benchmark_inputs.py create mode 100644 magi_compiler/profiling/runtime_estimator.py create mode 100644 tests/feature_tests/fsdp_overlap_helper/__init__.py create mode 100644 tests/feature_tests/fsdp_overlap_helper/e2e_helper.py create mode 100644 tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py create mode 100644 tests/feature_tests/fsdp_overlap_helper/reorder_helper.py delete mode 100644 tests/feature_tests/test_fsdp_collective_prefetch.py create mode 100644 tests/feature_tests/test_fsdp_overlap_bucket.py create mode 100644 tests/feature_tests/test_fsdp_overlap_e2e.py create mode 100644 tests/feature_tests/test_fsdp_overlap_lowering.py create mode 100644 tests/feature_tests/test_fsdp_overlap_reorder.py create mode 100644 tests/feature_tests/test_profiling_estimator.py create mode 100644 tests/feature_tests/test_profiling_registry.py diff --git a/magi_compiler/config.py b/magi_compiler/config.py index 725f5cb..7c12027 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -263,81 +263,59 @@ class CompileConfig(BaseSettings): "Each sub-graph between two splitting ops is compiled independently by Inductor." ), ) - enable_fsdp_all_gather_prefetch: bool = Field( + disable_graph_split: bool = Field( False, description=( - "Enable one-submod-ahead prefetch for SimpleFSDP/DTensor weight all-gather. " - "Before split_module, MagiCompiler moves prim_redistribute to the beginning of the previous submod " - "and leaves prim_to_local at the original weight-use point." + "Skip FX-level splitting at the custom subgraph-boundary ops (splitting_ops) and hand the " + "WHOLE graph to Inductor as a single piecewise submodule. The boundary custom ops (e.g. " + "gaga4_mh_moe / attention) stay in the graph as opaque extern calls that Inductor emits " + "verbatim -- there is exactly one compiled submod instead of one per compute region between " + "boundaries. Useful for comparing whole-graph Inductor codegen against the piecewise path. " + "Overrides splitting_ops (they are ignored). Incompatible with cudagraph_mode=PIECEWISE." ), ) - fsdp_all_gather_prefetch_distance: int = Field( - 1, - ge=0, - description="How many submods ahead to launch SimpleFSDP weight all-gather; the initial supported value is 1.", - ) - fsdp_prefetch_mode: str = Field( - "collective", - description=( - "Which FSDP all-gather prefetch pass to use when enable_fsdp_all_gather_prefetch is set. " - "'collective' (default): operate on explicit _c10d_functional.all_gather_into_tensor/wait_tensor " - "nodes, moving ONLY the launch to the previous submod and leaving the wait at the use site " - "(this is the one that actually overlaps). 'redistribute': legacy pass that moves the DTensor " - "prim_redistribute node (kept for reference; does not separate launch from wait). " - "NOTE: the 'collective' FX pass only fires when the pre-split graph already contains explicit " - "all_gather_into_tensor nodes. The real gaga4 SimpleFSDP graph instead contains opaque DTensor " - "prim_redistribute, so that pass is a no-op there; for real overlap on gaga4 use " - "enable_inductor_comm_reorder below (the within-submod Inductor reorder works regardless of " - "how the collectives originate)." - ), - ) - fsdp_lower_redistribute_to_collectives: bool = Field( + enable_fsdp_fullgraph_overlap: bool = Field( False, description=( - "Before the collective prefetch pass, decompose each SimpleFSDP weight DTensor " - "prim_redistribute/prim_to_local into explicit _c10d_functional.all_gather_into_tensor + " - "wait_tensor FX nodes. The real gaga4 graph only has opaque prim_redistribute (no explicit " - "collectives), so this is required for the collective prefetch pass to fire and to let the " - "all_gather launch be moved across the MoE boundary while the wait stays at the use site. " - "Shard(0) only; requires cudagraph_mode=NONE." + "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 " + "(the analytical roofline is unusable -- 0 for custom ops, 60x low for fused pointwise, 1500x high " + "for matmul; see scripts/demo/research_estimate_op_runtime_findings.md)." ), ) - fsdp_bucket_weight_all_gather: bool = Field( - False, + fsdp_fullgraph_bucket_mode: str = Field( + "none", description=( - "After lowering SimpleFSDP weight prim_redistribute to explicit collectives, coalesce each " - "submod's weight all-gathers into one all_gather_into_tensor_coalesced per (submod, group, " - "dtype) -- fewer launched NCCL kernels. Waits are kept separate (one per weight) so the " - "collective prefetch can still hoist the coalesced launch across the MoE boundary while each " - "wait stays at its use site. Requires fsdp_lower_redistribute_to_collectives=True and cudagraph NONE." + "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_bucket_mode: str = Field( - "concat", + fsdp_fullgraph_bucket_size_mib: int = Field( + 0, + ge=0, description=( - "Which bucketing implementation to use when fsdp_bucket_weight_all_gather is set. " - "'concat' (default): flatten each weight's local shard, cat them, do ONE " - "all_gather_into_tensor on the concatenated buffer, then reshape + split_with_sizes back to " - "each weight. The gathered buffer is rank-major while weights are weight-major, so the split " - "is lowered as a clone into a fresh buffer per weight (transient ~2x memory + a multi-GB cat " - "on the compute stream). 'coalesced': fuse the N launches into a single " - "all_gather_into_tensor_coalesced (one NCCL group, one weight-major output buffer per input), " - "then unpack each member with operator.getitem (zero-copy, no cat/clone/2x spike) and keep a " - "separate wait per member. The coalesced op returns a Tensor[] that cannot cross a piecewise " - "submod boundary, so the launch and its getitems are kept together in the launch submod; only " - "the unpacked single tensors cross. Requires fsdp_bucket_weight_all_gather=True and cudagraph NONE." + "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." ), ) - enable_inductor_comm_reorder: bool = Field( - False, + fsdp_overlap_slack_ns: float = Field( + 5000.0, + ge=0.0, description=( - "Enable Inductor's comm/compute overlap reorder for every piecewise submod " - "(reorder_for_compute_comm_overlap=True with passes raise_comms + sink_waits). This hoists each " - "submod's weight all-gather launches to the top of the generated code and sinks each wait to just " - "before its use, so a gather overlaps the compute that precedes its consumer. Works on lowered " - "_c10d_functional collectives inside each submod, so it applies to gaga4 SimpleFSDP regardless of " - "the prefetch pass. The collective prefetch pass also turns this on automatically when it moves a " - "launch across a submod boundary. Requires cudagraph_mode=NONE." + "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." ), ) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 4b72d20..2ea092f 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -34,12 +34,6 @@ from magi_compiler.magi_depyf.timeline import observe_lifecycle, observe_lifecycle_context from magi_compiler.offload.offload_warpper import OffloadWrapper from magi_compiler.passes import CustomJointGraphPartitionFn, FullGraphPassManager, PostGradPassManager, pass_context -from magi_compiler.passes.graph_split import ( - apply_fsdp_collective_prefetch, - bucket_weight_all_gather_coalesced_per_submod, - bucket_weight_all_gather_per_submod, - lower_prim_redistribute_to_collectives, -) from magi_compiler.utils import compilation_counter, compute_code_hash, compute_hash, magi_logger from magi_compiler.utils.visualize import save_fx_graph_visualization @@ -548,40 +542,156 @@ 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 _enable_inductor_comm_reorder(self) -> None: - """Turn on Inductor's comm/compute overlap reorder for piecewise submods. + def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: + """Whole-graph FSDP all-gather / compute overlap (disable_graph_split path). - ``raise_comms`` hoists each submod's collective launches to the top of - the generated code (Inductor otherwise sinks a collective whose result - is only a graph output to just before ``return``, defeating prefetch); - ``sink_waits`` pushes each ``wait_tensor`` down to just before its first - use. Together they let a weight all-gather overlap the compute that runs - before its consumer. Inductor's default is - ``reorder_for_compute_comm_overlap=False``. + 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. """ - self.inductor_compile_config.setdefault("reorder_for_compute_comm_overlap", True) - self.inductor_compile_config.setdefault( - "reorder_for_compute_comm_overlap_passes", ["raise_comms", "sink_waits"] + pass + + from magi_compiler.passes.fsdp_overlap import FsdpOverlapReorder, lower_and_bucket_full_graph + from magi_compiler.profiling import ProfilingRuntimeEstimator + + if self.compile_config.cudagraph_mode != CudaGraphMode.NONE: + raise ValueError("enable_fsdp_fullgraph_overlap requires cudagraph_mode=NONE") + + # The model's true subgraph-boundary ops delimit bucketing regions even + # though disable_graph_split emptied fx_split_ops. + boundary_ops = resolve_defined_ops(self.compile_config.splitting_ops or []) + 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, ) + # Cost model for the reorder's placement sweep. + # + # MULTI-RANK CORRECTNESS (critical): the reorder decides how far to hoist each + # weight all-gather from accumulated per-node costs. If those costs differ + # across ranks, the sweep places a different number of gathers before the eager + # CP all_to_all inside the attention op -> weight-PG vs CP-PG collectives + # interleave in rank-divergent order -> NCCL deadlock (verified on 8-GPU gaga4: + # per-rank profiling -> post-reorder graph `0037` differs across ranks -> hang + # at SeqNum=12). Costs must be rank-IDENTICAL. Three modes (env + # MAGI_COMPILE_FSDP_OVERLAP_COST_MODE = analytical | profile_sync | profile): + # * "analytical" : Inductor roofline (shapes+device only -> rank-identical). + # Deterministic, zero overhead, but less accurate. + # * "profile_sync" : REAL per-op profiling, re-measured in rank-lockstep + # (barrier + fixed iters) and MAX-reduced over gloo so the + # table is identical on every rank -- accurate AND safe. + # Requires the identical-graph precondition (guaranteed by + # the unconditional-pad lowering fix). + # * "profile" : plain per-rank profiling, NO sync. Rank-nondeterministic + # -> WILL deadlock multi-rank; only for world_size==1. + # Default: world_size>1 -> "profile_sync" (accurate + safe); ==1 -> "profile". + # Back-compat: MAGI_COMPILE_FSDP_OVERLAP_ANALYTICAL_COST=1 forces "analytical". + import os as _os + + import torch.distributed as _dist + + # DEFAULT: multi-rank -> "profile_sync"; single-rank -> "profile". + # profile_sync gives REAL measured costs AND a rank-identical schedule, which is + # now SAFE because the model builder replicates uneven-Shard(0) params + # (_replicate_uneven_shard_params in simple_fsdp.py) so the compiled graph is + # structurally identical on every rank -> same structural-key set -> the + # rank-lockstep warm_and_sync max-reduce fully reconciles costs. (An earlier + # attempt without the replicate fix hung even with profile_sync, because trailing + # ranks had extra constant_pad_nd nodes from uneven shards -> divergent placement + # that cost-sync couldn't fix; replicating those tiny params removes the source.) + # "analytical" (Inductor roofline, shapes+device only) is the deadlock-free + # FALLBACK if a model still has per-rank graph divergence -- rank-deterministic + # but less accurate. MAGI_COMPILE_FSDP_OVERLAP_COST_MODE=analytical|profile_sync| + # profile overrides; back-compat MAGI_COMPILE_FSDP_OVERLAP_ANALYTICAL_COST=1 + # forces analytical. + _world = _dist.get_world_size() if (_dist.is_available() and _dist.is_initialized()) else 1 + _mode = _os.environ.get("MAGI_COMPILE_FSDP_OVERLAP_COST_MODE") + if _mode is None: + if _os.environ.get("MAGI_COMPILE_FSDP_OVERLAP_ANALYTICAL_COST") == "1": + _mode = "analytical" + else: + _mode = "profile_sync" if _world > 1 else "profile" + + if _mode == "analytical": + self._fsdp_overlap_estimator = None + cost_fn = None # reorder pass defaults to Inductor's analytical estimate + magi_logger.info("FSDP fullgraph overlap: ANALYTICAL cost (world_size=%d)", _world) + else: + # Pass the estimator DIRECTLY to the reorder pass rather than installing it + # globally at inductor_config.estimate_op_runtime -- the global hook is + # consulted by Inductor's own scheduling/caching in ways that specialize + # dynamic shapes. In "profile_sync" the reorder detects warm_and_sync and + # reconciles costs across ranks; "profile" leaves them per-rank (single + # rank only). We set a flag the reorder reads to decide whether to sync. + self._fsdp_overlap_estimator = ProfilingRuntimeEstimator() + self._fsdp_overlap_estimator._sync_across_ranks = _mode == "profile_sync" + cost_fn = self._fsdp_overlap_estimator + magi_logger.info("FSDP fullgraph overlap: %s cost (world_size=%d)", _mode.upper(), _world) + + # Our reorder pass is the ONLY reorder pass. It moves each weight + # all-gather LAUNCH earlier -- into the upstream compute of the PREVIOUS + # region (a bounded distance~=1 prefetch) -- so the compute between the new + # launch position and the gather's (unmoved) wait hides the collective. + # We deliberately do NOT append Inductor's builtin ``sink_waits``: it defers + # every wait unconditionally, which lets every launch (they depend only on + # the param shards, so all are "ready" at graph start) float to the graph + # top -> all layers' weights resident at once -> OOM on real 8-way FSDP. + # The pass instead moves launches a bounded amount (only far enough that the + # accumulated upstream compute covers comm), so ~2 layers' weights are + # resident at a time. Crucially it ignores the artificial WeakDep ordering + # between collectives (FSDP gathers are independent) when computing each + # launch's earliest-legal position, else the WeakDep would pin the launch + # right after the previous gather and block the move. See + # scripts/demo/research_exposed_allgather_before_attn.md. + reorder = FsdpOverlapReorder(slack_ns=self.compile_config.fsdp_overlap_slack_ns, cost_fn=cost_fn) + 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.5: optionally lower SimpleFSDP weight DTensor prim_redistribute/ - # prim_to_local into explicit all_gather_into_tensor + wait_tensor nodes. - # Must run BEFORE Step 2 builds node_to_subgraph_id (it mutates the graph). - # This is what makes the 'collective' prefetch pass fire on the real gaga4 - # graph (which otherwise only contains opaque prim_redistribute) and lets - # the launch be moved across the MoE boundary while the wait stays put. - if self.compile_config.fsdp_lower_redistribute_to_collectives: - if self.compile_config.cudagraph_mode != CudaGraphMode.NONE: - raise ValueError("FSDP redistribute lowering currently requires cudagraph_mode=NONE") - lower_prim_redistribute_to_collectives(graph) + # 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 @@ -602,75 +712,6 @@ def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[Spli else: node_to_subgraph_id[node] = subgraph_id - # Step 2.4: coalesce each submod's lowered weight all-gathers into one - # all_gather_into_tensor_coalesced per (submod, group, dtype) -- fewer - # launched NCCL kernels. Must run after node_to_subgraph_id is built and - # BEFORE the prefetch pass (so prefetch sees one coalesced launch/submod). - if self.compile_config.fsdp_bucket_weight_all_gather: - if self.compile_config.cudagraph_mode != CudaGraphMode.NONE: - raise ValueError("FSDP all-gather bucketing currently requires cudagraph_mode=NONE") - bucket_mode = self.compile_config.fsdp_bucket_mode - if bucket_mode == "concat": - n_buckets = bucket_weight_all_gather_per_submod(graph, node_to_subgraph_id) - elif bucket_mode == "coalesced": - n_buckets = bucket_weight_all_gather_coalesced_per_submod(graph, node_to_subgraph_id) - else: - raise ValueError( - f"Unknown fsdp_bucket_mode={bucket_mode!r}; expected 'concat' or 'coalesced'" - ) - magi_logger.info("FSDP all-gather bucketing (%s) created %d buckets", bucket_mode, n_buckets) - - # Step 2.5: SimpleFSDP weight all-gather prefetch at submod granularity. - # Goal: launch a submod's weight all-gather during the *previous* submod's - # compute and only wait on it right before the weight is used, so the NCCL - # all-gather overlaps compute. - # - # 'collective' mode (default) operates on the explicit - # _c10d_functional.all_gather_into_tensor / wait_tensor nodes and moves - # ONLY the launch to the previous submod, leaving the wait at the use - # site. This is what actually produces overlap. 'redistribute' mode is - # the legacy pass that moves the opaque DTensor prim_redistribute node; - # it is kept for reference but does NOT separate launch from wait (Inductor - # lowers a single prim_redistribute into both the all_gather and its wait). - if self.compile_config.enable_fsdp_all_gather_prefetch: - if self.compile_config.cudagraph_mode != CudaGraphMode.NONE: - raise ValueError("FSDP all-gather prefetch currently requires cudagraph_mode=NONE") - distance = self.compile_config.fsdp_all_gather_prefetch_distance - if self.compile_config.fsdp_prefetch_mode == "collective": - moved = apply_fsdp_collective_prefetch(graph, node_to_subgraph_id, distance=distance) - magi_logger.info( - "FSDP collective prefetch moved %d all_gather_into_tensor launches by distance=%d", - moved, - distance, - ) - # The prefetch pass puts each hoisted all_gather at the FRONT of - # the previous submod's FX graph, but Inductor's per-submod - # scheduler re-sinks a collective that has no consumer inside the - # submod (its result is only a graph output) back to just before - # `return` -- so at runtime the launch is issued AFTER the - # submod's compute and cannot overlap. Enabling Inductor's - # comm/compute reorder with `raise_comms` hoists the launch back - # to the top of the generated code so it actually runs - # concurrently with this submod's compute on the NCCL stream. - if moved: - self._enable_inductor_comm_reorder() - else: - raise ValueError( - f"Unknown fsdp_prefetch_mode={self.compile_config.fsdp_prefetch_mode!r}; " - "expected 'collective' or 'redistribute'" - ) - - # Standalone Inductor comm/compute overlap reorder. Unlike the FX - # prefetch pass (which only fires when the pre-split graph has explicit - # all_gather_into_tensor nodes), this operates on the lowered collectives - # inside every submod, so it produces overlap on the real gaga4 SimpleFSDP - # graph (whose pre-split form is opaque DTensor prim_redistribute). - if self.compile_config.enable_inductor_comm_reorder: - if self.compile_config.cudagraph_mode != CudaGraphMode.NONE: - raise ValueError("Inductor comm reorder currently requires cudagraph_mode=NONE") - self._enable_inductor_comm_reorder() - magi_logger.info("Enabled Inductor comm/compute overlap reorder (raise_comms + sink_waits)") - # Step 3: split the graph based on node_to_subgraph_id # pytorch might reorder the nodes and the semantics of the graph will change when we have mutations in the graph, if we don't set keep_original_order=True split_gm = torch.fx.passes.split_module.split_module( diff --git a/magi_compiler/passes/fsdp_overlap/__init__.py b/magi_compiler/passes/fsdp_overlap/__init__.py new file mode 100644 index 0000000..c85747e --- /dev/null +++ b/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", +] diff --git a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py new file mode 100644 index 0000000..0c8c962 --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py @@ -0,0 +1,269 @@ +# 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. + +from __future__ import annotations + +import operator +from collections import defaultdict, deque + +import torch +import torch.fx as fx + +from magi_compiler.utils import magi_logger + +_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default +_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default +_WAIT = torch.ops._c10d_functional.wait_tensor.default + + +def _is_param_like(node: fx.Node) -> bool: + """True for a placeholder/get_attr that names a SimpleFSDP weight shard.""" + if node.op not in ("placeholder", "get_attr"): + return False + name = f"{node.name} {node.target}".lower() + return any(t in name for t in ("parameter", "parameters", "weight", "bias", "shard")) + + +def _gathers_a_weight(node: fx.Node) -> bool: + """Walk back from an all_gather through cheap producers (to_local / cast / + pad / view / reshape); return True if the gathered source is a weight shard. + + Used so the bucketing pass also fires on graphs whose weight gathers are NOT + tagged ``magi_fsdp_weight_ag`` (e.g. the demo, which emits explicit + ``all_gather_into_tensor`` directly off a ``*_shard`` parameter).""" + q: deque[fx.Node] = deque(node.all_input_nodes) + seen: set[fx.Node] = set() + while q: + dep = q.popleft() + if dep in seen: + continue + seen.add(dep) + if _is_param_like(dep): + return True + if dep.op == "call_method" and str(dep.target) in {"to_local", "contiguous", "to", "view", "reshape"}: + q.extend(dep.all_input_nodes) + elif dep.op == "call_function": + nm = getattr(dep.target, "__name__", "") or str(dep.target) + if any(t in nm for t in ("constant_pad_nd", "_to_copy", "convert_element_type", "view", "reshape", "clone")): + q.extend(dep.all_input_nodes) + return False + + +def _is_weight_all_gather(node: fx.Node) -> bool: + """A SimpleFSDP weight ``all_gather_into_tensor`` launch. + + Recognized either by the ``magi_fsdp_weight_ag`` tag (set by the redistribute + lowering pass on the real gaga4 graph) OR structurally, when the gathered + source traces back to a weight/param placeholder (covers untagged graphs such + as the demo).""" + if node.op != "call_function" or node.target is not _ALL_GATHER: + return False + if node.meta.get("magi_fsdp_weight_ag"): + return True + return _gathers_a_weight(node) + + +def _local_shard_bytes(ag: fx.Node) -> int: + """Bytes of the local (pre-gather) shard feeding a weight all_gather -- i.e. how + much this rank contributes to the collective. Used to cap coalesced bucket size.""" + loc = ag.args[0] + m = loc.meta.get("example_value") + if m is None: + return 0 + return int(m.numel()) * int(m.element_size()) + + +def _split_region_by_dtype_and_size( + ag_nodes: list[fx.Node], node_index: dict[fx.Node, int], bucket_size_bytes: int +) -> list[list[fx.Node]]: + """Split a region's weight all-gathers into buckets by PROGRAM-ADJACENCY (plan B). + + Walk ``ag_nodes`` in program order and start a NEW bucket whenever: + * the dtype changes (a different-dtype weight physically breaks the run -- so + e.g. bf16, bf16, fp32, bf16 -> {bf16,bf16}, {fp32}, {bf16}); or + * adding the next weight's local-shard bytes would exceed ``bucket_size_bytes`` + (only when the cap is > 0). + + A single weight larger than the cap forms its own bucket (we never split one + weight's all_gather). With ``bucket_size_bytes <= 0`` only the dtype-change rule + applies, so program-adjacent same-dtype weights still coalesce without a size cap. + Caller filters buckets of size < 2 (nothing to coalesce).""" + ordered = sorted(ag_nodes, key=lambda n: node_index[n]) + buckets: list[list[fx.Node]] = [] + cur: list[fx.Node] = [] + cur_dtype = None + cur_bytes = 0 + for ag in ordered: + dt = ag.meta["example_value"].dtype + b = _local_shard_bytes(ag) + breaks_dtype = cur and dt != cur_dtype + breaks_size = cur and bucket_size_bytes > 0 and (cur_bytes + b) > bucket_size_bytes + if breaks_dtype or breaks_size: + buckets.append(cur) + cur, cur_bytes = [], 0 + cur.append(ag) + cur_dtype = dt + cur_bytes += b + if cur: + buckets.append(cur) + return buckets + + +def _producer_chain(node: fx.Node) -> list[fx.Node]: + """The local-shard prep nodes feeding ``node`` (to_local / _to_copy / pad). + These depend only on the weight placeholder, so they may be hoisted. Returns + the movable producers (excludes placeholders/get_attr).""" + chain: list[fx.Node] = [] + seen: set[fx.Node] = set() + stack = [node] + while stack: + n = stack.pop() + for dep in n.all_input_nodes: + if dep in seen: + continue + if dep.op in ("call_function", "call_method"): + seen.add(dep) + chain.append(dep) + stack.append(dep) + return chain + + +def _coalesce_one_bucket( + graph: fx.GraphModule, + node_to_subgraph_id: dict[fx.Node, int], + node_index: dict[fx.Node, int], + sid: int, + ag_nodes: list[fx.Node], +) -> None: + """Merge ONE bucket of same-dtype weight all_gathers into a single + ``all_gather_into_tensor_coalesced`` (launch + per-member getitem + per-member + wait). ``ag_nodes`` must be same (group, dtype); caller guarantees len >= 2.""" + world = int(ag_nodes[0].args[1]) + ag_nodes = sorted(ag_nodes, key=lambda n: node_index[n]) + _, _world, group_name = ag_nodes[0].args + + locals_ = [ag.args[0] for ag in ag_nodes] + ag_metas = [ag.meta["example_value"] for ag in ag_nodes] # (W*chunk_i, *rest_i) + + first_ag = ag_nodes[0] + waits = [next(iter(ag.users)) for ag in ag_nodes] # sole user of each ag is its wait + + # Hoist every member's local-shard prep above the FIRST member's all_gather. + for loc in locals_: + chain = [loc, *_producer_chain(loc)] if loc.op in ("call_function", "call_method") else _producer_chain(loc) + for prod in sorted(chain, key=lambda n: node_index[n]): + first_ag.prepend(prod) + node_to_subgraph_id[prod] = sid + + with graph.graph.inserting_before(first_ag): + coalesced = graph.graph.call_function(_ALL_GATHER_COALESCED, (list(locals_), world, group_name)) + coalesced.meta["example_value"] = list(ag_metas) + coalesced.meta["magi_fsdp_weight_ag"] = True + coalesced.meta["magi_fsdp_weight_ag_coalesced"] = True + node_to_subgraph_id[coalesced] = sid + + outs = [] + for i, am in enumerate(ag_metas): + gi = graph.graph.call_function(operator.getitem, (coalesced, i)) + gi.meta["example_value"] = am + node_to_subgraph_id[gi] = sid + outs.append(gi) + + for out_i, old_wait, am in zip(outs, waits, ag_metas): + with graph.graph.inserting_before(old_wait): + wait_i = graph.graph.call_function(_WAIT, (out_i,)) + wait_i.meta["example_value"] = am + node_to_subgraph_id[wait_i] = node_to_subgraph_id.get(old_wait, sid) + old_wait.replace_all_uses_with(wait_i) + + for ag_old, wait_old in zip(ag_nodes, waits): + node_to_subgraph_id.pop(wait_old, None) + node_to_subgraph_id.pop(ag_old, None) + graph.graph.erase_node(wait_old) + graph.graph.erase_node(ag_old) + + +def bucket_weight_all_gather_coalesced_per_region( + graph: fx.GraphModule, node_to_subgraph_id: dict[fx.Node, int], bucket_size_bytes: int = 0 +) -> int: + """Coalesce, per compute region, the SimpleFSDP weight all-gathers into ONE + ``all_gather_into_tensor_coalesced`` per ``(region_id, group_name, dtype)`` + (or into several buckets when ``bucket_size_bytes > 0``). + + When ``bucket_size_bytes > 0``, a region is further split into multiple buckets: + the weight all_gathers are walked in PROGRAM ORDER and a new bucket starts + whenever the dtype changes or the accumulated local-shard bytes would exceed the + cap (see :func:`_split_region_by_dtype_and_size`). So only same-dtype, + program-adjacent weights within the byte budget coalesce; each bucket is one + ``all_gather_into_tensor_coalesced``. ``bucket_size_bytes == 0`` (default) keeps + the original behavior (one bucket per (region, group, dtype), no size cap). + + This does NOT cat the shards into one buffer. ``all_gather_into_tensor_coalesced`` + fuses the N launches into a single NCCL group while returning one + *weight-major* output buffer per input, so each member is recovered with a + zero-copy ``operator.getitem`` -- no ``cat`` on the compute stream, no + ``split_with_sizes`` clone, and no transient ~2x memory spike. + + For a group of N weight gathers (each a single ``all_gather_into_tensor`` whose + input is the padded local shard ``(chunk_i, *rest_i)``) this builds:: + + # launch side (the coalesced launch + its getitems stay together): + coalesced = all_gather_into_tensor_coalesced([local_0, ..., local_{N-1}], W, group) + out_i = getitem(coalesced, i) # (W*chunk_i, *rest_i), weight-major + # use side (one wait per member, left at its consumer): + wait_i = wait_tensor(out_i) + + ``out_i`` has exactly the shape of the member's old ``all_gather`` output, so + each member's existing downstream (slice / real use) is simply re-pointed from + its old ``wait_tensor`` to ``wait_i``. The ``coalesced`` launch and its + ``getitem`` nodes stay together; the scheduler-level ``FsdpOverlapReorder`` pass + later moves the launch as one unit for compute/comm overlap. + + Runs AFTER redistribute lowering (called from ``lower_and_bucket_full_graph``). + Returns the number of coalesced buckets created. + """ + node_index = {n: i for i, n in enumerate(graph.graph.nodes)} + + # Group by (region, group_name) ONLY -- dtype is NOT part of the key. Program + # order + the dtype-change rule in _split_region_by_dtype_and_size decide the + # actual buckets, so a different-dtype weight physically breaks a same-dtype run + # (plan B / strict program-adjacency). With bucket_size_bytes==0 this reduces to + # one bucket per (region, group, dtype) run. + groups: dict[tuple[int, str], list[fx.Node]] = defaultdict(list) + for node in graph.graph.nodes: + if not _is_weight_all_gather(node): + continue + sid = node_to_subgraph_id.get(node) + if sid is None: + continue + _, _world, group_name = node.args + groups[(sid, group_name)].append(node) + + buckets = 0 + for (sid, group_name), ag_nodes in groups.items(): + for sub in _split_region_by_dtype_and_size(ag_nodes, node_index, bucket_size_bytes): + if len(sub) < 2: + continue # single weight -> keep its own all_gather (nothing to coalesce) + _coalesce_one_bucket(graph, node_to_subgraph_id, node_index, sid, sub) + buckets += 1 + + if buckets: + graph.graph.lint() + graph.recompile() + magi_logger.info( + "FSDP weight all-gather bucketing (coalesced): created %d coalesced buckets across submods " "(bucket_size_bytes=%d)", + buckets, + bucket_size_bytes, + ) + return buckets diff --git a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py new file mode 100644 index 0000000..b13f9ac --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py @@ -0,0 +1,116 @@ +# 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. + +from __future__ import annotations + +"""Whole-graph FSDP weight all-gather lowering + bucketing. + +The piecewise pipeline (see ``magi_backend._split_graph``) lowers SimpleFSDP +weight ``prim_redistribute`` into explicit collectives and then buckets them +*per submod*, keyed by ``node_to_subgraph_id``. Under whole-graph compilation +(``disable_graph_split=True``) there is no split and therefore no +``node_to_subgraph_id`` -- the entire model is a single Inductor graph. + +This module reuses the exact same lowering/bucketing implementations, but drives +them with a **region-numbered** sid map computed the same way ``_split_graph`` +numbers subgraphs (increment at each boundary op) -- WITHOUT actually splitting. +That keeps bucketing granularity per compute-region: the weight gathers used by +one layer's compute (between two boundary ops, e.g. attention / MoE) coalesce +into one collective, but different regions do NOT collapse into a single +model-wide gather (which would force the whole unsharded model resident and OOM, +defeating FSDP). + +The scheduler-level :mod:`magi_compiler.passes.fsdp_overlap.reorder` pass then +places each (possibly coalesced) launch at its latest-safe position. Bucketing +here is intentionally graph-level (not scheduler-level) so the collective node +structure is fixed before Inductor lowering. +""" + +import torch +import torch.fx as fx + +from magi_compiler.utils import magi_logger + +from .bucket_all_gather import bucket_weight_all_gather_coalesced_per_region +from .redistribute_lowering import lower_prim_redistribute_to_collectives + + +def _build_region_sid_map(graph: fx.GraphModule, boundary_ops: list["torch._ops.OpOverload"]) -> dict[fx.Node, int]: + """Number graph nodes into regions delimited by ``boundary_ops``. + + Mirrors ``_split_graph`` Step 2's ``node_to_subgraph_id`` numbering (compute + regions get even ids, boundary ops get their own odd id), but is used only to + group weight all-gathers for bucketing -- the graph is never split. Weights + within the same region share a sid, so bucketing coalesces them; weights in + different regions stay in separate buckets. + """ + resolved = set(boundary_ops) + sid = 0 + mapping: dict[fx.Node, int] = {} + for node in graph.graph.nodes: + if node.op in ("output", "placeholder"): + continue + is_boundary = node.op == "call_function" and ( + node.target in resolved or (hasattr(node.target, "default") and node.target.default in resolved) + ) + if is_boundary: + sid += 1 + mapping[node] = sid + sid += 1 + else: + mapping[node] = sid + return mapping + + +def lower_and_bucket_full_graph( + graph: fx.GraphModule, + bucket_mode: str, + boundary_ops: list["torch._ops.OpOverload"] | None = None, + bucket_size_bytes: int = 0, +) -> int: + """Lower SimpleFSDP weight redistribute -> explicit collectives, then + optionally bucket them per compute-region across the WHOLE graph. + + ``bucket_mode``: + * ``"none"`` -- lowering only (N individual all_gather + N waits). + * ``"coalesced"`` -- per region: one all_gather_into_tensor_coalesced + (ONE launch, N getitems, N waits). + + ``boundary_ops`` are the subgraph-boundary op overloads (the model's + ``splitting_ops``); they delimit bucketing regions. When None/empty, all + same-(group, dtype) gathers fall in one region -- only safe for tiny graphs. + + ``bucket_size_bytes`` (coalesced mode only): when > 0, further split each region + into buckets of at most this many local-shard bytes, breaking at dtype changes + and the byte cap in program order (see + ``bucket_weight_all_gather_coalesced_per_region``). 0 = no cap (one bucket per + region/group/dtype). + + Returns the number of buckets created. + """ + lowered = lower_prim_redistribute_to_collectives(graph) + magi_logger.info("Whole-graph FSDP lowering: %d weight redistribute -> collectives", lowered) + + bucket_mode = (bucket_mode or "none").lower() + if bucket_mode == "none": + return 0 + + sid_map = _build_region_sid_map(graph, boundary_ops or []) + if bucket_mode == "coalesced": + n = bucket_weight_all_gather_coalesced_per_region(graph, sid_map, bucket_size_bytes=bucket_size_bytes) + else: + raise ValueError(f"Unknown bucket_mode={bucket_mode!r}; expected 'none' or 'coalesced'") + + magi_logger.info("Whole-graph FSDP bucketing (%s): created %d buckets", bucket_mode, n) + return n diff --git a/magi_compiler/passes/graph_split/fsdp_redistribute_lowering.py b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py similarity index 77% rename from magi_compiler/passes/graph_split/fsdp_redistribute_lowering.py rename to magi_compiler/passes/fsdp_overlap/redistribute_lowering.py index 4e28ff7..cf419eb 100644 --- a/magi_compiler/passes/graph_split/fsdp_redistribute_lowering.py +++ b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py @@ -1,3 +1,17 @@ +# 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. + from __future__ import annotations import math @@ -126,11 +140,23 @@ def lower_prim_redistribute_to_collectives(graph: fx.GraphModule) -> int: # Pad the local shard up to `chunk` rows when this rank owns fewer # (uneven Shard(0): trailing ranks get remainder/empty). # - # NOTE: meta example_values MUST be created via FakeTensor ops - # (``local.new_empty``), NOT ``torch.empty(..., device=cuda)``. A - # FakeTensor's ``.device`` is a real ``cuda:N``, so ``torch.empty`` - # allocates a REAL full-size buffer per weight; across all weights - # that leaks ~the whole unsharded model (tens of GB) and OOMs. + # NOTE (multi-rank): this `if L < chunk` branch makes trailing ranks' graphs + # STRUCTURALLY DIFFERENT (extra constant_pad_nd nodes) from full-chunk ranks. + # That per-rank structure does NOT by itself break the FSDP-overlap reorder, + # AS LONG AS the reorder's cost model is rank-deterministic (analytical) so it + # places gathers identically regardless of the pad-node count -- verified on + # 8-GPU gaga4: analytical cost -> gather placement identical across ranks + # (4==4 before the 1st attention) and runs clean 4/4, even though rank0 has 0 + # pads and rank4 has 1170. (Emitting the pad unconditionally does NOT help: + # Inductor DCEs the zero-width pad on full-chunk ranks, so the structure + # diverges again post-lowering anyway.) PER-RANK PROFILING is the thing that + # breaks it -- its per-rank costs make the placement diverge (7 vs 5) and + # deadlock; MagiBackend therefore uses analytical cost for world_size > 1. + # + # meta example_values MUST be created via FakeTensor ops (``local.new_empty``), + # NOT ``torch.empty(..., device=cuda)``: a FakeTensor's ``.device`` is a real + # ``cuda:N`` so torch.empty allocates a REAL full-size buffer per weight + # (~the whole unsharded model across all weights) and OOMs. if L < chunk: pad = [0, 0] * (local.dim() - 1) + [0, chunk - L] padded = graph.graph.call_function(_PAD, (cur, pad, 0.0)) diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py new file mode 100644 index 0000000..7a7dc4e --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -0,0 +1,527 @@ +# 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. + +from __future__ import annotations + +"""Latest-safe-launch FSDP all-gather / compute overlap reorder pass. + +Installed into ``torch._inductor.config.reorder_for_compute_comm_overlap_passes`` +as a callable ``list[BaseSchedulerNode] -> list[BaseSchedulerNode]``. Runs on the +WHOLE Inductor graph (MagiCompiler ``disable_graph_split=True``), replacing +PyTorch's builtin ``raise_comms``/``sink_waits``. + +Objective (opposite of ``raise_comms``, which schedules comms as EARLY as +possible): for each FSDP weight ``all_gather`` launch, find its wait / first real +consumer and place the launch at the LATEST position whose downstream compute +still hides the collective, i.e. the latest slot where:: + + sum(compute runtime between launch and first-consumer) >= comm_runtime + slack + +If there is not enough upstream compute, fall back to as-early-as-legal (so we +overlap what we can; never worse than raise_comms). + +TWO-POINTER back-to-front sweep (like the offload HeuristicScheduler's reverse +walk for "latest start of transmission", offload/scheduler.py:319): + * comm pointer -- weight all-gathers in REVERSE original order (last first); + * compute pointer -- a SINGLE index that walks backward CONTINUOUSLY over the + whole graph and is NEVER reset per gather. +Each gather consumes a contiguous run of compute nodes (accumulating their cost) +until it covers its (scaled) comm; its launch target is that stopping point. The +NEXT (earlier) gather resumes the compute pointer from where it stopped -- NOT +from just before its own consumer -- so no two gathers claim the same compute +(this serializes the single NCCL stream) and the collectives keep their original +relative order automatically (targets only decrease). All moves are then applied +in one stable-sort rebuild and validated once (``_validate_full``). + +The Inductor driver (``comms.reorder_compute_and_comm_for_overlap``) does NOT +validate or repair the returned order and does NOT re-sink collectives, so the +returned list MUST be a valid topological order. We guarantee that by only ever +inserting a launch group inside ``[earliest-legal, first-consumer)`` and +asserting the move keeps every producer before / every consumer after the group +(else we abort the move and leave the launch in place). + +Handles two graph-level forms produced by +``fsdp_overlap.lower_and_bucket.lower_and_bucket_full_graph``: +* no-bucket : 1 all_gather -> 1 wait -> 1 consumer; +* coalesced : 1 packed all_gather_into_tensor_coalesced (MultiOutputLayout) + -> N MultiOutput members -> N waits -> N consumers. The launch is + ONE snode with N waits; the packed collective + its N MultiOutput + members move together as one contiguous block, before any wait. +""" + +from collections import defaultdict + +import torch +from torch._inductor.comms import _is_fake_dep +from torch._inductor.ir import MultiOutput +from torch._inductor.scheduler import BaseSchedulerNode +from torch._inductor.utils import contains_collective, contains_wait, is_collective + +from magi_compiler.utils import magi_logger + + +def _debug_enabled() -> bool: + """MagiLogger has no isEnabledFor; check the underlying std logger.""" + import logging + + return logging.getLogger("magi_compiler").isEnabledFor(logging.DEBUG) + + +_AG = torch.ops._c10d_functional.all_gather_into_tensor.default +_AG_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default +_WEIGHT_AG_OPS = (_AG, _AG_COALESCED) + +# Default extra headroom (ns) added to each collective's runtime when sizing the +# compute window, absorbing estimator error + kernel-launch latency so the wait +# rarely stalls. Overridable via the reorder pass constructor. +_DEFAULT_SLACK_NS = 5_000.0 + + +def _leaf_collective_node(snode: BaseSchedulerNode): + """The underlying collective IR node for a (possibly grouped) snode, or None.""" + node = getattr(snode, "node", None) + if node is not None and is_collective(node): + return node + # GroupedSchedulerNode: find the collective child. + for child in getattr(snode, "snodes", []) or []: + cn = getattr(child, "node", None) + if cn is not None and is_collective(cn): + return cn + return None + + +def _is_weight_gather(snode: BaseSchedulerNode) -> bool: + node = _leaf_collective_node(snode) + return node is not None and getattr(node, "op_overload", None) in _WEIGHT_AG_OPS + + +def _is_multi_output(snode: BaseSchedulerNode) -> bool: + node = getattr(snode, "node", None) + return type(node) is MultiOutput + + +class FsdpOverlapReorder: + """Callable reorder pass (see module docstring).""" + + def __init__(self, slack_ns: float = _DEFAULT_SLACK_NS, cost_fn=None) -> None: + self.slack_ns = slack_ns + # cost_fn: snode -> ns. Default uses Inductor's estimate_op_runtime hook, + # which MagiBackend points at the profiling estimator. + if cost_fn is None: + from torch._inductor.comms import estimate_op_runtime + + cost_fn = estimate_op_runtime + self._cost_fn = cost_fn + # Per-call cost cache keyed by snode. Reset at the start of every + # __call__ (snodes are unique per compile). NEVER let it survive into a + # deepcopy: Inductor serializes config.reorder_for_compute_comm_overlap_passes + # into the fx-graph cache key via deepcopy, and snode keys hold FakeTensors + # whose data_ptr access raises. __deepcopy__ below returns a clean instance. + self._cost_cache: dict[BaseSchedulerNode, float] = {} + + def __deepcopy__(self, memo): + # Return a fresh, cache-free instance so config serialization (fx-graph + # cache key) never deepcopies snode/FakeTensor state. cost_fn is shared + # by reference (it is itself deepcopy-safe: see ProfilingRuntimeEstimator). + new = FsdpOverlapReorder.__new__(FsdpOverlapReorder) + new.slack_ns = self.slack_ns + new._cost_fn = self._cost_fn + new._cost_cache = {} + memo[id(self)] = new + return new + + # -- cost ------------------------------------------------------------- + def _cost(self, snode: BaseSchedulerNode) -> float: + c = self._cost_cache.get(snode) + if c is None: + try: + c = max(0.0, float(self._cost_fn(snode))) + except Exception: # noqa: BLE001 + c = 0.0 + self._cost_cache[snode] = c + return c + + @staticmethod + def _is_compute(snode: BaseSchedulerNode) -> bool: + return not contains_collective(snode) and not contains_wait(snode) + + # -- main ------------------------------------------------------------- + def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: + self._cost_cache = {} # fresh per compile; snodes are unique + order = list(snodes) + launches = [s for s in order if _is_weight_gather(s)] + if not launches: + return order + + buf_to_snode = {b: s for s in order for b in s.get_buffer_names()} + op_to_snode: dict[str, BaseSchedulerNode] = {} + for s in order: + for op in s.get_operation_names(): + op_to_snode[op] = s + op_to_snode[s.get_name()] = s + users: dict[str, set] = defaultdict(set) + for s in order: + for d in s.unmet_dependencies: + if not _is_fake_dep(d): + users[d.name].add(s) + + index_of = {s: i for i, s in enumerate(order)} + + # ---- MULTI-RANK DETERMINISM (why this schedule is rank-identical) ---- + # This placement sweep MUST produce an IDENTICAL schedule on every rank, else + # the weight all-gathers (weight PG) interleave with the eager CP all_to_all + # (CP PG, inside the attention boundary op) in a rank-divergent order -> + # cross-PG NCCL collective-order mismatch -> deadlock (verified via flight + # recorder on 8-GPU gaga4: hang at SeqNum=12, rank0 races weight-PG ahead + # while peers block on the CP all_to_all). Rank-consistency needs BOTH: + # 1. IDENTICAL INPUT GRAPH per rank -- guaranteed by the FSDP redistribute + # lowering now emitting 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 (different snode-list length -> + # divergent placement). Verified: stage `0008` byte-identical across ranks. + # 2. RANK-DETERMINISTIC COSTS -- the sweep accumulates `_cost` to decide how + # far to hoist each launch. Profiling benchmarks PER RANK (timing noise + + # replays the attention op's internal CP all_to_all per rank), so with + # identical graphs the post-reorder `0037` STILL diverged (rank0=7 vs + # rank4=5 gathers before the 1st attention) and hung. Two ways to make + # costs rank-identical, both supported (MagiBackend selects via cost_fn): + # (a) ANALYTICAL cost_fn (roofline, shapes+device only) -- deterministic + # by construction, zero overhead, but less accurate; + # (b) SYNCHRONIZED PROFILING -- the estimator exposes `warm_and_sync`, + # which re-measures every op in rank-lockstep (barrier + fixed iters) + # and MAX-reduces over gloo, giving REAL costs that are still + # identical on every rank. Safe now that the graph is identical + # (same key set on all ranks -> symmetric barriers/all_gather). + # Below: if the cost_fn supports warm_and_sync, warm the table on all + # compute nodes then sync it; on failure, bail (leave graph unchanged = + # overlap off, no hang). Analytical cost_fn has no warm_and_sync -> skip. + if hasattr(self._cost_fn, "warm_and_sync") and getattr(self._cost_fn, "_sync_across_ranks", False): + try: + for s in order: + if self._is_compute(s) or contains_collective(s): + self._cost(s) # populate estimator._table (per structural key) + n_changed = self._cost_fn.warm_and_sync() + self._cost_cache = {} # drop pass-local cache -> re-read synced costs + magi_logger.info( + "FSDP overlap reorder: rank-synchronized profiling done (%d cost entries reconciled)", n_changed + ) + except Exception as exc: # noqa: BLE001 + magi_logger.warning("FSDP overlap reorder: synchronized profiling failed (%s); leaving graph unchanged", exc) + return order + + # ---- TWO-POINTER back-to-front sweep (like offload HeuristicScheduler) ---- + # comm pointer: weight all-gathers in REVERSE original order (last first). + # compute pointer: a SINGLE index that walks backward CONTINUOUSLY across the + # whole graph and is NEVER reset per gather -- each gather consumes a + # contiguous run of compute nodes to cover its (scaled) comm, and the next + # (earlier) gather resumes from where the pointer stopped (NOT from just + # before the consumer again). This (a) serializes the single NCCL stream -- + # no two gathers claim the same compute -- and (b) keeps collectives in their + # original relative order automatically (targets only decrease). Mirrors + # offload/scheduler.py's reverse `i` walk for "latest start of transmission". + launches_in_order = sorted(launches, key=lambda s: index_of[s]) # original program order + + # Per-gather static facts (on the STABLE original order). + plans = [] # (launch, group, fc_idx, comm_runtime, lower) + for launch in launches_in_order: + group = self._launch_group(launch, order, buf_to_snode, users) + fc_idx = self._first_consumer_index(launch, group, order, users) + if fc_idx is None: + continue + comm_runtime = self._cost(launch) + lower = self._earliest_legal_index(group, order, index_of, buf_to_snode, op_to_snode) + plans.append((launch, group, fc_idx, comm_runtime, lower)) + + # ---- Sweep the compute pointer backward (UPSTREAM-of-launch scan) ---- + # For each gather (LATEST first) accumulate the compute that sits IMMEDIATELY + # BEFORE its launch, walking backward, until the accumulated compute covers + # comm; move the launch to just before that compute block. Because this pass + # moves ONLY the launch (never the wait) and lowering emits ``all_gather; + # wait`` back-to-back (wait at launch+1), the compute we walk past ends up in + # the [new_launch, wait] window and therefore actually overlaps the gather. + # Scanning UPSTREAM (from the launch) -- not downstream from the consumer -- + # is the fix: the old downstream scan counted compute that runs AFTER the + # (unmoved) wait and can never overlap, so single-weight gathers were judged + # "covered" and left glued to their wait -> fully exposed. + # + # A SINGLE compute pointer walks backward CONTINUOUSLY and is NEVER reset per + # gather (each gather resumes from where the previous, later gather stopped), + # so no two gathers claim the same compute -- this serializes the single NCCL + # stream and keeps collectives in their original relative order (targets only + # decrease). A gather with no upstream compute before `lower` (e.g. the + # graph's first gather) can't move and stays where it is (structural bubble). + targets: dict = {} # launch -> target index (in original order space) + compute_idx = len(order) # scan compute strictly below this + for launch, group, fc_idx, comm_runtime, lower in reversed(plans): + cur = index_of[launch] + # Start just before the launch, but no later than where the previous + # (later) gather already consumed compute down to. + compute_idx = min(compute_idx, cur) + need = comm_runtime + self.slack_ns + acc = 0.0 + t = compute_idx + while t > lower: + s = order[t - 1] + if self._is_compute(s): + acc += self._cost(s) + t -= 1 + if acc >= need: + break + # `t` is the earliest position whose downstream compute (up to the launch) + # covers comm. target < cur moves the launch earlier; target == cur means + # no upstream compute was available (graph head / previous gather took it) + # -> leave in place. target >= lower always (producers stay before it). + target = max(lower, t) + targets[launch] = (target, group) + compute_idx = target # next (earlier) gather resumes from actual placement + if _debug_enabled(): + magi_logger.debug( + "FSDP overlap: launch cur=%d -> target=%d fc=%d lower=%d | comm=%.1fus " "acc_upstream=%.1fus %s", + cur, + target, + fc_idx, + lower, + comm_runtime / 1e3, + acc / 1e3, + "hidden" if acc >= need else "COMPUTE-LIMITED", + ) + + # ---- CROSS-RANK DETERMINISM: enforce weight-gather relative order ---- + # CRITICAL for multi-rank correctness. All FSDP weight all-gathers run on + # the SAME process group, so NCCL matches the Nth all-gather call across + # ranks positionally -- the ranks MUST issue them in an identical relative + # order or they deadlock (observed: "Watchdog caught collective operation + # timeout ... SeqNum=7 COALESCED"). The reverse sweep's ``target = + # max(lower, t)`` can INVERT two gathers' order: a later gather may move far + # up (small target, derived from PER-RANK profiled compute costs), while an + # earlier gather is pinned by its real-dep floor ``lower`` to a LARGER target + # -> the earlier gather sorts after the later one. Because the targets + # depend on per-rank kernel timings (ProfilingRuntimeEstimator benchmarks on + # each rank independently), whether the inversion happens differs per rank -> + # divergent collective order -> hang. Fix: clamp targets to be + # NON-DECREASING in original program index (a graph property, identical on + # every rank). Walk launches in original order with a running max; a gather + # can never be placed before an earlier-issued gather. This preserves the + # original weight-gather subsequence on all ranks regardless of cost jitter, + # at the cost of occasionally not moving a launch as early as its compute + # window would allow (it clusters just after the prior gather instead). + running = -1 + for launch in launches_in_order: + if launch not in targets: + continue + target, group = targets[launch] + if target < running: + target = running + targets[launch] = (target, group) + running = target + + # ---- apply all moves in ONE rebuild (targets are in the ORIGINAL index + # space; applying moves incrementally would shift those indices). Assign a + # sort key per node: non-group nodes keep their original index; each launch + # group is inserted just before the node originally at `target` (key + # target-0.5), members ordered among themselves by original index. A stable + # sort then realizes every move at once while preserving all other nodes' + # relative order and each group's internal order. Collective order is kept + # because targets are monotonic (the two-pointer only decreases). ---- + group_members: dict = {} + for launch, (target, group) in targets.items(): + for m in group: + group_members[m] = target + moved = sum(1 for launch, (target, _g) in targets.items() if index_of[launch] != target) + + def _key(s): + if s in group_members: + return (group_members[s] - 0.5, index_of[s]) + return (index_of[s], 0.0) + + new_order = sorted(order, key=_key) + # Validate the rebuilt order is a valid topological order; only commit if so. + if self._validate_full(new_order, op_to_snode, buf_to_snode, users): + order[:] = new_order + else: + magi_logger.warning("FSDP overlap reorder: rebuilt order failed validation; leaving graph unchanged") + moved = 0 + + measured = getattr(self._cost_fn, "n_measured", None) + cache_hits = getattr(self._cost_fn, "n_cache_hits", None) + n_distinct = len(getattr(self._cost_fn, "_table", {}) or {}) + magi_logger.info( + "FSDP overlap reorder: repositioned %d/%d weight all-gather launch(es) " + "(cost table: %d distinct ops, measured=%s reused=%s)", + moved, + len(launches), + n_distinct, + measured, + cache_hits, + ) + # Full op->time table at DEBUG, or to a file when MAGI_COMPILE_FSDP_OVERLAP_DUMP + # is set (so the estimate table can be captured WITHOUT global DEBUG logging, + # which can trip torch's PT2_COMPILE chromium-event assertion on some builds). + if hasattr(self._cost_fn, "summary"): + import os as _os + + if _debug_enabled(): + magi_logger.debug("FSDP overlap %s", self._cost_fn.summary()) + dump = _os.environ.get("MAGI_COMPILE_FSDP_OVERLAP_DUMP") + if dump: + try: + rank = _os.environ.get("RANK", "0") + with open(f"{dump}.rank{rank}", "a") as f: + f.write(self._cost_fn.summary() + "\n") + except Exception as exc: # noqa: BLE001 + magi_logger.warning("FSDP overlap: could not write cost dump: %s", exc) + return order + + # -- group detection -------------------------------------------------- + def _launch_group(self, launch, order, buf_to_snode, users) -> list[BaseSchedulerNode]: + """The snodes that must move together with the launch. + + Coalesced: packed collective + its MultiOutput members (they depend on the + packed buffer and must stay immediately after it, before any wait). + no-bucket: just the launch (the wait stays put). + """ + group = [launch] + node = _leaf_collective_node(launch) + if node is not None and getattr(node, "op_overload", None) is _AG_COALESCED: + produced = set(launch.get_buffer_names()) + for s in order: + if _is_multi_output(s) and any((not _is_fake_dep(d)) and d.name in produced for d in s.unmet_dependencies): + group.append(s) + return group + + # -- consumer discovery ---------------------------------------------- + def _wait_snodes(self, group, order, users) -> list[BaseSchedulerNode]: + produced: set[str] = set() + for s in group: + produced |= set(s.get_buffer_names()) + waits = [] + seen = set() + for b in produced: + for u in users.get(b, ()): # readers of the launch/member buffers + if u in seen: + continue + seen.add(u) + if contains_wait(u): + waits.append(u) + return waits + + def _first_consumer_index(self, launch, group, order, users) -> int | None: + """min over all waits of the earliest real (non-transparent) consumer index.""" + index_of = {s: i for i, s in enumerate(order)} + waits = self._wait_snodes(group, order, users) + if not waits: + return None + best = None + for w in waits: + fc = self._first_real_consumer_index(w, index_of, users) + if fc is not None: + best = fc if best is None else min(best, fc) + return best + + def _first_real_consumer_index(self, wait, index_of, users) -> int | None: + """Forward BFS from a wait through transparent forwarders (cost~0 view / + getitem / MultiOutput / split) to the first genuine compute consumer.""" + seen = set() + stack = list(wait.get_buffer_names()) + best = None + while stack: + b = stack.pop() + for u in users.get(b, ()): + if u in seen: + continue + seen.add(u) + if self._is_transparent(u): + stack.extend(u.get_buffer_names()) + else: + idx = index_of.get(u) + if idx is not None: + best = idx if best is None else min(best, idx) + return best + + def _is_transparent(self, snode: BaseSchedulerNode) -> bool: + """A forwarder that doesn't count as the weight's real use: waits, + MultiOutput unpacks, and ~zero-cost view/reshape/getitem kernels.""" + if contains_wait(snode) or _is_multi_output(snode): + return True + # Treat vanishingly cheap nodes (views, getitems, splits) as transparent. + return self._cost(snode) <= 1.0 + + # -- repositioning ---------------------------------------------------- + def _earliest_legal_index(self, group, order, index_of, buf_to_snode, op_to_snode) -> int: + """1 + max index of any REAL (data-dependency) producer the group needs. + + Uses ONLY non-fake buffer dependencies, walked transitively. We must NOT + use ``snode.ancestors``: that transitive-closure set is polluted by the + artificial ``WeakDep`` edges Inductor inserts between collectives to + serialize them onto one comm stream (verified: a single-weight qkv/o gather + lists the PREVIOUS layer's coalesced all-gather as an ancestor via a + ``WeakDep`` on its output buffer -- but it does NOT read that buffer). FSDP + weight all-gathers gather INDEPENDENT param shards; there is no real + gather->gather data dependency. Counting the WeakDep pinned ``lower`` right + after the previous collective, which wrongly forbade moving the launch up + past the (real, hideable) compute between the two gathers -- the exact reason + the attention gathers stayed exposed. A gather's only real producer is its + own weight-shard placeholder (+ any to_local/pad/cast chain), so real + ``lower`` is ~0; the launch is then free to move up into earlier compute. + """ + group_set = set(group) + lo = 0 + for s in group: + for d in s.unmet_dependencies: # buffer names + if _is_fake_dep(d): # WeakDep / StarDep -- ordering hint, not data + continue + prod = buf_to_snode.get(d.name) + if prod is None or prod in group_set: + continue + lo = max(lo, index_of.get(prod, 0) + 1) + return lo + + def _validate_full(self, new_order, op_to_snode, buf_to_snode, users) -> bool: + """Check the rebuilt order is a valid topological order w.r.t. REAL data + dependencies: every node's non-fake buffer producers precede it (the + Inductor driver does NOT repair the order, so a real-dep violation would + silently miscompile). O(nodes * deps). + + We deliberately do NOT validate against ``snode.ancestors``: that set is the + transitive closure of ALL deps including the artificial ``WeakDep`` ordering + edges Inductor inserts between collectives (see ``_earliest_legal_index``). + Our pass intentionally reorders a weight all-gather across such a WeakDep + (there is no real gather->gather data dependency), so an ancestors-based + check would false-reject that legal move and no-op the whole reorder. The + per-node real-buffer-dep check below is itself a complete topological + validation over the real data-dependency DAG: if every node's direct real + producers precede it, the order is a valid real-dep topological order. + WeakDep ordering is advisory (stream serialization / memory) and is NOT a + correctness constraint, so violating it is safe.""" + pos = {s: i for i, s in enumerate(new_order)} + for s in new_order: + sp = pos[s] + for d in s.unmet_dependencies: # buffer names + if _is_fake_dep(d): # WeakDep / StarDep -- advisory ordering, not data + continue + prod = buf_to_snode.get(d.name) + if prod is s: # fused snode may name its own internal buffers + continue + if prod is not None and pos.get(prod, -1) >= sp: + if _debug_enabled(): + magi_logger.debug( + "validate fail: %s@%d needs buffer-dep %s@%d (buf %s)", + s.get_name(), + sp, + prod.get_name(), + pos.get(prod, -1), + d.name, + ) + return False + return True diff --git a/magi_compiler/passes/graph_split/__init__.py b/magi_compiler/passes/graph_split/__init__.py deleted file mode 100644 index e5f74d2..0000000 --- a/magi_compiler/passes/graph_split/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from .fsdp_bucket_all_gather import ( - bucket_weight_all_gather_coalesced_per_submod, - bucket_weight_all_gather_per_submod, -) -from .fsdp_collective_prefetch import apply_fsdp_collective_prefetch -from .fsdp_redistribute_lowering import lower_prim_redistribute_to_collectives - -__all__ = [ - "apply_fsdp_collective_prefetch", - "bucket_weight_all_gather_per_submod", - "bucket_weight_all_gather_coalesced_per_submod", - "lower_prim_redistribute_to_collectives", -] \ No newline at end of file diff --git a/magi_compiler/passes/graph_split/fsdp_bucket_all_gather.py b/magi_compiler/passes/graph_split/fsdp_bucket_all_gather.py deleted file mode 100644 index cbc8d50..0000000 --- a/magi_compiler/passes/graph_split/fsdp_bucket_all_gather.py +++ /dev/null @@ -1,342 +0,0 @@ -from __future__ import annotations - -import math -import operator -from collections import defaultdict, deque - -import torch -import torch.fx as fx - -from magi_compiler.utils import magi_logger - -_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default -_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default -_WAIT = torch.ops._c10d_functional.wait_tensor.default -_CAT = torch.ops.aten.cat.default -_RESHAPE = torch.ops.aten.reshape.default -_SPLIT = torch.ops.aten.split_with_sizes.default - - -def _is_param_like(node: fx.Node) -> bool: - """True for a placeholder/get_attr that names a SimpleFSDP weight shard.""" - if node.op not in ("placeholder", "get_attr"): - return False - name = f"{node.name} {node.target}".lower() - return any(t in name for t in ("parameter", "parameters", "weight", "bias", "shard")) - - -def _gathers_a_weight(node: fx.Node) -> bool: - """Walk back from an all_gather through cheap producers (to_local / cast / - pad / view / reshape); return True if the gathered source is a weight shard. - - Used so the bucketing pass also fires on graphs whose weight gathers are NOT - tagged ``magi_fsdp_weight_ag`` (e.g. the demo, which emits explicit - ``all_gather_into_tensor`` directly off a ``*_shard`` parameter).""" - q: deque[fx.Node] = deque(node.all_input_nodes) - seen: set[fx.Node] = set() - while q: - dep = q.popleft() - if dep in seen: - continue - seen.add(dep) - if _is_param_like(dep): - return True - if dep.op == "call_method" and str(dep.target) in {"to_local", "contiguous", "to", "view", "reshape"}: - q.extend(dep.all_input_nodes) - elif dep.op == "call_function": - nm = getattr(dep.target, "__name__", "") or str(dep.target) - if any(t in nm for t in ("constant_pad_nd", "_to_copy", "convert_element_type", "view", "reshape", "clone")): - q.extend(dep.all_input_nodes) - return False - - -def _is_weight_all_gather(node: fx.Node) -> bool: - """A SimpleFSDP weight ``all_gather_into_tensor`` launch. - - Recognized either by the ``magi_fsdp_weight_ag`` tag (set by the redistribute - lowering pass on the real gaga4 graph) OR structurally, when the gathered - source traces back to a weight/param placeholder (covers untagged graphs such - as the demo).""" - if node.op != "call_function" or node.target is not _ALL_GATHER: - return False - if node.meta.get("magi_fsdp_weight_ag"): - return True - return _gathers_a_weight(node) - - -def _producer_chain(node: fx.Node) -> list[fx.Node]: - """The local-shard prep nodes feeding ``node`` (to_local / _to_copy / pad). - These depend only on the weight placeholder, so they may be hoisted. Returns - the movable producers (excludes placeholders/get_attr).""" - chain: list[fx.Node] = [] - seen: set[fx.Node] = set() - stack = [node] - while stack: - n = stack.pop() - for dep in n.all_input_nodes: - if dep in seen: - continue - if dep.op in ("call_function", "call_method"): - seen.add(dep) - chain.append(dep) - stack.append(dep) - return chain - - -def bucket_weight_all_gather_per_submod( - graph: fx.GraphModule, - node_to_subgraph_id: dict[fx.Node, int], -) -> int: - """Coalesce, per submod, the SimpleFSDP weight all-gathers into ONE - ``all_gather_into_tensor`` per ``(subgraph_id, group_name, dtype)`` using the - torch.compile-style flatten/cat/gather/split merge (single collective, single - wait), instead of N individual gathers. - - For a group of N weight gathers (each a single ``all_gather_into_tensor`` from - the redistribute-lowering pass, tagged ``magi_fsdp_weight_ag``, input = padded - local shard of shape ``(chunk_i, *rest_i)``), this builds:: - - # launch side (movable by the prefetch pass): - flat_i = reshape(local_i, [-1]) for each member - cat_in = cat([flat_0, ..., flat_{N-1}]) # (sum_numel,) - ag = all_gather_into_tensor(cat_in, W, group) # (W*sum_numel,) - # use side (wait stays before the consuming compute): - waited = wait_tensor(ag) # ONE wait - resh = reshape(waited, [W, sum_numel]) - splits = split_with_sizes(resh, [numel_0, ...], dim=1) - out_i = reshape(splits[i], [W*chunk_i, *rest_i]) - - ``out_i`` has exactly the shape of the member's old ``all_gather`` output, so - each member's existing downstream (the optional ``slice`` back to the true - full size, then the real use) is simply re-pointed from its old - ``wait_tensor`` to ``out_i``. The old per-member ``all_gather`` and - ``wait_tensor`` are erased. - - Only the single ``ag`` tensor crosses submod boundaries (the ``split_with_sizes`` - list and its getitems live entirely on the use side), so this avoids the AOT - output-spec mismatch that the list-returning ``all_gather_into_tensor_coalesced`` - op triggers under piecewise split. - - Runs AFTER redistribute lowering and BEFORE the collective prefetch pass. - Returns the number of coalesced buckets created. - """ - node_index = {n: i for i, n in enumerate(graph.graph.nodes)} - - groups: dict[tuple[int, str, torch.dtype], list[fx.Node]] = defaultdict(list) - for node in graph.graph.nodes: - if not _is_weight_all_gather(node): - continue - sid = node_to_subgraph_id.get(node) - if sid is None: - continue - _, _world, group_name = node.args - dtype = node.meta["example_value"].dtype - groups[(sid, group_name, dtype)].append(node) - - buckets = 0 - for (sid, group_name, dtype), ag_nodes in groups.items(): - if len(ag_nodes) < 2: - continue # nothing to coalesce - - world = int(ag_nodes[0].args[1]) - ag_nodes.sort(key=lambda n: node_index[n]) - - # Per-member: the padded local shard (ag input) and its full gathered meta. - locals_ = [ag.args[0] for ag in ag_nodes] - ag_metas = [ag.meta["example_value"] for ag in ag_nodes] # (W*chunk_i, *rest_i) - local_metas = [loc.meta["example_value"] for loc in locals_] # (chunk_i, *rest_i) - numels = [int(math.prod(m.shape)) for m in local_metas] - sum_numel = sum(numels) - dev = local_metas[0].device - - first_ag = ag_nodes[0] - # Sole user of each member's all_gather is its wait_tensor. - waits = [next(iter(ag.users)) for ag in ag_nodes] - - # Hoist every member's local-shard prep (to_local / _to_copy / pad chain) - # above the FIRST member's all_gather, so all `locals_` are defined before - # the coalesced launch we insert there. These chains depend only on the - # weight placeholder, so they can always move up. - for loc in locals_: - chain = [loc, *_producer_chain(loc)] if loc.op in ("call_function", "call_method") else _producer_chain(loc) - for prod in sorted(chain, key=lambda n: node_index[n]): - first_ag.prepend(prod) - node_to_subgraph_id[prod] = sid - - # ---- launch side: flatten each local, concat, single all_gather ---- - # Insert just before the FIRST member's all_gather (all locals now precede it). - with graph.graph.inserting_before(first_ag): - flats = [] - for loc, lm, n in zip(locals_, local_metas, numels): - fl = graph.graph.call_function(_RESHAPE, (loc, [-1])) - fl.meta["example_value"] = lm.reshape(-1) - node_to_subgraph_id[fl] = sid - flats.append(fl) - cat_in = graph.graph.call_function(_CAT, (flats, 0)) - cat_in.meta["example_value"] = local_metas[0].new_empty((sum_numel,)) - node_to_subgraph_id[cat_in] = sid - - ag = graph.graph.call_function(_ALL_GATHER, (cat_in, world, group_name)) - ag.meta["example_value"] = local_metas[0].new_empty((world * sum_numel,)) - ag.meta["magi_fsdp_weight_ag"] = True # so the prefetch pass moves it - node_to_subgraph_id[ag] = sid - - # ---- use side: one wait, reshape, split back to per-member tensors ---- - # Insert before the EARLIEST member's wait (its use site within this submod). - first_wait = min(waits, key=lambda n: node_index[n]) - with graph.graph.inserting_before(first_wait): - waited = graph.graph.call_function(_WAIT, (ag,)) - waited.meta["example_value"] = ag.meta["example_value"] - node_to_subgraph_id[waited] = sid - - resh = graph.graph.call_function(_RESHAPE, (waited, [world, sum_numel])) - resh.meta["example_value"] = local_metas[0].new_empty((world, sum_numel)) - node_to_subgraph_id[resh] = sid - - split = graph.graph.call_function(_SPLIT, (resh, numels, 1)) - split.meta["example_value"] = [local_metas[0].new_empty((world, n)) for n in numels] - node_to_subgraph_id[split] = sid - - for i, (am, lm) in enumerate(zip(ag_metas, local_metas)): - gi = graph.graph.call_function(operator.getitem, (split, i)) - gi.meta["example_value"] = local_metas[0].new_empty((world, numels[i])) - node_to_subgraph_id[gi] = sid - out_i = graph.graph.call_function(_RESHAPE, (gi, list(am.shape))) - out_i.meta["example_value"] = am # (W*chunk_i, *rest_i): same as old ag out - node_to_subgraph_id[out_i] = sid - - # Re-point the member's downstream (slice / real use) from its old - # wait to out_i, then drop the old wait + all_gather. - waits[i].replace_all_uses_with(out_i) - - for ag_old, wait_old in zip(ag_nodes, waits): - node_to_subgraph_id.pop(wait_old, None) - node_to_subgraph_id.pop(ag_old, None) - graph.graph.erase_node(wait_old) - graph.graph.erase_node(ag_old) - - buckets += 1 - - if buckets: - graph.graph.lint() - graph.recompile() - magi_logger.info( - "FSDP weight all-gather bucketing (concat): created %d coalesced buckets across submods", buckets - ) - return buckets - - -def bucket_weight_all_gather_coalesced_per_submod( - graph: fx.GraphModule, - node_to_subgraph_id: dict[fx.Node, int], -) -> int: - """Coalesce, per submod, the SimpleFSDP weight all-gathers into ONE - ``all_gather_into_tensor_coalesced`` per ``(subgraph_id, group_name, dtype)``. - - Unlike :func:`bucket_weight_all_gather_per_submod` (the ``concat`` strategy), - this does NOT cat the shards into one buffer. ``all_gather_into_tensor_coalesced`` - fuses the N launches into a single NCCL group while returning one - *weight-major* output buffer per input, so each member is recovered with a - zero-copy ``operator.getitem`` -- no ``cat`` on the compute stream, no - ``split_with_sizes`` clone, and no transient ~2x memory spike that the concat - path incurs. - - For a group of N weight gathers (each a single ``all_gather_into_tensor`` whose - input is the padded local shard ``(chunk_i, *rest_i)``) this builds:: - - # launch side (kept together in ONE submod -- the list must not cross a - # piecewise split boundary): - coalesced = all_gather_into_tensor_coalesced([local_0, ..., local_{N-1}], W, group) - out_i = getitem(coalesced, i) # (W*chunk_i, *rest_i), weight-major - # use side (one wait per member, left at its consumer): - wait_i = wait_tensor(out_i) - - ``out_i`` has exactly the shape of the member's old ``all_gather`` output, so - each member's existing downstream (slice / real use) is simply re-pointed from - its old ``wait_tensor`` to ``wait_i``. The ``coalesced`` launch and all its - ``getitem`` nodes stay in the launch submod; only the per-member ``out_i`` - single tensors cross to the consumer submod. The prefetch pass moves the - launch together with its getitems (see ``apply_fsdp_collective_prefetch``). - - Runs AFTER redistribute lowering and BEFORE the collective prefetch pass. - Returns the number of coalesced buckets created. - """ - node_index = {n: i for i, n in enumerate(graph.graph.nodes)} - - groups: dict[tuple[int, str, torch.dtype], list[fx.Node]] = defaultdict(list) - for node in graph.graph.nodes: - if not _is_weight_all_gather(node): - continue - sid = node_to_subgraph_id.get(node) - if sid is None: - continue - _, _world, group_name = node.args - dtype = node.meta["example_value"].dtype - groups[(sid, group_name, dtype)].append(node) - - buckets = 0 - for (sid, group_name, dtype), ag_nodes in groups.items(): - if len(ag_nodes) < 2: - continue # nothing to coalesce - - world = int(ag_nodes[0].args[1]) - ag_nodes.sort(key=lambda n: node_index[n]) - - locals_ = [ag.args[0] for ag in ag_nodes] - ag_metas = [ag.meta["example_value"] for ag in ag_nodes] # (W*chunk_i, *rest_i) - - first_ag = ag_nodes[0] - # Sole user of each member's all_gather is its wait_tensor. - waits = [next(iter(ag.users)) for ag in ag_nodes] - - # Hoist every member's local-shard prep (to_local / _to_copy / pad chain) - # above the FIRST member's all_gather so all `locals_` precede the - # coalesced launch we insert there. These depend only on the weight - # placeholder, so they can always move up. - for loc in locals_: - chain = [loc, *_producer_chain(loc)] if loc.op in ("call_function", "call_method") else _producer_chain(loc) - for prod in sorted(chain, key=lambda n: node_index[n]): - first_ag.prepend(prod) - node_to_subgraph_id[prod] = sid - - # ---- launch side: ONE coalesced all_gather + per-member getitem unpack ---- - # Insert before the FIRST member's all_gather (all locals now precede it). - with graph.graph.inserting_before(first_ag): - coalesced = graph.graph.call_function(_ALL_GATHER_COALESCED, (list(locals_), world, group_name)) - # The coalesced op returns a list[Tensor]; meta mirrors the per-member - # gathered shapes (weight-major, same as each old all_gather output). - coalesced.meta["example_value"] = list(ag_metas) - coalesced.meta["magi_fsdp_weight_ag"] = True # so prefetch moves it - coalesced.meta["magi_fsdp_weight_ag_coalesced"] = True # move getitems too - node_to_subgraph_id[coalesced] = sid - - outs = [] - for i, am in enumerate(ag_metas): - gi = graph.graph.call_function(operator.getitem, (coalesced, i)) - gi.meta["example_value"] = am # (W*chunk_i, *rest_i) - node_to_subgraph_id[gi] = sid - outs.append(gi) - - # ---- use side: one wait per member, left before its consumer ---- - for i, (out_i, old_wait) in enumerate(zip(outs, waits)): - with graph.graph.inserting_before(old_wait): - wait_i = graph.graph.call_function(_WAIT, (out_i,)) - wait_i.meta["example_value"] = ag_metas[i] - node_to_subgraph_id[wait_i] = node_to_subgraph_id.get(old_wait, sid) - old_wait.replace_all_uses_with(wait_i) - - for ag_old, wait_old in zip(ag_nodes, waits): - node_to_subgraph_id.pop(wait_old, None) - node_to_subgraph_id.pop(ag_old, None) - graph.graph.erase_node(wait_old) - graph.graph.erase_node(ag_old) - - buckets += 1 - - if buckets: - graph.graph.lint() - graph.recompile() - magi_logger.info( - "FSDP weight all-gather bucketing (coalesced): created %d coalesced buckets across submods", buckets - ) - return buckets diff --git a/magi_compiler/passes/graph_split/fsdp_collective_prefetch.py b/magi_compiler/passes/graph_split/fsdp_collective_prefetch.py deleted file mode 100644 index fe0d838..0000000 --- a/magi_compiler/passes/graph_split/fsdp_collective_prefetch.py +++ /dev/null @@ -1,367 +0,0 @@ -from __future__ import annotations - -from collections import deque - -import torch -import torch.fx as fx - -from magi_compiler.utils import magi_logger - -# The two functional-collective ops a SimpleFSDP weight gather lowers to once -# the DTensor ``redistribute``/``to_local`` prims have been turned into explicit -# collectives. This pass operates on *these* nodes (not on the opaque DTensor -# prims), because only here are the launch (``all_gather_into_tensor``) and the -# wait (``wait_tensor``) separate, individually movable FX nodes. -_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default -_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default -_WAIT_TENSOR = torch.ops._c10d_functional.wait_tensor.default - - -def _target_name(target: object) -> str: - if hasattr(target, "name") and callable(getattr(target, "name")): - try: - return str(target.name()) - except Exception: - pass - name = getattr(target, "__name__", None) - if name is not None: - return str(name) - return str(target) - - -def _is_all_gather(node: fx.Node) -> bool: - # Both the single launch and the per-submod coalesced launch are movable. - return node.op == "call_function" and node.target in (_ALL_GATHER, _ALL_GATHER_COALESCED) - - -def _is_wait_tensor(node: fx.Node) -> bool: - return node.op == "call_function" and node.target is _WAIT_TENSOR - - -def _is_param_like(node: fx.Node) -> bool: - """True for a placeholder/get_attr that names a SimpleFSDP weight shard.""" - if node.op not in ("placeholder", "get_attr"): - return False - name = f"{node.name} {node.target}".lower() - return any(token in name for token in ("parameter", "parameters", "weight", "bias")) - - -def _is_transparent_producer(node: fx.Node) -> bool: - """Cheap, structure-only ops that may sit *between* a weight placeholder and - the ``all_gather`` (dtype cast / pad / view). Used when walking backwards - from the launch to (a) decide it gathers a weight and (b) collect the - input-producing chain that must move together with the launch.""" - # The redistribute-lowering pass extracts the local shard with - # ``call_method("to_local", (weight_placeholder,))`` before the gather, so - # the chain from all_gather back to the weight passes through it. - if node.op == "call_method" and str(node.target) in {"to_local", "contiguous", "to", "view", "reshape"}: - return True - if node.op != "call_function": - return False - name = _target_name(node.target) - return any( - token in name - for token in ( - "constant_pad_nd", - "_to_copy", - "aten.to", - "aten::to", - "convert_element_type", - # bucketing inserts cat(reshape(local),...) before the coalesced launch - "aten.cat", - "aten::cat", - "aten.view", - "aten::view", - "aten.reshape", - "aten::reshape", - "_unsafe_view", - "aten.clone", - "aten::clone", - "aten.detach", - "aten::detach", - ) - ) - - -def _is_transparent_consumer(node: fx.Node) -> bool: - """Nodes that forward the gathered weight to its real use without being the - real use themselves: the ``wait_tensor`` plus cheap view/cast/slice ops.""" - if _is_wait_tensor(node): - return True - if node.op == "call_method" and str(node.target) in { - "view", - "reshape", - "flatten", - "squeeze", - "unsqueeze", - "permute", - "transpose", - "t", - "contiguous", - "to", - "detach", - "clone", - }: - return True - if node.op != "call_function": - return False - name = _target_name(node.target) - return any( - token in name - for token in ( - "aten.view", - "aten::view", - "aten.reshape", - "aten::reshape", - "_unsafe_view", - "aten.squeeze", - "aten::squeeze", - "aten.unsqueeze", - "aten::unsqueeze", - "aten.permute", - "aten::permute", - "aten.transpose", - "aten::transpose", - "aten.t", - "aten::t", - "aten.slice", - "aten::slice", - # bucketing's use-side chain: wait -> reshape -> split_with_sizes -> getitem -> reshape - "split_with_sizes", - "aten.detach", - "aten::detach", - "aten.clone", - "aten::clone", - "aten.to", - "aten::to", - "_to_copy", - "convert_element_type", - "operator.getitem", - "getitem", - ) - ) - - -def _gathers_a_weight(all_gather: fx.Node) -> bool: - """Walk backwards from the launch through transparent producers; return True - if the gathered source is a weight/bias placeholder.""" - q: deque[fx.Node] = deque(all_gather.all_input_nodes) - seen: set[fx.Node] = set() - while q: - dep = q.popleft() - if dep in seen: - continue - seen.add(dep) - if _is_param_like(dep): - return True - if _is_transparent_producer(dep): - q.extend(dep.all_input_nodes) - return False - - -def _find_consumer_subgraph_id( - all_gather: fx.Node, node_to_subgraph_id: dict[fx.Node, int] -) -> int | None: - """Forward BFS through wait/transparent nodes to the first real consumer; - return the minimum subgraph id among real consumers.""" - q: deque[fx.Node] = deque(all_gather.users) - seen: set[fx.Node] = set() - consumer_ids: list[int] = [] - while q: - user = q.popleft() - if user in seen: - continue - seen.add(user) - if _is_transparent_consumer(user): - q.extend(user.users) - continue - sid = node_to_subgraph_id.get(user) - if sid is not None: - consumer_ids.append(sid) - if not consumer_ids: - return None - return min(consumer_ids) - - -def _collect_launch_input_chain(all_gather: fx.Node) -> list[fx.Node]: - """Backwards from the launch, collect transparent producer nodes that feed - *only* into this launch (their results are not used elsewhere). These must - move with the launch so it is self-contained in the previous submod. - - The weight placeholder/get_attr itself is NOT included (placeholders are not - movable graph nodes); only intermediate compute that prepares the shard. - """ - chain: list[fx.Node] = [] - q: deque[fx.Node] = deque(n for n in all_gather.all_input_nodes if _is_transparent_producer(n)) - seen: set[fx.Node] = set() - frontier = {all_gather} - while q: - node = q.popleft() - if node in seen: - continue - # Movable only if every consumer is within the chain we are moving. - if not all(u in frontier or u in seen or u is all_gather for u in node.users): - continue - seen.add(node) - chain.append(node) - frontier.add(node) - for dep in node.all_input_nodes: - if _is_transparent_producer(dep): - q.append(dep) - return chain - - -def _build_first_compute_node_by_subgraph_id( - graph: fx.GraphModule, node_to_subgraph_id: dict[fx.Node, int] -) -> dict[int, fx.Node]: - """First *compute* node of each submod, used as the prefetch anchor. - - The anchor MUST be a node that genuinely STAYS in the target submod — i.e. a - real compute node (matmul, custom op, norm, ...), NOT any weight-gather-chain - node. After redistribute lowering a gather is a chain - ``to_local -> [_to_copy] -> [pad] -> all_gather -> wait -> [slice]``; every one - of those is itself relocatable by this pass. If we anchored on, say, the - first ``to_local`` of the target submod, that node could be hoisted away to an - even earlier submod, and prepending another launch before the (now-moved) - anchor would drag it to the very front of the graph — corrupting submod - ordering and making ``split_module`` raise KeyError on partition inputs. - Skipping the whole gather-chain vocabulary keeps the anchor stable. - """ - first: dict[int, fx.Node] = {} - for node in graph.graph.nodes: - if node.op in ("placeholder", "output"): - continue - # Skip every node that is part of a (relocatable) weight-gather chain. - if ( - _is_all_gather(node) - or _is_wait_tensor(node) - or _is_transparent_producer(node) - or _is_transparent_consumer(node) - ): - continue - sid = node_to_subgraph_id.get(node) - if sid is not None and sid not in first: - first[sid] = node - return first - - -def _deps_available_before_anchor( - nodes: list[fx.Node], anchor: fx.Node, node_index: dict[fx.Node, int] -) -> bool: - """Every external dependency of the to-be-moved nodes must already be - defined before the anchor (so the move does not read an undefined value).""" - moved = set(nodes) - anchor_index = node_index[anchor] - for node in nodes: - for dep in node.all_input_nodes: - if dep in moved: - continue - if dep.op in ("placeholder", "get_attr"): - continue - if node_index.get(dep, anchor_index + 1) > anchor_index: - return False - return True - - -def apply_fsdp_collective_prefetch( - graph: fx.GraphModule, - node_to_subgraph_id: dict[fx.Node, int], - *, - distance: int = 1, -) -> int: - """Prefetch SimpleFSDP weight all-gather across a submod boundary. - - For each ``all_gather_into_tensor`` whose gathered result feeds a weight use - in submod ``N``, move ONLY the launch node (and the transparent producer - chain that feeds exclusively into it) to the beginning of submod - ``N - distance``. The ``wait_tensor`` and everything after it are left - untouched at the original use site in submod ``N``. - - Because ``split_module(keep_original_order=True)`` runs afterwards, the - gathered tensor automatically becomes an output of submod ``N - distance`` - and an input of submod ``N`` -- so the collective is launched during the - previous submod's compute and only waited on right before use, giving real - compute/communication overlap. - - Unlike :func:`apply_fsdp_all_gather_prefetch`, this pass never moves the - wait: at the explicit-collective level the launch and wait are distinct FX - nodes, so leaving the wait behind actually keeps it at the consumer. - - IMPORTANT -- this FX move is necessary but NOT sufficient for overlap. The - hoisted ``all_gather`` has no consumer inside submod ``N - distance`` (its - result is only a graph output), so Inductor's per-submod scheduler re-sinks - it to just before ``return`` -- at runtime the launch is then issued AFTER - that submod's compute and cannot overlap. The caller must therefore also - enable Inductor's comm/compute reorder with the ``raise_comms`` pass - (``reorder_for_compute_comm_overlap=True``, - ``reorder_for_compute_comm_overlap_passes=["raise_comms", "sink_waits"]``), - which hoists the launch back to the top of the generated code. MagiBackend - wires this automatically when this pass moves any launch. - """ - if distance <= 0: - return 0 - - first_node_by_subgraph_id = _build_first_compute_node_by_subgraph_id(graph, node_to_subgraph_id) - node_index = {node: idx for idx, node in enumerate(graph.graph.nodes)} - moved = 0 - changed = False - - for node in list(graph.graph.nodes): - if not _is_all_gather(node): - continue - if not _gathers_a_weight(node): - continue - - current_id = node_to_subgraph_id.get(node) - if current_id is None: - continue - - consumer_id = _find_consumer_subgraph_id(node, node_to_subgraph_id) - if consumer_id is None: - continue - - target_id = max(0, consumer_id - distance) - if current_id <= target_id: - # Already launched at or before the desired submod. - continue - - anchor = first_node_by_subgraph_id.get(target_id) - if anchor is None: - continue - - # Move the launch plus the producer chain that feeds only into it. - chain = _collect_launch_input_chain(node) - # For a coalesced launch the op returns a list[Tensor] that must NOT cross - # a submod boundary; its operator.getitem unpackers (whose only input is - # the launch) move together with it so the list stays in the target submod - # and only the per-member single tensors cross to the consumer. - getitems = [] - if node.target is _ALL_GATHER_COALESCED or node.meta.get("magi_fsdp_weight_ag_coalesced"): - getitems = [ - u for u in node.users - if u.op == "call_function" and getattr(u.target, "__name__", "") == "getitem" - ] - # Order from earliest to latest so prepend keeps dependency order. - move_group = sorted([node, *chain, *getitems], key=lambda n: node_index[n]) - - if not _deps_available_before_anchor(move_group, anchor, node_index): - magi_logger.debug( - "Skip collective prefetch for %s: deps not available before %s", node.name, anchor.name - ) - continue - - for moved_node in move_group: - node_to_subgraph_id[moved_node] = target_id - anchor.prepend(moved_node) - - node.meta["magi_fsdp_prefetch_from_subgraph"] = current_id - node.meta["magi_fsdp_prefetch_to_subgraph"] = target_id - node.meta["magi_fsdp_prefetch_for_consumer_subgraph"] = consumer_id - - moved += 1 - changed = True - - if changed: - graph.graph.lint() - graph.recompile() - return moved diff --git a/magi_compiler/profiling/__init__.py b/magi_compiler/profiling/__init__.py new file mode 100644 index 0000000..45c7883 --- /dev/null +++ b/magi_compiler/profiling/__init__.py @@ -0,0 +1,31 @@ +# 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. + +"""Per-op runtime cost model (profiling) -- general infrastructure, not a pass. + +Provides: +* :class:`.runtime_estimator.ProfilingRuntimeEstimator` -- a ``snode -> nanoseconds`` + cost model that MEASURES each scheduler node's real kernel time (Triton via + ``benchmark_fused_nodes``, extern/custom via eager replay, collectives via a + rank-lockstep replay), replacing Inductor's unreliable analytical roofline. Any + pass that needs accurate per-op timing (e.g. the FSDP-overlap reorder) uses it. +* :mod:`.benchmark_inputs` -- a model-agnostic registry so the owning code (e.g. + athena's gaga4) can supply valid replay inputs / lockstep flags for opaque custom + ops that can't be replayed from generic size-hinted tensors alone. +""" + +from .benchmark_inputs import get_benchmark_inputs_hook, op_has_internal_collective, register_benchmark_inputs +from .runtime_estimator import ProfilingRuntimeEstimator + +__all__ = ["ProfilingRuntimeEstimator", "register_benchmark_inputs", "get_benchmark_inputs_hook", "op_has_internal_collective"] diff --git a/magi_compiler/profiling/benchmark_inputs.py b/magi_compiler/profiling/benchmark_inputs.py new file mode 100644 index 0000000..bdf7a56 --- /dev/null +++ b/magi_compiler/profiling/benchmark_inputs.py @@ -0,0 +1,71 @@ +# 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. + +"""Per-op benchmark-input hooks for the profiling runtime estimator. + +Some opaque boundary custom ops cannot be replayed for timing from generic +size-hinted tensors alone, because they carry VALUE-DEPENDENT metadata that must +be self-consistent (not just right-shaped) or they raise -- e.g. +``gaga4_fa3_with_sink_cp`` takes a ``cp_split_sizes`` tensor whose values must sum +to the sequence length and equal the per-rank seqlen, else its internal CP +all_to_all asserts and the op falls back to a 0 cost estimate. + +The owning code (which knows the op's semantics -- e.g. athena's gaga4 model) +registers a hook here that builds a VALID, RANK-DETERMINISTIC set of replay inputs. +MagiCompiler stays free of model-specific knowledge: ``_measure_extern`` just looks +up the op by name and, if a hook exists, uses it instead of the generic realizer. + +A hook is ``fn(fx_node, realize) -> (args, kwargs) | None``: + * ``fx_node`` -- the ``torch.fx.Node`` for the op (read its ``args``/``kwargs`` + and their ``meta['val']`` for shapes); + * ``realize`` -- MagiCompiler's default arg realizer (fx.Node/SymInt/container + -> concrete tensor/int), so a hook can reuse it for the plain + tensor args and only special-case the metadata ones; + * returns ``(args, kwargs)`` of real objects to call the op with, or ``None`` to + fall back to the generic path. + +Hooks MUST produce rank-identical inputs (derive everything from shapes / static +sizes, no per-rank timing or state) so the rank-lockstep ``warm_and_sync`` measure +issues any internal collective in lockstep across ranks. +""" + +from __future__ import annotations + +from typing import Callable + +# op name (``torch.ops.ns.op`` overload string, e.g. "athena::gaga4_fa3_with_sink_cp") +# -> hook. Also records which ops issue an internal collective (need fixed-iter +# lockstep replay) -- so MagiCompiler no longer hardcodes model-specific op names. +_BENCHMARK_INPUT_HOOKS: dict[str, Callable] = {} +_INTERNAL_COLLECTIVE_OPS: set[str] = set() + + +def register_benchmark_inputs(op_name: str, fn: Callable, *, has_internal_collective: bool = False) -> None: + """Register a replay-input builder for ``op_name`` (see module docstring). + + ``has_internal_collective``: mark this op as issuing an internal collective, so + the estimator replays it with a FIXED iteration count under barriers (keeps every + rank in lockstep -- a duration-adaptive count would desync the internal NCCL op). + """ + _BENCHMARK_INPUT_HOOKS[op_name] = fn + if has_internal_collective: + _INTERNAL_COLLECTIVE_OPS.add(op_name) + + +def get_benchmark_inputs_hook(op_name: str) -> Callable | None: + return _BENCHMARK_INPUT_HOOKS.get(op_name) + + +def op_has_internal_collective(op_name: str) -> bool: + return op_name in _INTERNAL_COLLECTIVE_OPS diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py new file mode 100644 index 0000000..af7ad81 --- /dev/null +++ b/magi_compiler/profiling/runtime_estimator.py @@ -0,0 +1,758 @@ +# 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. + +from __future__ import annotations + +"""Profiling-based ``estimate_op_runtime`` replacement. + +Inductor's default ``BaseSchedulerNode.get_estimated_runtime`` is a pure +analytical roofline. It is unreliable for exactly the nodes our FSDP overlap +reorder pass must size: + +* fused pointwise/reduction kernels -> ``estimate_flops()`` is None, so the + estimate degrades to ``bytes / dram_bw`` and IGNORES the compute fused into the + kernel (measured 60x under-estimate for a 32-deep fusion); +* matmul -> the device TFLOPS table is wrong on this box (0.5 vs ~990), giving a + ~1500x over-estimate; +* custom ops (Triton / flash-attn) -> ``ExternKernelSchedulerNode`` with + ``MultiOutputLayout`` -> dtype is None -> the estimate is silently 0. + +Full write-up + demos: ``scripts/demo/research_estimate_op_runtime_findings.md``. + +Instead we MEASURE each scheduler node's real kernel time: +* Triton (pointwise/reduction/template) snodes -> ``scheduler.benchmark_fused_nodes``, + which codegens the SAME fused kernel production emits and ``do_bench``es it + (verified 0.995-1.15x of the compiled-graph kernel); +* extern snodes (matmul / custom op) -> replay the aten/custom op on real inputs + rebuilt from the fx fake-tensor meta, timed with the Inductor benchmarker. + +COLLECTIVES are NOT benchmarked here -- they always use the analytical +``estimate_nccl_collective_runtime``. Benchmarking a real collective during compile +is not multi-rank-safe: the reorder pass runs per-rank during independent, +non-co-scheduled compilation, so issuing an NCCL op desyncs ranks -> watchdog hang +(seen on 8-GPU gaga4: rank0 raced to FSDP seqNum~194 while peers never matched). +For real measured comm, calibrate OUTSIDE compile at a synchronized point. How far +the analytical estimate is from reality is measured by +``scripts/demo/verify_collective_estimate.py`` (a standalone, properly-synchronized +harness -- not the compile path). + +Op -> time table +---------------- +The estimator maintains ``self._table: dict[key -> ProfileEntry]`` -- the +persistent op->time structure for COMPUTE nodes. The KEY is the op's STRUCTURAL +identity ``(target op, tuple[(input shape, dtype)])`` (``_structural_key``), NOT the +snode's unique name (``buf0``/``op42`` are unique per node and would defeat reuse). +Each distinct key is benchmarked ONCE on first encounter; every later isomorphic +snode (same op + same input shapes, e.g. the same matmul in every layer) reuses the +entry (``reuse_count++``). A 40-layer model thus measures O(#distinct kernels), not +O(#nodes). ``ProfileEntry`` carries ``(ns, kind, label, measured, reuse_count)``; +``estimator.summary()`` prints the whole table (see it at DEBUG log level). + +Passed to the reorder pass as ``cost_fn`` (a callable ``snode -> nanoseconds``). +The extern measurement path is ShapeEnv-isolated (real tensors from +size_hints, eager call) so they run on the dynamic base compile; only benchmark_fused_nodes +(fused Triton) would specialize the dynamic dim, so that path stays analytical while +free symbols exist. +""" + +import dataclasses +from typing import Any + +import torch +from torch._inductor.runtime.benchmarking import benchmarker +from torch._inductor.scheduler import BaseSchedulerNode, ExternKernelSchedulerNode, FusedSchedulerNode +from torch._inductor.utils import contains_collective, contains_wait +from torch._inductor.virtualized import V + +from magi_compiler.utils import magi_logger + +from .benchmark_inputs import get_benchmark_inputs_hook, op_has_internal_collective + +# Dedicated GLOO (CPU) group for exchanging profile metadata across ranks (built +# once, cached). A CPU/gloo group keeps the cost sync off the NCCL process groups +# the forward uses, so it can never interleave with / desync the weight-gather or +# CP collectives. +_COST_SYNC_GROUP = "uninit" + + +def _get_cost_sync_group(): + global _COST_SYNC_GROUP + import torch.distributed as dist + + if _COST_SYNC_GROUP != "uninit": + return _COST_SYNC_GROUP + try: + _COST_SYNC_GROUP = dist.new_group(backend="gloo") + except Exception as exc: # noqa: BLE001 + magi_logger.warning("cost-sync: gloo group unavailable (%s); using default group", exc) + _COST_SYNC_GROUP = None + return _COST_SYNC_GROUP + + +@dataclasses.dataclass +class ProfileEntry: + """One row of the op -> time table.""" + + ns: float # measured (or analytical-fallback) runtime, nanoseconds + kind: str # "compute" | "extern" | "collective" + label: str # human-readable op identity (target + shapes), for logs + measured: bool # True if really benchmarked, False if analytical fallback + reuse_count: int = 0 # how many later snodes reused this entry + + +def _snode_label(snode: BaseSchedulerNode, max_shapes: int = 3) -> str: + """Short human-readable identity of an snode for the profile table / logs: + the op target plus its first few input shapes. (The snode's own NAME, e.g. + 'op123', is deliberately NOT the cache key -- it is unique per node and would + defeat cross-layer reuse; the key is the structural identity below.)""" + node = getattr(snode, "node", None) + origin = node.get_origin_node() if (node is not None and hasattr(node, "get_origin_node")) else None + target = str(getattr(origin, "target", type(node).__name__ if node is not None else "?")) + target = target.split("(")[0].split(" ")[-1][-40:] + shapes = [] + if origin is not None: + for a in (*origin.args, *getattr(origin, "kwargs", {}).values()): + ev = a.meta.get("val") if isinstance(a, torch.fx.Node) else None + if isinstance(ev, torch.Tensor): + shapes.append("x".join(str(x) for x in _static(ev.shape))) + if len(shapes) >= max_shapes: + break + return f"{target}[{','.join(shapes)}]" if shapes else target + + +def _structural_key(snode: BaseSchedulerNode) -> tuple | None: + """A cache key that is identical for isomorphic kernels (same op set + same + input shapes/dtypes) so repeated layers share one measurement. Returns None + when we can't build a stable key (then we don't cache).""" + parts: list[Any] = [] + for n in snode.get_nodes(): + node = getattr(n, "node", None) + if node is None: + return None + origin = node.get_origin_node() if hasattr(node, "get_origin_node") else None + target = str(getattr(origin, "target", type(node).__name__)) + shapes: list[Any] = [] + if origin is not None: + for a in (*origin.args, *origin.kwargs.values()): + ev = a.meta.get("val") if isinstance(a, torch.fx.Node) else None + if isinstance(ev, torch.Tensor): + shapes.append((tuple(_static(ev.shape)), str(ev.dtype))) + parts.append((target, tuple(shapes))) + return tuple(parts) + + +def _is_symbolic(s) -> bool: + return isinstance(s, torch.SymInt) or hasattr(s, "node") + + +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).""" + out = [] + for s in shape: + if _is_symbolic(s): + out.append(str(s)) + else: + out.append(int(s)) + return tuple(out) + + +def _concrete_size(s, fallback: int = 1) -> int: + """A concrete size for building real benchmark inputs, WITHOUT specializing: + use Inductor's size_hint (reads the hint, adds no guard).""" + if _is_symbolic(s): + try: + return int(V.graph.sizevars.size_hint(s, fallback=fallback)) + except Exception: # noqa: BLE001 + return fallback + return int(s) + + +def _realize_arg(v): + """Turn an fx arg into a concrete replay input (rank-deterministic). + + fx.Node(tensor) -> right-shaped tensor from size-hints; fx.Node/bare SymInt -> + concrete int; list/tuple/dict -> recursively realized PLAIN container. The + container de-immutabilization matters: FX stores list args as + torch.fx.immutable_collections.immutable_list, and the custom-op C++ arg parser + requires a plain List[int] for a ``SymInt[]`` arg -- an immutable_list of still- + symbolic / nested elements is rejected ("Expected List[int] ... found + immutable_list"), making the op fall back to a 0 cost.""" + if isinstance(v, torch.fx.Node): + ev = v.meta.get("val") + if isinstance(ev, torch.Tensor): + shape = tuple(_concrete_size(s) for s in ev.shape) + if ev.is_floating_point(): + return torch.randn(shape, device=ev.device, dtype=ev.dtype) + return torch.zeros(shape, device=ev.device, dtype=ev.dtype) + if _is_symbolic(ev) or isinstance(ev, int): + return _concrete_size(ev) # a Node carrying a scalar -> concrete hint + return v + if _is_symbolic(v): + return _concrete_size(v) + if isinstance(v, (list, tuple)): + realized = [_realize_arg(x) for x in v] + return type(v)(realized) if type(v) in (list, tuple) else list(realized) + if isinstance(v, dict): + return {k: _realize_arg(x) for k, x in v.items()} + return v + + +def _measure_extern(snode: ExternKernelSchedulerNode, fixed_iters: bool = False) -> float: + """Time an extern (matmul / custom-op) snode by replaying its aten op. + + ``fixed_iters``: when True, time a CONSTANT number of iterations with CUDA events + instead of the duration-adaptive ``benchmark_gpu``. REQUIRED for ops that issue + an INTERNAL collective (e.g. the CP ``all_to_all`` inside gaga4_fa3_with_sink_cp): + the adaptive benchmarker runs a rank-dependent iteration count, so different ranks + would issue different numbers of that internal collective -> NCCL count mismatch -> + deadlock. A fixed count keeps every rank in lockstep (mirrors _measure_collective_op). + + Ops whose replay needs VALUE-CONSISTENT metadata (not just right-shaped tensors) + can register a hook via ``benchmark_inputs.register_benchmark_inputs`` that builds + a valid ``(args, kwargs)``; otherwise args come from the generic ``_realize_arg``.""" + fx_node = snode.node.get_origin_node() + if fx_node is None: + return 0.0 + target = fx_node.target + + hook = get_benchmark_inputs_hook(_op_name(target)) + built = hook(fx_node, _realize_arg) if hook is not None else None + if built is not None: + args, kwargs = built + else: + args = tuple(_realize_arg(a) for a in fx_node.args) + kwargs = {k: _realize_arg(v) for k, v in fx_node.kwargs.items()} + + # Replay EAGERLY, decoupled from the enclosing compile: + # * torch._dynamo.disable(): _measure_extern runs while the OUTER graph is being + # compiled (Dynamo/Inductor active). A custom boundary op whose impl contains + # torch.compile'd regions (e.g. gaga4_fa3_with_sink_cp's sink-correction path) + # would otherwise RE-ENTER Dynamo on these concrete tensors, re-trace with + # DYNAMIC shapes (symbolic s*), and blow up ("Dynamo failed to run FX node ... + # broadcast" / "can't pickle cyclic objects") -> the op fell back to a 0 cost. + # We want the EAGER kernel time, so disable Dynamo for the replay. + # * no_grad: the compiled forward runs under inference_mode; some ops branch on + # torch.is_grad_enabled() (fa3 takes a training-only flash-attn wrapper path when + # grad is on, incompatible with the installed flash-attn). Match inference. + @torch._dynamo.disable + def _call(): + return target(*args, **kwargs) + + def fn(): + with torch.no_grad(): + return _call() + + if not fixed_iters: + fn() # warmup / correctness + return benchmarker.benchmark_gpu(fn) * 1e6 # ms -> ns + # Fixed-iteration timing (lockstep-safe for internal collectives). + _WARMUP, _ITERS = 3, 10 + for _ in range(_WARMUP): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(_ITERS): + fn() + end.record() + torch.cuda.synchronize() + return (start.elapsed_time(end) / _ITERS) * 1e6 # ms/iter -> ns + + +def _op_name(target) -> str: + """Overload-qualified op name (e.g. 'athena::gaga4_fa3_with_sink_cp'), or '' for + a non-op target. Used to look ops up in the benchmark-input registry.""" + name = getattr(target, "name", None) + if callable(name): + try: + return name() # OpOverload.name() -> 'ns::op' + except Exception: # noqa: BLE001 + return "" + return "" + + +def _extern_has_internal_collective(snode: BaseSchedulerNode) -> bool: + """True for opaque boundary ops that issue collectives internally (CP attention / + MoE), so we must measure them with fixed iterations under a barrier. The set of + such ops is declared by the owning code via ``register_benchmark_inputs(..., + has_internal_collective=True)`` -- MagiCompiler holds no model-specific op names.""" + node = getattr(snode, "node", None) + origin = node.get_origin_node() if (node is not None and hasattr(node, "get_origin_node")) else None + target = getattr(origin, "target", None) if origin is not None else None + return op_has_internal_collective(_op_name(target)) if target is not None else False + + +# ---- collective (weight all-gather) benchmarking -------------------------- +_AG = torch.ops._c10d_functional.all_gather_into_tensor.default +_AG_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default +_WAIT = torch.ops._c10d_functional.wait_tensor.default + + +def _leaf_collective(snode: BaseSchedulerNode): + """The underlying _CollectiveKernel IR node (unwraps GroupedSchedulerNode).""" + from torch._inductor.utils import is_collective + + node = getattr(snode, "node", None) + if node is not None and is_collective(node): + return node + for child in getattr(snode, "snodes", []) or []: + cn = getattr(child, "node", None) + if cn is not None and is_collective(cn): + return cn + return None + + +def _collective_spec(node): + """(op_overload, group_name, group_size, [(shape, dtype, device), ...]) for a + collective IR node, or None if it isn't a benchmarkable all-gather.""" + op = getattr(node, "op_overload", None) + if op not in (_AG, _AG_COALESCED): + return None + group_name = node.constant_args[-1] # (..., group_size, group_name) + from torch.distributed.distributed_c10d import _get_group_size_by_name + + group_size = _get_group_size_by_name(group_name) + specs = [] + for inp in node.inputs: + shape = tuple(_concrete_size(s) for s in inp.layout.size) + specs.append((shape, inp.layout.dtype, inp.layout.device)) + return op, group_name, group_size, specs + + +def _collective_label(snode: BaseSchedulerNode) -> str: + """Readable identity of a collective: op name, world size, #inputs + first shape.""" + node = _leaf_collective(snode) + spec = _collective_spec(node) if node is not None else None + if spec is None: + return _snode_label(snode) + _op, _group, group_size, specs = spec + shape0 = "x".join(str(x) for x in specs[0][0]) if specs else "?" + return f"all_gather(ws={group_size},n={len(specs)},{shape0})" + + +def _measure_collective_op(snode: BaseSchedulerNode) -> float: + """Replay the functional all-gather (+wait) on real tensors and time it.""" + node = _leaf_collective(snode) + if node is None: + return 0.0 + spec = _collective_spec(node) + if spec is None: + return 0.0 + op, group_name, group_size, specs = spec + + ins = [torch.empty(shape, dtype=dt, device=dev) for shape, dt, dev in specs] + if op is _AG_COALESCED: + + def fn(): + outs = _AG_COALESCED(ins, group_size, group_name) + for o in outs: + _WAIT(o) + + else: + + def fn(): + _WAIT(_AG(ins[0], group_size, group_name)) + + # FIXED iteration counts across ALL ranks. A duration-based benchmarker + # (e.g. benchmark_gpu shrinks benchmark_iters by per-rank estimated runtime) + # would issue a DIFFERENT number of collectives on different ranks -> + # NCCL count mismatch -> deadlock. A constant loop keeps every rank in lockstep. + _WARMUP, _ITERS = 3, 10 + for _ in range(_WARMUP): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(_ITERS): + fn() + end.record() + torch.cuda.synchronize() + return (start.elapsed_time(end) / _ITERS) * 1e6 # ms/iter -> ns + + +class ProfilingRuntimeEstimator: + """Callable ``snode -> ns`` for ``config.estimate_op_runtime``. + + Measures compute nodes with real kernels (memoized); defers collectives and + waits to the analytical ``get_estimated_runtime`` (the reorder pass sizes + comm separately via ``estimate_nccl_collective_runtime`` or a calibrated + value). Never raises -- on any failure returns the analytical estimate (or + 0), so it degrades to today's behaviour rather than breaking compilation. + """ + + def __init__(self) -> None: + # The op -> time table. Key is the STRUCTURAL identity of the op + # (target op + input shapes/dtypes for compute, or (coll, group, group_size, + # in-shapes) for collectives) -- NOT the snode's unique name, so isomorphic + # ops across repeated layers share ONE measurement. Each distinct key is + # profiled once on first encounter; later encounters reuse the entry. + self._table: dict[tuple, ProfileEntry] = {} + self.n_measured = 0 + self.n_cache_hits = 0 + # Set True by MagiBackend for multi-rank runs: the reorder pass then calls + # warm_and_sync() to reconcile costs across ranks (rank-identical schedule). + self._sync_across_ranks = False + # Transient {structural_key -> representative snode}, used ONLY by + # warm_and_sync() to re-measure in rank-lockstep order. Kept OFF the + # ProfileEntry (which Inductor pickles into the fx-graph cache key) because + # snodes hold FakeTensors that are unpicklable / cyclic. Never serialized. + self._key_snode: dict = {} + + def __deepcopy__(self, memo): + # Shared by reference from the reorder pass; when config serialization + # deepcopies the pass list, return a clean instance (the memoized + # measurements are transient and hold no tensors, but avoid copying them). + new = ProfilingRuntimeEstimator() + new._sync_across_ranks = self._sync_across_ranks + memo[id(self)] = new + return new + + # Backward-compat alias: some callers/tests read `.table`. + @property + def table(self) -> "dict[tuple, ProfileEntry]": + return self._table + + def warm_and_sync(self) -> int: + """Rank-LOCKSTEP profiling of every distinct op, so multi-rank runs get REAL + measured costs (more accurate than the analytical roofline) AND a rank- + identical cost table (required so the FSDP-overlap reorder produces the same + schedule on every rank -- else weight-PG gathers interleave with the eager CP + all_to_all in rank-divergent order -> deadlock). + + PRECONDITION (caller guarantees): the graph is structurally IDENTICAL on all + ranks (the unconditional-pad lowering fix ensures this). Therefore every rank + has exactly the same set of ``_structural_key`` keys already populated in the + table by the reorder pass's warm-up loop, so: + * the key iteration order (sorted) is identical on every rank; + * for a collective-containing op (attention/MoE), every rank measures it at + the same step, so the barrier-wrapped fixed-iteration replay issues the + op's internal CP all_to_all in lockstep (no NCCL count mismatch); + * the max-reduce over gloo is symmetric. + + Steps: for each key in sorted order, barrier -> (re)measure locally with FIXED + iters -> barrier; then all_gather_object the {key: ns} maps and take the MAX. + Fixed-iteration timing (``_measure_one``) is mandatory -- a duration-adaptive + benchmark would run a per-rank iteration count and desync the internal + collective. Returns the number of table entries whose cost changed.""" + import torch.distributed as dist + + if not (dist.is_available() and dist.is_initialized()): + return 0 + world = dist.get_world_size() + if world <= 1: + return 0 + group = _get_cost_sync_group() + + # The reorder warm-up already populated self._table (one entry per distinct + # structural key) and self._key_snode (a representative snode per key). + # Re-measure each in a rank-uniform (sorted) order under barriers so any + # internal collective is issued in lockstep across ranks. + keys = sorted(self._table.keys(), key=repr) + local_ns: dict = {} + for k in keys: + snode = self._key_snode.get(k) + dist.barrier(group=group) + if snode is not None: + local_ns[k] = self._measure_one(snode) + else: + local_ns[k] = self._table[k].ns # no cached snode -> keep prior measurement + dist.barrier(group=group) + + # keys re-measured on THIS rank (had a representative snode) -- gather across + # ranks so a key measured on any rank is flagged measured (the graph is + # identical, so this is the same set everywhere; the union is just defensive). + measured_here = set(self._key_snode.keys()) + + gathered: list = [None] * world + dist.all_gather_object(gathered, local_ns, group=group) + gathered_measured: list = [None] * world + dist.all_gather_object(gathered_measured, list(measured_here), group=group) + merged: dict = {} + for d in gathered: + for k, ns in (d or {}).items(): + if k not in merged or ns > merged[k]: + merged[k] = ns + measured_keys = set() + for mk in gathered_measured: + measured_keys.update(mk or []) + n = 0 + for k, e in self._table.items(): + if k in measured_keys: + e.measured = True # reconciled from a real rank-lockstep measurement + m = merged.get(k) + if m is not None and m != e.ns: + e.ns = m + n += 1 + self._key_snode.clear() # drop snode refs (unpicklable) once sync is done + return n + + def _measure_one(self, snode: BaseSchedulerNode) -> float: + """Measure a single snode with FIXED iterations (lockstep-safe), never raising + -- falls back to the analytical estimate. Collective-containing externs + (attention / MoE) use fixed-iter replay so every rank issues the same number + of internal collectives.""" + try: + if contains_collective(snode): + return _measure_collective_op(snode) + if isinstance(snode, ExternKernelSchedulerNode): + fixed = _extern_has_internal_collective(snode) + with _shapeenv_sandbox(), _suppress_guards(): + ns = _measure_extern(snode, fixed_iters=fixed) + self.n_measured += 1 + return ns + return self._measure(snode) + except BaseException as exc: # noqa: BLE001 + magi_logger.debug("warm/sync measure fell back to analytical for %s: %s", snode.get_name(), exc) + return _safe_analytical(snode) + + def summary(self) -> str: + """One line per distinct profiled op: kind, label, per-call time, #calls, + aggregate time (per-call * #calls). Each op also gets a machine-parseable + ``ESTLINE`` tag (kind|label|per_call_us|calls|total_us|measured) so a run's + estimates can be diffed against the real nsys kernel trace -- see + scripts/demo/compare_estimate_vs_nsys.py.""" + lines = [] + for e in sorted(self._table.values(), key=lambda e: -e.ns * (e.reuse_count + 1)): + calls = e.reuse_count + 1 # first encounter + reuses + per_us = e.ns / 1e3 + total_us = per_us * calls + meas = "measured" if e.measured else "analytical" + lines.append(f" [{e.kind:10}] {e.label:<48} {per_us:9.2f}us/call x{calls:<4} " f"= {total_us:11.2f}us ({meas})") + # grep-friendly: ESTLINE|kind|label|per_call_us|calls|total_us|measured + lines.append(f" ESTLINE|{e.kind}|{e.label}|{per_us:.3f}|{calls}|{total_us:.3f}|{meas}") + return ( + f"profile table: {len(self._table)} distinct ops, " + f"{self.n_measured} measured, {self.n_cache_hits} reuses\n" + "\n".join(lines) + ) + + def __call__(self, snode: BaseSchedulerNode) -> float: + # A wait_tensor kernel itself takes ~0 time (the collective's cost is + # attributed to the launch); keep it analytical (returns 0). + if contains_wait(snode) and not contains_collective(snode): + return _safe_analytical(snode) + + # COLLECTIVES: never benchmark a real collective HERE. The Inductor reorder + # pass runs per-rank during independent, non-co-scheduled compilation, so + # issuing an NCCL op in __call__ desyncs ranks -> watchdog hang (observed on + # 8-GPU gaga4: rank0 raced to FSDP seqNum~194 while peers never matched). + # SEED the cost with Inductor's static analytical NCCL estimate (a pure + # function of shapes+device -> rank-deterministic). In profile_sync mode we + # ALSO stash the snode so the rank-lockstep warm_and_sync() later OVERRIDES + # this seed with a REAL measured time (the only multi-rank-safe place to touch + # NCCL) -- which is where the accuracy comes from. Collectives enter the same + # op->time table as compute so warm_and_sync can find them. + if contains_collective(snode): + cnode = _leaf_collective(snode) + spec = _collective_spec(cnode) if cnode is not None else None + if spec is None: + return _safe_analytical(snode) # non-AG / unparseable -> old behaviour + op, _group_name, group_size, specs = spec + ckey = ("collective", str(op), group_size, tuple((tuple(shape), str(dt)) for shape, dt, _dev in specs)) + entry = self._table.get(ckey) + if entry is not None: + entry.reuse_count += 1 + self.n_cache_hits += 1 + return entry.ns + ns = _safe_analytical(snode) # Inductor static estimate as the seed + self._table[ckey] = ProfileEntry(ns=ns, kind="collective", label=_collective_label(snode), measured=False) + if self._sync_across_ranks: + self._key_snode[ckey] = snode # warm_and_sync -> real measured override + return ns + + is_extern = isinstance(snode, ExternKernelSchedulerNode) + + # Benchmark compute nodes on the dynamic graph without specializing the + # dynamic dim. Extern (matmul / custom op) is ShapeEnv-isolated (replays the + # aten op on size-hinted real tensors) so it is safe even with free symbols. + # Fused Triton pointwise/reduction goes through benchmark_fused_nodes which + # re-enters Inductor codegen bound to the live ShapeEnv and WOULD specialize, + # so that path stays analytical while the graph is dynamic. Matmul/attention/ + # MoE (the dominant costs, and the source of the bogus roofline) are extern. + if not is_extern and _graph_has_free_symbols(): + return _safe_analytical(snode) + + # --- op -> time table: profile a distinct key ONCE, reuse afterwards --- + key = _structural_key(snode) + if key is not None: + entry = self._table.get(key) + if entry is not None: + entry.reuse_count += 1 + self.n_cache_hits += 1 + return entry.ns + + # First encounter of this key -> measure it. Measuring must NEVER break + # compilation: any failure (unbenchmarkable extern, FakeTensor deepcopy in + # the benchmark harness, ...) falls back to the analytical estimate. + measured = True + try: + ns = self._measure_extern_safe(snode) if is_extern else self._measure(snode) + except BaseException as exc: # noqa: BLE001 + magi_logger.debug("Profiling estimator fell back to analytical for %s: %s", snode.get_name(), exc) + ns = _safe_analytical(snode) + measured = False + + if key is not None: + kind = "extern" if is_extern else "compute" + label = _snode_label(snode) + self._table[key] = ProfileEntry(ns=ns, kind=kind, label=label, measured=measured) + # Remember a representative snode (transient, unpicklable) so + # warm_and_sync() can re-measure this key in rank-lockstep order. + if self._sync_across_ranks: + self._key_snode[key] = snode + magi_logger.debug( + "profile[%s] %s -> %.2fus%s", kind, label, ns / 1e3, "" if measured else " (analytical fallback)" + ) + return ns + + def _measure_extern_safe(self, snode: BaseSchedulerNode) -> float: + with _shapeenv_sandbox(), _suppress_guards(): + ns = _measure_extern(snode) + self.n_measured += 1 + return ns + + def _measure(self, snode: BaseSchedulerNode) -> float: + # Benchmarking runs real kernels at concrete (hinted) shapes. Under + # dynamic-shape compilation that would add ``Eq(sym, hint)`` guards and + # replacements into the ShapeEnv and SPECIALIZE the dynamic dim (breaking + # the compile for other seq lens). Suppress guard creation AND snapshot / + # restore the ShapeEnv's mutable specialization state so nothing leaks. + with _shapeenv_sandbox(), _suppress_guards(): + return self._measure_inner(snode) + + def _measure_inner(self, snode: BaseSchedulerNode) -> float: + try: + if isinstance(snode, ExternKernelSchedulerNode): + self.n_measured += 1 + return _measure_extern(snode) + scheduler = V.graph.scheduler + nodes = list(snode.get_nodes()) if isinstance(snode, FusedSchedulerNode) else [snode] + ms, _ = scheduler.benchmark_fused_nodes(nodes) + self.n_measured += 1 + return ms * 1e6 + except Exception as exc: # noqa: BLE001 + magi_logger.debug("Profiling estimator fell back to analytical for %s: %s", snode.get_name(), exc) + return _safe_analytical(snode) + + +def _safe_analytical(snode: BaseSchedulerNode) -> float: + try: + return snode.get_estimated_runtime() + except Exception: # noqa: BLE001 + return 0.0 + + +def _graph_has_free_symbols() -> bool: + """True if the graph being compiled still has symbolic (dynamic) shapes. + + Benchmarking real kernels is only safe when everything is concrete; a + dynamic-shape compile has free symbols and must use the analytical estimate.""" + try: + shape_env = V.graph.sizevars.shape_env + except Exception: # noqa: BLE001 + return False + if shape_env is None: + return False + try: + # A fully-static graph has no unbacked/free symbols with a non-singleton + # range. If ANY symbol has a range wider than one value and is not yet a + # constant replacement, the graph is dynamic and must not be benchmarked. + replacements = getattr(shape_env, "replacements", {}) + for sym, vr in shape_env.var_to_range.items(): + if sym in replacements: + continue # already specialized to a constant + lower, upper = vr.lower, vr.upper + # int_oo / unbounded upper -> definitely dynamic. Guard the compare. + try: + same = bool(lower == upper) + except Exception: # noqa: BLE001 + same = False + if not same: + return True + except Exception: # noqa: BLE001 + # Cannot prove static -> assume dynamic (safe: fall back to analytical). + return True + return False + + +def _suppress_guards(): + """Context manager that suppresses ShapeEnv guard creation while we run real + benchmark kernels at hinted shapes, so measurement never specializes a + dynamic dim. Falls back to a no-op when no ShapeEnv is active.""" + from contextlib import nullcontext + + try: + shape_env = V.graph.sizevars.shape_env + if shape_env is not None: + return shape_env.suppress_guards() + except Exception: # noqa: BLE001 + pass + return nullcontext() + + +# ShapeEnv mutable fields that benchmarking could pollute with a `s -> hint` +# specialization; snapshot/restore them so a measurement never persists a guard. +_SHAPEENV_STATE_FIELDS = ( + "guards", + "axioms", + "replacements", + "replacements_slocs", + "var_to_range", + "deferred_runtime_asserts", + "num_deferred_runtime_asserts", + "specializations", +) + + +class _shapeenv_sandbox: + """Snapshot the live ShapeEnv's specialization state on enter, restore it on + exit. Shallow-copies the mutable dict/list fields (their *contents* are + replaced wholesale by the compile, never mutated in place after we restore), + so benchmarking a kernel at a hinted concrete shape cannot leak an + ``Eq(sym, hint)`` replacement/guard into the real compile.""" + + def __init__(self) -> None: + self._env = None + self._saved: dict = {} + + def __enter__(self): + try: + self._env = V.graph.sizevars.shape_env + except Exception: # noqa: BLE001 + self._env = None + if self._env is None: + return self + import copy + + for f in _SHAPEENV_STATE_FIELDS: + if hasattr(self._env, f): + val = getattr(self._env, f) + try: + self._saved[f] = copy.copy(val) if isinstance(val, (dict, list, set)) else val + except Exception: # noqa: BLE001 + pass + return self + + def __exit__(self, *exc): + if self._env is None: + return False + for f, val in self._saved.items(): + try: + setattr(self._env, f, val) + except Exception: # noqa: BLE001 + pass + return False diff --git a/tests/feature_tests/fsdp_overlap_helper/__init__.py b/tests/feature_tests/fsdp_overlap_helper/__init__.py new file mode 100644 index 0000000..3eaa44a --- /dev/null +++ b/tests/feature_tests/fsdp_overlap_helper/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py b/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py new file mode 100644 index 0000000..3d550c3 --- /dev/null +++ b/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py @@ -0,0 +1,172 @@ +# Copyright (c) 2025 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. + +"""torchrun entrypoint: exercise the WHOLE FSDP-overlap chain end to end on a small +SimpleFSDP-style model, and print machine-checkable markers to stdout for the pytest +driver (tests/feature_tests/test_fsdp_overlap_e2e.py). + +Chain under test (magi_backend._apply_fsdp_fullgraph_overlap): + redistribute lowering -> per-region bucketing (coalesced) -> FsdpOverlapReorder + +Model: two Linear layers whose params are Shard(0) DTensors (via torchtitan +``data_parallel(mode="fully_shard")`` -- the same mechanism gaga4 uses), with an +opaque ``@magi_register_custom_op(is_subgraph_boundary=True)`` op between them so the +graph has a real subgraph boundary (delimits bucketing regions) and the weight +``prim_redistribute`` nodes the lowering pass rewrites. + +Run: torchrun --nproc_per_node=N tests/feature_tests/fsdp_overlap_helper/e2e_helper.py + [--cost-mode analytical|profile_sync] + +Markers printed on rank 0 (grepped by the test): + E2E_CONFIG world= cost_mode= + E2E_COMPILED (compiled callable produced, forward ran) + E2E_NUMERIC rel= ok= (compiled vs eager output) + E2E_PASS / E2E_FAIL +""" + +from __future__ import annotations + +import argparse +import os + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.device_mesh import init_device_mesh + +from magi_compiler import magi_compile, magi_register_custom_op +from magi_compiler.config import CompileMode, CudaGraphMode + + +# An opaque boundary op: forces a real subgraph boundary in the graph (delimits the +# bucketing regions). Elementwise so it is trivially correct. +@magi_register_custom_op(name="fsdp_e2e::boundary_gelu", is_subgraph_boundary=True) +def boundary_gelu(x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.gelu(x) + + +class Block(nn.Module): + """One transformer-ish block: two Linears with an opaque subgraph-boundary op + between them. Multiple blocks give the graph MULTIPLE boundary-delimited regions, + each with several weight all-gathers -- so the bucketing pass coalesces per region + and the reorder pass has real cross-region compute to hide gathers behind.""" + + def __init__(self, hidden: int): + super().__init__() + self.fc1 = nn.Linear(hidden, hidden, bias=False) + self.fc2 = nn.Linear(hidden, hidden, bias=False) + + def forward(self, x): + x = self.fc1(x) + x = boundary_gelu(x) + x = self.fc2(x) + return x + + +class TinyModel(nn.Module): + """A multi-layer stack of Blocks (default 4) -> many Shard(0) weights across + several regions, exercising the full lower->bucket->reorder chain at scale.""" + + def __init__(self, hidden: int, n_layers: int = 4): + super().__init__() + self.layers = nn.ModuleList(Block(hidden) for _ in range(n_layers)) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--cost-mode", default="analytical", choices=["analytical", "profile_sync", "profile"]) + ap.add_argument("--hidden", type=int, default=256) + ap.add_argument("--n-layers", type=int, default=4) + args = ap.parse_args() + + dist.init_process_group("cpu:gloo,cuda:nccl") + rank = dist.get_rank() + world = dist.get_world_size() + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", rank))) + dev = torch.cuda.current_device() + torch.manual_seed(0) + + os.environ["MAGI_COMPILE_FSDP_OVERLAP_COST_MODE"] = args.cost_mode + # Surface the backend's INFO logs (redistribute lowering / bucketing / reorder + # "repositioned N/M") to stderr so the pytest driver can assert the chain ran. + # These come back through torch's logging even from standalone_compile. + os.environ.setdefault("MAGI_LOGGING_LEVEL", "INFO") + + if rank == 0: + print(f"E2E_CONFIG world={world} cost_mode={args.cost_mode}", flush=True) + + from torchtitan.experiments.simple_fsdp.simple_fsdp import data_parallel + + mesh = init_device_mesh("cuda", (world,)) + + # --- eager reference (plain replicated weights) --- + hidden = args.hidden + ref = TinyModel(hidden, n_layers=args.n_layers).to(dev).to(torch.bfloat16) + x = torch.randn(64, hidden, device=dev, dtype=torch.bfloat16) + with torch.no_grad(): + eager_out = ref(x) + + # --- SimpleFSDP-sharded + magi_compile with the overlap chain --- + model = TinyModel(hidden, n_layers=args.n_layers).to(dev).to(torch.bfloat16) + # copy the reference weights so outputs are comparable (all layers) + with torch.no_grad(): + for (_, dst), (_, src) in zip(model.named_parameters(), ref.named_parameters()): + dst.copy_(src) + model = data_parallel(model, mesh, mode="fully_shard") + + def _patch(cfg): + cfg.compile_mode = CompileMode.MAGI_COMPILE + cfg.cudagraph_mode = CudaGraphMode.NONE + cfg.disable_graph_split = True + cfg.enable_fsdp_fullgraph_overlap = True + cfg.fsdp_fullgraph_bucket_mode = "coalesced" + return cfg + + # dim 0 of the input (token count) is the dynamic dim. + compiled = magi_compile(model, config_patch=_patch, dynamic_arg_dims={"x": 0}) + + with torch.no_grad(): + out = compiled(x) + torch.cuda.synchronize() + if rank == 0: + print("E2E_COMPILED", flush=True) + + # numeric check vs eager (bf16 tolerance; world=1 all_gather is identity, world>1 + # gathers the sharded weight back -> same math). + out_f = out.float() + ref_f = eager_out.float() + rel = ((out_f - ref_f).norm() / (ref_f.norm() + 1e-6)).item() + ok = bool(torch.isfinite(out_f).all().item()) and rel < 5e-2 + + # confirm every rank agrees (no rank-divergent hang/output) + ok_t = torch.tensor([1 if ok else 0], device=dev) + dist.all_reduce(ok_t) + all_ok = int(ok_t.item()) == world + + if rank == 0: + print(f"E2E_NUMERIC rel={rel:.5f} ok={ok}", flush=True) + print("E2E_PASS" if all_ok else "E2E_FAIL", flush=True) + + dist.barrier() + dist.destroy_process_group() + raise SystemExit(0 if all_ok else 1) + + +if __name__ == "__main__": + main() diff --git a/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py b/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py new file mode 100644 index 0000000..c4b53df --- /dev/null +++ b/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py @@ -0,0 +1,166 @@ +# Copyright (c) 2025 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. + +"""torchrun entrypoint: verify the profile_sync COLLECTIVE measurement accuracy on a +real multi-rank all-gather, for the pytest driver (test_profiling_estimator.py). + +profile_sync measures collectives with ``_measure_collective_op`` (barrier + fixed +iters, rank-lockstep). Here we: + 1. compile a tiny all_gather fn and capture its real _CollectiveKernel snode (via the + same Inductor reorder hook the pass uses); + 2. run ``_measure_collective_op`` on that snode -> the estimator's collective time; + 3. INDEPENDENTLY time the same all_gather with CUDA events (rank-lockstep); + 4. assert the estimate is within a tolerance band of the independent measurement. +Also exercises ``ProfilingRuntimeEstimator.warm_and_sync`` (the profile_sync entry) to +confirm it runs rank-lockstep without deadlock and reconciles a collective entry. + +Run: torchrun --nproc_per_node=2 .../estimator_collective_helper.py + +Markers (rank 0): + COLL_MEASURED est_us= real_us= ratio= + COLL_ACCURATE ok= + COLL_WARMSYNC reconciled= ok= + COLL_PASS / COLL_FAIL +""" + +from __future__ import annotations + +import os + +import torch +import torch._inductor.config as inductor_config +import torch.distributed as dist +from torch._inductor.utils import contains_collective + +from magi_compiler.profiling import ProfilingRuntimeEstimator +from magi_compiler.profiling.runtime_estimator import _measure_collective_op + +_AG = torch.ops._c10d_functional.all_gather_into_tensor.default +_WAIT = torch.ops._c10d_functional.wait_tensor.default + + +class _CaptureColl: + """Grab the collective snode(s) of one compiled graph.""" + + def __init__(self): + self.snodes = [] + + def __call__(self, snodes): + for s in snodes: + if contains_collective(s): + self.snodes.append(s) + return snodes + + +def _real_allgather_ns(shard, world, grp_name, iters=20, warmup=5): + """Independent rank-lockstep CUDA-event timing of all_gather(+wait).""" + + def fn(): + _WAIT(_AG(shard, world, grp_name)) + + for _ in range(warmup): + fn() + torch.cuda.synchronize() + dist.barrier() + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + s.record() + for _ in range(iters): + fn() + e.record() + torch.cuda.synchronize() + return (s.elapsed_time(e) / iters) * 1e6 # ns + + +def main() -> None: + dist.init_process_group("cpu:gloo,cuda:nccl") + rank = dist.get_rank() + world = dist.get_world_size() + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", rank))) + dev = torch.cuda.current_device() + grp_name = dist.group.WORLD.group_name + + numel = 4 * 1024 * 1024 # 4M bf16 = 8 MiB local shard -- big enough to time stably + shard = torch.randn(numel, device=dev, dtype=torch.bfloat16) + + # 1. capture the real all_gather snode from a compiled fn. + cap = _CaptureColl() + prev = ( + inductor_config.reorder_for_compute_comm_overlap, + inductor_config.reorder_for_compute_comm_overlap_passes, + inductor_config.force_disable_caches, + ) + inductor_config.reorder_for_compute_comm_overlap = True + inductor_config.reorder_for_compute_comm_overlap_passes = [cap] + inductor_config.force_disable_caches = True + try: + torch._dynamo.reset() + torch.compile(lambda s: _WAIT(_AG(s, world, grp_name)), dynamic=False)(shard) + torch.cuda.synchronize() + finally: + ( + inductor_config.reorder_for_compute_comm_overlap, + inductor_config.reorder_for_compute_comm_overlap_passes, + inductor_config.force_disable_caches, + ) = prev + + assert cap.snodes, "no collective snode captured" + coll_snode = cap.snodes[0] + + # 2. estimator's profile_sync collective measurement (barrier + fixed iters). + dist.barrier() + est_ns = _measure_collective_op(coll_snode) + dist.barrier() + + # 3. independent CUDA-event ground truth of the SAME all_gather. + real_ns = _real_allgather_ns(shard, world, grp_name) + dist.barrier() + + ratio = est_ns / real_ns if real_ns > 0 else 0.0 + # both measure the same op with fixed-iter CUDA events; expect close (allow + # 0.5-2x for iteration-count / warmup / noise differences). + accurate = est_ns > 0 and real_ns > 0 and 0.5 <= ratio <= 2.0 + + # 4. warm_and_sync: put the collective in the table (mimic the reorder warm-up), + # stash its snode, then reconcile across ranks -- must not deadlock. + est = ProfilingRuntimeEstimator() + est._sync_across_ranks = True + est(coll_snode) # __call__ seeds the table + stashes the snode (sync mode) + n_reconciled = est.warm_and_sync() + # the collective entry should now be marked measured (real value) + coll_entries = [e for e in est.table.values() if e.kind == "collective"] + warmsync_ok = len(coll_entries) >= 1 and all(e.measured for e in coll_entries) + + # gather agreement across ranks + ok_local = accurate and warmsync_ok + t = torch.tensor([1 if ok_local else 0], device=dev) + dist.all_reduce(t) + all_ok = int(t.item()) == world + + if rank == 0: + print(f"COLL_MEASURED est_us={est_ns/1e3:.1f} real_us={real_ns/1e3:.1f} ratio={ratio:.2f}", flush=True) + print(f"COLL_ACCURATE ok={accurate}", flush=True) + print(f"COLL_WARMSYNC reconciled={n_reconciled} ok={warmsync_ok}", flush=True) + print("COLL_PASS" if all_ok else "COLL_FAIL", flush=True) + rc = 0 if all_ok else 1 + else: + rc = 0 + + dist.barrier() + dist.destroy_process_group() + raise SystemExit(rc) + + +if __name__ == "__main__": + main() diff --git a/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py b/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py new file mode 100644 index 0000000..34d27e7 --- /dev/null +++ b/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py @@ -0,0 +1,121 @@ +# Copyright (c) 2025 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. + +"""torchrun entrypoint: drive ``FsdpOverlapReorder`` through a real ``torch.compile`` +on a tiny graph that contains upstream compute + a weight all-gather + a consumer, +and print markers for the pytest driver (test_fsdp_overlap_reorder.py). + +The reorder pass is an Inductor ``reorder_for_compute_comm_overlap_passes`` callback; +it needs a process group (its multi-rank-determinism warmup calls dist.get_rank()). +We build a fn: y = (x @ w0).relu() # upstream compute + g = all_gather(shard) ; wait # a weight gather to hoist + out = y + gathered_use # consumer after the compute +and wrap the pass so we can assert it RAN and returned a valid schedule. + +Run: torchrun --nproc_per_node=1 tests/feature_tests/fsdp_overlap_helper/reorder_helper.py + +Markers (rank 0): + REORDER_CALLED gathers= + REORDER_OK moved= (pass returned; N launches repositioned) + REORDER_FINITE ok= (compiled output finite + matches eager) + REORDER_PASS / REORDER_FAIL +""" + +from __future__ import annotations + +import os + +import torch +import torch._inductor.config as inductor_config +import torch.distributed as dist + +from magi_compiler.passes.fsdp_overlap import FsdpOverlapReorder +from magi_compiler.passes.fsdp_overlap import reorder as _ro + + +def main() -> None: + dist.init_process_group("cpu:gloo,cuda:nccl") + rank = dist.get_rank() + world = dist.get_world_size() + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", rank))) + dev = torch.cuda.current_device() + grp = dist.group.WORLD.group_name + torch.manual_seed(0) + + _AG = torch.ops._c10d_functional.all_gather_into_tensor.default + _WAIT = torch.ops._c10d_functional.wait_tensor.default + + H = 512 + w0 = torch.randn(H, H, device=dev, dtype=torch.bfloat16) + shard = torch.randn(H, H, device=dev, dtype=torch.bfloat16) + + def fn(x, w0, shard): + y = (x @ w0).relu() # upstream compute the gather can hide behind + g = _WAIT(_AG(shard, world, grp)) # weight all-gather + wait + gathered = g.reshape(world * H, H)[:H] # use the gathered weight + return y @ gathered + + # instrument the pass: count how many times it runs and how many launches move. + calls = {"n": 0, "gathers": 0, "moved": 0} + orig_call = FsdpOverlapReorder.__call__ + + def spy(self, snodes): + calls["n"] += 1 + calls["gathers"] += sum(1 for s in snodes if _ro._is_weight_gather(s)) + out = orig_call(self, snodes) + return out + + FsdpOverlapReorder.__call__ = spy + + reorder = FsdpOverlapReorder(slack_ns=5000.0) # default cost_fn (Inductor analytical) + prev_flag = inductor_config.reorder_for_compute_comm_overlap + prev_passes = inductor_config.reorder_for_compute_comm_overlap_passes + prev_cache = inductor_config.force_disable_caches + inductor_config.reorder_for_compute_comm_overlap = True + inductor_config.reorder_for_compute_comm_overlap_passes = [reorder] + inductor_config.force_disable_caches = True + try: + torch._dynamo.reset() + x = torch.randn(H, H, device=dev, dtype=torch.bfloat16) + eager = fn(x, w0, shard) + compiled = torch.compile(fn, dynamic=False) + out = compiled(x, w0, shard) + torch.cuda.synchronize() + finally: + inductor_config.reorder_for_compute_comm_overlap = prev_flag + inductor_config.reorder_for_compute_comm_overlap_passes = prev_passes + inductor_config.force_disable_caches = prev_cache + FsdpOverlapReorder.__call__ = orig_call + + finite = bool(torch.isfinite(out).all().item()) + rel = ((out.float() - eager.float()).norm() / (eager.float().norm() + 1e-6)).item() + numeric_ok = finite and rel < 5e-2 + + if rank == 0: + print(f"REORDER_CALLED gathers={calls['gathers']}", flush=True) + print(f"REORDER_OK ran={calls['n'] > 0}", flush=True) + print(f"REORDER_FINITE ok={numeric_ok} rel={rel:.5f}", flush=True) + ok = calls["n"] > 0 and calls["gathers"] >= 1 and numeric_ok + print("REORDER_PASS" if ok else "REORDER_FAIL", flush=True) + rc = 0 if ok else 1 + else: + rc = 0 + + dist.barrier() + dist.destroy_process_group() + raise SystemExit(rc) + + +if __name__ == "__main__": + main() diff --git a/tests/feature_tests/test_fsdp_collective_prefetch.py b/tests/feature_tests/test_fsdp_collective_prefetch.py deleted file mode 100644 index 139b00e..0000000 --- a/tests/feature_tests/test_fsdp_collective_prefetch.py +++ /dev/null @@ -1,229 +0,0 @@ -# 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. - -""" -Unit test for apply_fsdp_collective_prefetch. - -Builds a tiny FX graph that mirrors the gaga4 split structure: - - submod 0 (compute): a = x @ w0 - submod 1 (boundary): b = relu(a) # stands in for a custom-op boundary - submod 2 (compute): g = all_gather(w2) - wt = wait_tensor(g) - out = b @ wt # weight `w2` is used here - -The weight all-gather for `w2` lives in submod 2 alongside its use. With -distance=2 the pass should move ONLY the `all_gather_into_tensor` launch into -submod 0 (the previous *compute* submod, skipping the odd boundary submod), -while leaving `wait_tensor` and the matmul in submod 2. After split_module that -makes the collective launch during submod 0's compute and wait right before use. -""" - -import operator - -import torch -import torch.fx as fx - -from magi_compiler.passes.graph_split import ( - apply_fsdp_collective_prefetch, - bucket_weight_all_gather_coalesced_per_submod, -) - -_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default -_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default -_WAIT = torch.ops._c10d_functional.wait_tensor.default -_CAT = torch.ops.aten.cat.default -_SPLIT = torch.ops.aten.split_with_sizes.default - - -def _build_graph(): - g = fx.Graph() - x = g.placeholder("x") - w0 = g.placeholder("l_self_modules_proj0_parameters_weight_") - w2 = g.placeholder("l_self_modules_proj2_parameters_weight_") - - # submod 0: compute - a = g.call_function(torch.ops.aten.mm.default, (x, w0)) - # submod 1: boundary op (semantics irrelevant to the pass) - b = g.call_function(torch.ops.aten.relu.default, (a,)) - # submod 2: weight gather + wait + use - ag = g.call_function(_ALL_GATHER, (w2, 8, "0")) - wt = g.call_function(_WAIT, (ag,)) - out = g.call_function(torch.ops.aten.mm.default, (b, wt)) - g.output((out,)) - - gm = fx.GraphModule(torch.nn.Module(), g) - node_to_subgraph_id = {a: 0, b: 1, ag: 2, wt: 2, out: 2} - return gm, node_to_subgraph_id, {"a": a, "b": b, "ag": ag, "wt": wt, "out": out} - - -def test_launch_moves_wait_stays(): - gm, mapping, nodes = _build_graph() - - moved = apply_fsdp_collective_prefetch(gm, mapping, distance=2) - - assert moved == 1, f"expected exactly one launch moved, got {moved}" - # Launch relocated to the previous compute submod ... - assert mapping[nodes["ag"]] == 0, mapping[nodes["ag"]] - # ... but the wait and the real use stay at the consumer site. - assert mapping[nodes["wt"]] == 2, mapping[nodes["wt"]] - assert mapping[nodes["out"]] == 2, mapping[nodes["out"]] - - # Launch is tagged with provenance metadata. - assert nodes["ag"].meta.get("magi_fsdp_prefetch_to_subgraph") == 0 - assert nodes["ag"].meta.get("magi_fsdp_prefetch_for_consumer_subgraph") == 2 - - # The all_gather now precedes the first node of submod 0 in graph order. - order = list(gm.graph.nodes) - assert order.index(nodes["ag"]) < order.index(nodes["a"]), "launch must come before submod-0 compute" - assert order.index(nodes["wt"]) > order.index(nodes["b"]), "wait must remain after the boundary" - - # Graph stays valid. - gm.graph.lint() - - -def test_no_move_when_already_early(): - """If the consumer is in submod 0 there is no previous submod to move to.""" - gm, mapping, nodes = _build_graph() - # Pretend everything is in submod 0. - for n in nodes.values(): - mapping[n] = 0 - moved = apply_fsdp_collective_prefetch(gm, mapping, distance=2) - assert moved == 0 - assert mapping[nodes["ag"]] == 0 - - -def test_non_weight_all_gather_is_ignored(): - """An all_gather of a non-weight (activation) tensor must not be moved.""" - g = fx.Graph() - x = g.placeholder("x") # activation, not a weight - w0 = g.placeholder("l_self_modules_proj0_parameters_weight_") - a = g.call_function(torch.ops.aten.mm.default, (x, w0)) - b = g.call_function(torch.ops.aten.relu.default, (a,)) - ag = g.call_function(_ALL_GATHER, (b, 8, "0")) # gathers an activation - wt = g.call_function(_WAIT, (ag,)) - out = g.call_function(torch.ops.aten.relu.default, (wt,)) - g.output((out,)) - gm = fx.GraphModule(torch.nn.Module(), g) - mapping = {a: 0, b: 1, ag: 2, wt: 2, out: 2} - - moved = apply_fsdp_collective_prefetch(gm, mapping, distance=2) - assert moved == 0 - assert mapping[ag] == 2 - - -# --------------------------------------------------------------------------- # -# Coalesced bucketing: N same-submod weight gathers -> ONE -# all_gather_into_tensor_coalesced unpacked with operator.getitem (no cat, no -# split_with_sizes clone), one wait per member. -# --------------------------------------------------------------------------- # -def _build_coalesced_graph(world: int = 4): - """Two weight gathers in submod 2 (a compute submod after a boundary), - each off its own Shard(0) weight placeholder, plus a use of each.""" - from torch._subclasses.fake_tensor import FakeTensorMode - - fake = FakeTensorMode() - g = fx.Graph() - x = g.placeholder("x") - w0 = g.placeholder("l_self_modules_proj0_parameters_weight_") - wa = g.placeholder("l_self_modules_mlp_parameters_weight_") - wb = g.placeholder("l_self_modules_mlp_o_parameters_weight_") - - a = g.call_function(torch.ops.aten.mm.default, (x, w0)) # submod 0 compute - b = g.call_function(torch.ops.aten.relu.default, (a,)) # submod 1 boundary - # submod 2: two weight gathers + waits + uses - ag_a = g.call_function(_ALL_GATHER, (wa, world, "0")) - wt_a = g.call_function(_WAIT, (ag_a,)) - ag_b = g.call_function(_ALL_GATHER, (wb, world, "0")) - wt_b = g.call_function(_WAIT, (ag_b,)) - use_a = g.call_function(torch.ops.aten.mm.default, (b, wt_a)) - out = g.call_function(torch.ops.aten.mm.default, (use_a, wt_b)) - g.output((out,)) - - # Attach meta example_values (the bucket pass reads dtype + shapes). - with fake: - # local shards (chunk, rest) and gathered (world*chunk, rest) - loc_a = torch.empty((8, 16), dtype=torch.bfloat16) - loc_b = torch.empty((16, 8), dtype=torch.bfloat16) - wa.meta["example_value"] = loc_a - wb.meta["example_value"] = loc_b - ag_a.meta["example_value"] = loc_a.new_empty((world * 8, 16)) - ag_b.meta["example_value"] = loc_b.new_empty((world * 16, 8)) - wt_a.meta["example_value"] = ag_a.meta["example_value"] - wt_b.meta["example_value"] = ag_b.meta["example_value"] - - gm = fx.GraphModule(torch.nn.Module(), g) - node_to_subgraph_id = {a: 0, b: 1, ag_a: 2, wt_a: 2, ag_b: 2, wt_b: 2, use_a: 2, out: 2} - nodes = {"a": a, "b": b, "ag_a": ag_a, "ag_b": ag_b, "wt_a": wt_a, "wt_b": wt_b, "use_a": use_a, "out": out} - return gm, node_to_subgraph_id, nodes - - -def _count(gm, target): - return sum(1 for n in gm.graph.nodes if n.op == "call_function" and n.target is target) - - -def test_coalesced_bucket_unpacks_with_getitem(): - gm, mapping, _ = _build_coalesced_graph() - n = bucket_weight_all_gather_coalesced_per_submod(gm, mapping) - assert n == 1, f"expected one coalesced bucket, got {n}" - # exactly one coalesced launch, two getitem, two waits ... - assert _count(gm, _ALL_GATHER_COALESCED) == 1 - assert _count(gm, _WAIT) == 2 - assert sum(1 for x in gm.graph.nodes if x.op == "call_function" and x.target is operator.getitem) == 2 - # ... and NONE of the concat-path machinery (no cat, no split_with_sizes, no - # leftover per-member single all_gather). - assert _count(gm, _CAT) == 0 - assert _count(gm, _SPLIT) == 0 - assert _count(gm, _ALL_GATHER) == 0 - # launch + getitems all tagged / in submod 2; waits in submod 2 too. - coal = next(x for x in gm.graph.nodes if x.target is _ALL_GATHER_COALESCED) - assert coal.meta.get("magi_fsdp_weight_ag_coalesced") is True - assert mapping[coal] == 2 - gm.graph.lint() - - -def test_coalesced_launch_and_getitems_move_together(): - gm, mapping, _ = _build_coalesced_graph() - bucket_weight_all_gather_coalesced_per_submod(gm, mapping) - moved = apply_fsdp_collective_prefetch(gm, mapping, distance=2) - assert moved == 1, f"expected the coalesced launch moved once, got {moved}" - - coal = next(x for x in gm.graph.nodes if x.target is _ALL_GATHER_COALESCED) - gis = [x for x in gm.graph.nodes if x.op == "call_function" and x.target is operator.getitem] - waits = [x for x in gm.graph.nodes if x.op == "call_function" and x.target is _WAIT] - - # launch + both getitems hoisted into submod 0 (the previous compute submod) ... - assert mapping[coal] == 0 - for gi in gis: - assert mapping[gi] == 0, f"getitem must move with the coalesced launch, got {mapping[gi]}" - # ... but the waits stay at the use site in submod 2. - for wt in waits: - assert mapping[wt] == 2, f"wait must stay at consumer, got {mapping[wt]}" - - # The list (coalesced launch) and its getitems never become a submod boundary - # edge: only the getitem outputs cross. Order check: launch precedes submod-0 - # compute. - order = list(gm.graph.nodes) - a = next(x for x in gm.graph.nodes if x.op == "call_function" and x.target is torch.ops.aten.mm.default) - assert order.index(coal) < order.index(a) - gm.graph.lint() - - -if __name__ == "__main__": - test_launch_moves_wait_stays() - test_no_move_when_already_early() - test_non_weight_all_gather_is_ignored() - test_coalesced_bucket_unpacks_with_getitem() - test_coalesced_launch_and_getitems_move_together() - print("PASS: fsdp_collective_prefetch unit tests") diff --git a/tests/feature_tests/test_fsdp_overlap_bucket.py b/tests/feature_tests/test_fsdp_overlap_bucket.py new file mode 100644 index 0000000..19e4eb8 --- /dev/null +++ b/tests/feature_tests/test_fsdp_overlap_bucket.py @@ -0,0 +1,204 @@ +# Copyright (c) 2025 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. + +"""Unit tests for the FSDP-overlap weight all-gather bucketing pass +(``magi_compiler.passes.fsdp_overlap.bucket_all_gather`` + +``lower_and_bucket_full_graph``). + +Pure-CPU: the bucketing pass operates on ``all_gather_into_tensor`` fx nodes tagged +``magi_fsdp_weight_ag`` (it never runs the ops), so we drive it with SYNTHETIC fx +graphs built with meta-tensor ``example_value``s -- no DTensor / distributed / GPU. +""" + +from collections import Counter + +import pytest +import torch +import torch.fx as fx + +from magi_compiler.passes.fsdp_overlap import bucket_weight_all_gather_coalesced_per_region, lower_and_bucket_full_graph + +_AG = torch.ops._c10d_functional.all_gather_into_tensor.default +_AG_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default +_WAIT = torch.ops._c10d_functional.wait_tensor.default + + +# --------------------------------------------------------------------------- +# synthetic graph builder +# --------------------------------------------------------------------------- +def _build_ag_graph(specs, world=2, group="grp0"): + """Build an fx graph of independent weight all-gathers. + + ``specs``: list of dicts, each ``{shape, dtype, boundary_after?}``. For each spec + we emit weight_shard placeholder -> all_gather_into_tensor (tagged + magi_fsdp_weight_ag) -> wait_tensor. A spec with ``boundary`` inserts an opaque + boundary op (aten.relu here) between gathers so ``_build_region_sid_map`` splits + regions. ALL placeholders are declared first (as in a real traced graph) so the + hoisted coalesced launch stays topologically valid. + + Returns (gm, boundary_op_or_None). + """ + g = fx.Graph() + locs = [] + for i, s in enumerate(specs): + loc = g.placeholder(f"w{i}_weight_shard") + loc.meta["example_value"] = torch.empty(*s["shape"], dtype=s["dtype"], device="meta") + locs.append(loc) + + outs = [] + for i, (s, loc) in enumerate(zip(specs, locs)): + if s.get("boundary_before"): + # opaque boundary op consuming the previous wait (keeps it in the graph) + b = g.call_function(torch.ops.aten.relu.default, (outs[-1],)) if outs else None + if b is not None: + b.meta["example_value"] = outs[-1].meta["example_value"] + chunk = s["shape"][0] + rest = s["shape"][1:] + gathered = torch.empty(chunk * world, *rest, dtype=s["dtype"], device="meta") + ag = g.call_function(_AG, (loc, world, group)) + ag.meta["example_value"] = gathered + ag.meta["magi_fsdp_weight_ag"] = True + w = g.call_function(_WAIT, (ag,)) + w.meta["example_value"] = gathered + outs.append(w) + g.output(tuple(outs)) + gm = fx.GraphModule(torch.nn.Module(), g) + return gm + + +def _op_counts(gm) -> Counter: + return Counter(str(n.target) for n in gm.graph.nodes if n.op == "call_function") + + +def _n(gm, target) -> int: + return sum(1 for n in gm.graph.nodes if n.op == "call_function" and n.target is target) + + +# --------------------------------------------------------------------------- +# bucket_weight_all_gather_coalesced_per_region +# --------------------------------------------------------------------------- +def test_coalesce_same_region_merges_to_one(): + """N same-(region,group,dtype) gathers -> ONE coalesced + N getitems + N waits.""" + gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 3) + sid_map = {n: 0 for n in gm.graph.nodes} + n_buckets = bucket_weight_all_gather_coalesced_per_region(gm, sid_map) + + assert n_buckets == 1 + assert _n(gm, _AG_COALESCED) == 1 + assert _n(gm, _AG) == 0 # all plain gathers replaced + assert _n(gm, _WAIT) == 3 # one wait per member, preserved + assert sum(1 for n in gm.graph.nodes if n.op == "call_function" and "getitem" in str(n.target)) == 3 + gm.graph.lint() # topologically valid + + +def test_single_gather_not_coalesced(): + """A lone weight gather has nothing to coalesce -> untouched, 0 buckets.""" + gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}]) + sid_map = {n: 0 for n in gm.graph.nodes} + n_buckets = bucket_weight_all_gather_coalesced_per_region(gm, sid_map) + assert n_buckets == 0 + assert _n(gm, _AG) == 1 + assert _n(gm, _AG_COALESCED) == 0 + + +def test_dtype_change_breaks_bucket(): + """bf16, bf16, fp32, bf16 (program order) -> {bf16,bf16}, {fp32 alone}, {bf16 alone} + => only the leading same-dtype run of >=2 coalesces -> 1 bucket.""" + gm = _build_ag_graph( + [ + {"shape": (4, 8), "dtype": torch.bfloat16}, + {"shape": (4, 8), "dtype": torch.bfloat16}, + {"shape": (4, 8), "dtype": torch.float32}, + {"shape": (4, 8), "dtype": torch.bfloat16}, + ] + ) + sid_map = {n: 0 for n in gm.graph.nodes} + n_buckets = bucket_weight_all_gather_coalesced_per_region(gm, sid_map) + assert n_buckets == 1 # only the {bf16,bf16} run + assert _n(gm, _AG_COALESCED) == 1 + assert _n(gm, _AG) == 2 # the fp32 and the trailing bf16 stay individual + + +def test_bucket_size_bytes_caps_run(): + """Same dtype, but bucket_size_bytes forces a split. 4x8 bf16 shard = 512 B each; + cap at 512 B means each gather is its own bucket -> no coalescing (0 buckets>=2).""" + gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 4) + sid_map = {n: 0 for n in gm.graph.nodes} + # 4*8*2 = 64 B local shard each; cap at 64 B => each starts a new bucket -> singletons + n_buckets = bucket_weight_all_gather_coalesced_per_region(gm, sid_map, bucket_size_bytes=64) + assert n_buckets == 0 + assert _n(gm, _AG) == 4 + + # cap at 128 B => 2 shards per bucket => 2 buckets of 2 + gm2 = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 4) + sid2 = {n: 0 for n in gm2.graph.nodes} + n2 = bucket_weight_all_gather_coalesced_per_region(gm2, sid2, bucket_size_bytes=128) + assert n2 == 2 + assert _n(gm2, _AG_COALESCED) == 2 + + +def test_different_regions_not_merged(): + """Gathers in different sids (regions) never share a bucket.""" + gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 4) + ags = [n for n in gm.graph.nodes if n.op == "call_function" and n.target is _AG] + # put first 2 gathers (+their locals/waits) in region 0, last 2 in region 2 + sid_map = {} + for n in gm.graph.nodes: + sid_map[n] = 0 + for ag in ags[2:]: + sid_map[ag] = 2 + n_buckets = bucket_weight_all_gather_coalesced_per_region(gm, sid_map) + assert n_buckets == 2 # one per region + assert _n(gm, _AG_COALESCED) == 2 + + +# --------------------------------------------------------------------------- +# lower_and_bucket_full_graph (entry point) +# --------------------------------------------------------------------------- +def test_lower_and_bucket_mode_none_returns_zero(): + """mode 'none' -> lowering only, no bucketing -> 0 buckets, gathers untouched.""" + gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 3) + n = lower_and_bucket_full_graph(gm, "none") + assert n == 0 + assert _n(gm, _AG) == 3 + assert _n(gm, _AG_COALESCED) == 0 + + +def test_lower_and_bucket_mode_coalesced(): + """mode 'coalesced' with no boundary_ops -> all gathers in one region -> 1 bucket.""" + gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 3) + n = lower_and_bucket_full_graph(gm, "coalesced") + assert n == 1 + assert _n(gm, _AG_COALESCED) == 1 + + +def test_lower_and_bucket_boundary_ops_split_regions(): + """A boundary op between gathers splits them into separate buckets.""" + gm = _build_ag_graph( + [ + {"shape": (4, 8), "dtype": torch.bfloat16}, + {"shape": (4, 8), "dtype": torch.bfloat16}, + {"shape": (4, 8), "dtype": torch.bfloat16, "boundary_before": True}, + {"shape": (4, 8), "dtype": torch.bfloat16}, + ] + ) + n = lower_and_bucket_full_graph(gm, "coalesced", boundary_ops=[torch.ops.aten.relu.default]) + assert n == 2 # {g0,g1} before boundary, {g2,g3} after + assert _n(gm, _AG_COALESCED) == 2 + + +def test_lower_and_bucket_unknown_mode_raises(): + gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 2) + with pytest.raises(ValueError, match="expected 'none' or 'coalesced'"): + lower_and_bucket_full_graph(gm, "concat") diff --git a/tests/feature_tests/test_fsdp_overlap_e2e.py b/tests/feature_tests/test_fsdp_overlap_e2e.py new file mode 100644 index 0000000..fd6ac0f --- /dev/null +++ b/tests/feature_tests/test_fsdp_overlap_e2e.py @@ -0,0 +1,95 @@ +# Copyright (c) 2025 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. + +"""End-to-end FSDP-overlap chain test: a small SimpleFSDP-sharded model through +``magi_compile`` with ``enable_fsdp_fullgraph_overlap`` -- exercising the whole +chain (redistribute lowering -> bucketing -> FsdpOverlapReorder) in one real compile. + +Driven via a ``torchrun`` subprocess helper (fsdp_overlap_helper/e2e_helper.py); we +assert on stdout markers + returncode. + +- SINGLE rank (world=1): chain runs, compiles, output matches eager. (At world=1 the + weights are Replicate() so there is nothing to gather -- this only proves the chain + is wired and does not break the compile.) +- MULTI rank (world=2, needs >=2 GPUs): weights are real Shard(0) DTensors, so the + graph has weight all-gathers -- the reorder pass actually repositions them (asserted + via the backend's "repositioned N/M" INFO log). This is the true full-chain test. +""" + +import os +import shutil +import subprocess +from pathlib import Path + +import pytest +import torch + +_HELPER = Path(__file__).parent / "fsdp_overlap_helper" / "e2e_helper.py" + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +requires_torchrun = pytest.mark.skipif(shutil.which("torchrun") is None, reason="requires torchrun") + + +def _run(nproc: int, cost_mode: str, port: str) -> subprocess.CompletedProcess: + env = os.environ.copy() + env["MAGI_LOGGING_LEVEL"] = "info" # so the backend's chain INFO logs are captured + return subprocess.run( + ["torchrun", f"--nproc_per_node={nproc}", f"--master_port={port}", str(_HELPER), "--cost-mode", cost_mode], + env=env, + capture_output=True, + text=True, + timeout=900, + ) + + +@requires_cuda +@requires_torchrun +def test_e2e_single_rank(): + """world=1: whole overlap chain compiles + runs + output matches eager.""" + p = _run(1, "analytical", "29641") + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-4000:]}" + assert "E2E_COMPILED" in p.stdout, out[-4000:] + assert "E2E_PASS" in p.stdout, out[-4000:] + # the chain was actually invoked + assert "FSDP fullgraph overlap" in out, out[-4000:] + + +@requires_cuda +@requires_torchrun +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") +def test_e2e_multi_rank_full_chain(): + """world=2: real Shard(0) weights -> the reorder repositions real weight + all-gathers. Asserts the full chain ran (no deadlock, output correct, and the + reorder pass touched >=1 gather).""" + p = _run(2, "analytical", "29642") + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-4000:]}" + assert "E2E_PASS" in p.stdout, out[-4000:] + # the reorder pass ran on a graph WITH weight all-gathers (Shard(0) at world>1) + assert "FSDP overlap reorder: repositioned" in out, ( + "reorder pass did not run on any weight all-gather at world=2\n" + out[-4000:] + ) + + +@requires_cuda +@requires_torchrun +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") +def test_e2e_multi_rank_profile_sync(): + """world=2 with profile_sync cost mode: exercises warm_and_sync's rank-lockstep + real measurement path end to end (the gaga4 default).""" + p = _run(2, "profile_sync", "29643") + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-4000:]}" + assert "E2E_PASS" in p.stdout, out[-4000:] diff --git a/tests/feature_tests/test_fsdp_overlap_lowering.py b/tests/feature_tests/test_fsdp_overlap_lowering.py new file mode 100644 index 0000000..9374f18 --- /dev/null +++ b/tests/feature_tests/test_fsdp_overlap_lowering.py @@ -0,0 +1,144 @@ +# Copyright (c) 2025 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. + +"""Unit tests for the SimpleFSDP weight redistribute lowering pass +(``magi_compiler.passes.fsdp_overlap.lower_prim_redistribute_to_collectives``). + +The pass matches ``prim_redistribute`` + ``prim_to_local`` fx nodes (by +``target.__name__``) whose input is a weight placeholder carrying a ``Shard(0)`` +DTensor ``example_value``, and rewrites them into explicit +``all_gather_into_tensor`` + ``wait_tensor``. ``prim_redistribute`` is a +torch.compile-internal prim that can't be constructed directly, so we build a +minimal synthetic graph with plain functions named ``prim_redistribute`` / +``prim_to_local`` and REAL 1-rank DTensor metas (this drives the exact matching / +rewrite logic without depending on Dynamo capturing the prim). + +Uses a 1-rank process group + device mesh (GPU required). +""" + +import os + +import pytest +import torch +import torch.fx as fx + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.fixture(scope="module") +def dist_1rank(): + """A single-rank process group + cuda device mesh (module-scoped).""" + import torch.distributed as dist + + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29661") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + created = False + if not dist.is_initialized(): + dist.init_process_group("gloo") + created = True + torch.cuda.set_device(0) + from torch.distributed.device_mesh import init_device_mesh + + mesh = init_device_mesh("cuda", (1,)) + yield mesh + if created: + dist.destroy_process_group() + + +# named functions whose __name__ is EXACTLY "prim_redistribute" / "prim_to_local" +# (the pass matches node.target.__name__ against those strings). +def _make_prim(name): + def f(x): + return x + + f.__name__ = name + return f + + +prim_redistribute = _make_prim("prim_redistribute") +prim_to_local = _make_prim("prim_to_local") + + +def _build_redistribute_graph(mesh, weight_name, dtype=torch.bfloat16, rows=8, cols=4): + """weight(Shard0 DTensor) -> prim_redistribute(Replicate) -> prim_to_local -> out.""" + from torch.distributed.tensor import Replicate, Shard, distribute_tensor + + full = torch.randn(rows, cols, device="cuda", dtype=dtype) + sharded = distribute_tensor(full, mesh, [Shard(0)]) + replicated = distribute_tensor(full, mesh, [Replicate()]) + + g = fx.Graph() + w = g.placeholder(weight_name) + w.meta["example_value"] = sharded + rd = g.call_function(prim_redistribute, (w,)) + rd.meta["example_value"] = replicated + tl = g.call_function(prim_to_local, (rd,)) + tl.meta["example_value"] = replicated._local_tensor + g.output((tl,)) + return fx.GraphModule(torch.nn.Module(), g) + + +_AG = torch.ops._c10d_functional.all_gather_into_tensor.default +_WAIT = torch.ops._c10d_functional.wait_tensor.default + + +def _targets(gm): + return [n.target for n in gm.graph.nodes if n.op == "call_function"] + + +@requires_cuda +def test_lowering_rewrites_shard0_weight(dist_1rank): + from magi_compiler.passes.fsdp_overlap import lower_prim_redistribute_to_collectives + + gm = _build_redistribute_graph(dist_1rank, "model_fc1_weight_parameter") + n = lower_prim_redistribute_to_collectives(gm) + + assert n == 1, "one Shard(0) weight redistribute should be lowered" + targets = _targets(gm) + assert _AG in targets + assert _WAIT in targets + # the prims are gone, replaced by explicit collectives + assert prim_redistribute not in targets and prim_to_local not in targets + # the gather is tagged so the bucketing pass can find it + tagged = [x for x in gm.graph.nodes if x.meta.get("magi_fsdp_weight_ag")] + assert len(tagged) == 1 + + +@requires_cuda +def test_lowering_skips_non_weight_input(dist_1rank): + """A redistribute whose input placeholder is NOT weight-named is left untouched.""" + from magi_compiler.passes.fsdp_overlap import lower_prim_redistribute_to_collectives + + gm = _build_redistribute_graph(dist_1rank, "some_activation") # not weight/param/bias + n = lower_prim_redistribute_to_collectives(gm) + assert n == 0 + targets = _targets(gm) + assert _AG not in targets + assert prim_redistribute in targets # untouched + + +@requires_cuda +def test_lowering_gather_and_wait_are_separate_nodes(dist_1rank): + """Launch and wait must be DISTINCT nodes (so the reorder can move the launch + independently of the wait) -- the whole point of lowering.""" + from magi_compiler.passes.fsdp_overlap import lower_prim_redistribute_to_collectives + + gm = _build_redistribute_graph(dist_1rank, "layer_weight") + lower_prim_redistribute_to_collectives(gm) + ag = [n for n in gm.graph.nodes if n.op == "call_function" and n.target is _AG] + wait = [n for n in gm.graph.nodes if n.op == "call_function" and n.target is _WAIT] + assert len(ag) == 1 and len(wait) == 1 + assert wait[0].args[0] is ag[0] # wait consumes the gather's output diff --git a/tests/feature_tests/test_fsdp_overlap_reorder.py b/tests/feature_tests/test_fsdp_overlap_reorder.py new file mode 100644 index 0000000..42a8adb --- /dev/null +++ b/tests/feature_tests/test_fsdp_overlap_reorder.py @@ -0,0 +1,70 @@ +# Copyright (c) 2025 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. + +"""FsdpOverlapReorder integration test. + +The reorder pass is an Inductor scheduler callback that needs a process group + a +real ``torch.compile``, so we drive it via a ``torchrun`` subprocess helper +(fsdp_overlap_helper/reorder_helper.py) and assert on its stdout markers -- same +subprocess pattern as test_autograd_function_cache_flag.py. +""" + +import os +import shutil +import subprocess +from pathlib import Path + +import pytest +import torch + +_HELPER = Path(__file__).parent / "fsdp_overlap_helper" / "reorder_helper.py" + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +requires_torchrun = pytest.mark.skipif(shutil.which("torchrun") is None, reason="requires torchrun") + + +def _run(nproc: int) -> subprocess.CompletedProcess: + env = os.environ.copy() + env["MAGI_LOGGING_LEVEL"] = env.get("MAGI_LOGGING_LEVEL", "info") + return subprocess.run( + ["torchrun", f"--nproc_per_node={nproc}", "--master_port=29631", str(_HELPER)], + env=env, + capture_output=True, + text=True, + timeout=600, + ) + + +@requires_cuda +@requires_torchrun +def test_reorder_single_rank(): + """world=1: the reorder pass runs inside a real compile, sees a weight gather, + reorders without error, and the compiled output matches eager.""" + p = _run(1) + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" + assert "REORDER_CALLED gathers=1" in p.stdout, out[-3000:] + assert "REORDER_OK ran=True" in p.stdout, out[-3000:] + assert "REORDER_PASS" in p.stdout, out[-3000:] + + +@requires_cuda +@requires_torchrun +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") +def test_reorder_multi_rank(): + """world=2: same, with a real 2-rank all-gather group.""" + p = _run(2) + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" + assert "REORDER_PASS" in p.stdout, out[-3000:] diff --git a/tests/feature_tests/test_profiling_estimator.py b/tests/feature_tests/test_profiling_estimator.py new file mode 100644 index 0000000..8de54db --- /dev/null +++ b/tests/feature_tests/test_profiling_estimator.py @@ -0,0 +1,373 @@ +# Copyright (c) 2025 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. + +"""Unit tests for the profiling runtime estimator +(``magi_compiler.profiling.runtime_estimator``): the arg-realizer, the cache-key +shape helper, the estimator's memoization/deepcopy, and the extern replay measure. +""" + +import copy +import os + +import pytest +import torch +import torch.fx as fx +from torch.fx.immutable_collections import immutable_list + +from magi_compiler.profiling import ProfilingRuntimeEstimator +from magi_compiler.profiling.runtime_estimator import ProfileEntry, _measure_extern, _realize_arg, _static + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +# --------------------------------------------------------------------------- +# _realize_arg -- fx args -> concrete replay inputs +# --------------------------------------------------------------------------- +def _node_with_val(val): + """A bare fx.Node carrying ``val`` as meta['val'] (no graph needed for realize).""" + g = fx.Graph() + n = g.placeholder("p") + n.meta["val"] = val + return n + + +@requires_cuda +def test_realize_tensor_node_builds_matching_tensor(): + dev = torch.cuda.current_device() + node = _node_with_val(torch.empty(4, 8, device=dev, dtype=torch.bfloat16)) + out = _realize_arg(node) + assert isinstance(out, torch.Tensor) + assert out.shape == (4, 8) + assert out.dtype == torch.bfloat16 + assert out.device.type == "cuda" + + +def test_realize_scalar_int_node(): + node = _node_with_val(16) # a Node carrying a plain int scalar + assert _realize_arg(node) == 16 + + +def test_realize_immutable_list_of_int_nodes_to_plain_list(): + """The grouped-linear m_splits case: immutable_list of Node(int) -> plain list[int] + (the strict List[int] C++ check rejects immutable_list / symbolic elements).""" + elems = immutable_list([_node_with_val(16), _node_with_val(16), _node_with_val(16)]) + out = _realize_arg(elems) + assert type(out) is list # NOT immutable_list + assert out == [16, 16, 16] + assert all(type(x) is int for x in out) + + +def test_realize_plain_containers_preserved(): + assert _realize_arg([1, 2, 3]) == [1, 2, 3] + assert type(_realize_arg((1, 2))) is tuple + assert _realize_arg({"a": 1}) == {"a": 1} + # non-node scalars pass through + assert _realize_arg(3.5) == 3.5 + + +# --------------------------------------------------------------------------- +# _static -- cache-key shape (symbolic dims stringified, never int()'d) +# --------------------------------------------------------------------------- +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()) + + +# --------------------------------------------------------------------------- +# ProfilingRuntimeEstimator -- memoization + deepcopy +# --------------------------------------------------------------------------- +def test_table_property_and_initial_state(): + est = ProfilingRuntimeEstimator() + assert est.table == {} + assert est.n_measured == 0 + assert est.n_cache_hits == 0 + assert est._sync_across_ranks is False + + +def test_deepcopy_returns_clean_instance(): + """Inductor deep-copies the pass list into the fx-graph cache key; the estimator's + __deepcopy__ must return a fresh instance WITHOUT the (FakeTensor-holding) table.""" + est = ProfilingRuntimeEstimator() + est._sync_across_ranks = True + est._table[("k",)] = ProfileEntry(ns=123.0, kind="extern", label="x", measured=True) + est._key_snode[("k",)] = object() + + clone = copy.deepcopy(est) + assert isinstance(clone, ProfilingRuntimeEstimator) + assert clone._table == {} # transient state dropped + assert clone._key_snode == {} + assert clone._sync_across_ranks is True # flag preserved + + +def test_profile_entry_fields(): + e = ProfileEntry(ns=5.0, kind="collective", label="all_gather(...)", measured=False) + assert e.ns == 5.0 and e.kind == "collective" and e.measured is False + assert e.reuse_count == 0 + + +# --------------------------------------------------------------------------- +# _measure_extern -- eager replay timing of an extern op +# --------------------------------------------------------------------------- +class _FakeIR: + def __init__(self, fx_node): + self._fx = fx_node + + def get_origin_node(self): + return self._fx + + +class _FakeSnode: + def __init__(self, fx_node): + self.node = _FakeIR(fx_node) + + def get_name(self): + return "op_test" + + +@requires_cuda +def test_measure_extern_matmul_positive(): + """_measure_extern replays a plain aten.mm on real tensors and returns >0 ns.""" + dev = torch.cuda.current_device() + g = fx.Graph() + a = g.placeholder("a") + b = g.placeholder("b") + a.meta["val"] = torch.empty(512, 512, device=dev, dtype=torch.bfloat16) + b.meta["val"] = torch.empty(512, 512, device=dev, dtype=torch.bfloat16) + mm = g.call_function(torch.ops.aten.mm.default, (a, b)) + mm.meta["val"] = torch.empty(512, 512, device=dev, dtype=torch.bfloat16) + + ns = _measure_extern(_FakeSnode(mm), fixed_iters=False) + assert ns > 0.0 + + +import torch._inductor.config as inductor_config # noqa: E402 + +# =========================================================================== +# ACCURACY: estimator vs INDEPENDENT CUDA-event ground truth, for the op +# categories the FSDP-overlap reorder actually meets -- torch-native matmul, +# a fused elementwise chain (FusedSchedulerNode), and a user @triton.jit op. +# +# The estimator's fused-Triton path calls V.graph.scheduler.benchmark_fused_nodes, +# valid only INSIDE Inductor codegen with a live V.graph, so we run the estimator +# from a capture callable installed on the SAME hook the real pass uses +# (reorder_for_compute_comm_overlap_passes) -- it receives the real scheduler nodes. +# The hook calls dist.get_rank(), so a (1-rank) process group is required. +# =========================================================================== +import triton # noqa: E402 +import triton.language as tl # noqa: E402 +from torch._inductor.scheduler import ExternKernelSchedulerNode # noqa: E402 +from torch._inductor.utils import contains_collective, contains_wait # noqa: E402 + + +@pytest.fixture(scope="module") +def pg_1rank(): + """Single-rank process group -- the Inductor reorder hook calls dist.get_rank().""" + import torch.distributed as dist + + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29671") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + created = False + if not dist.is_initialized(): + dist.init_process_group("gloo") + created = True + torch.cuda.set_device(0) + yield + if created: + dist.destroy_process_group() + + +def _cuda_time_ns(fn, warmup=10, iters=50): + """Median per-call GPU time (ns) via CUDA events -- independent ground truth + (NOT the estimator's own benchmark helpers).""" + import statistics + + for _ in range(warmup): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(iters): + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + s.record() + fn() + e.record() + torch.cuda.synchronize() + samples.append(s.elapsed_time(e) * 1e6) + return statistics.median(samples) + + +class _Capture: + """Runs a fresh estimator on every snode of one compiled graph; records the + (kind, est_ns) rows so a test can pick the op it cares about.""" + + def __init__(self): + self.est = ProfilingRuntimeEstimator() + self.rows = [] # (kind, est_ns) + + def __call__(self, snodes): + for s in snodes: + if contains_wait(s) and not contains_collective(s): + continue + kind = ( + "collective" if contains_collective(s) else "extern" if isinstance(s, ExternKernelSchedulerNode) else "fused" + ) + self.rows.append((kind, float(self.est(s)))) + return snodes + + def best(self, kind): + vals = [ns for k, ns in self.rows if k == kind and ns > 0] + return max(vals) if vals else None + + +def _compile_and_capture(fn, *args): + cap = _Capture() + prev = ( + inductor_config.reorder_for_compute_comm_overlap, + inductor_config.reorder_for_compute_comm_overlap_passes, + inductor_config.force_disable_caches, + ) + inductor_config.reorder_for_compute_comm_overlap = True + inductor_config.reorder_for_compute_comm_overlap_passes = [cap] + inductor_config.force_disable_caches = True # else a cached graph skips the scheduler + try: + torch._dynamo.reset() + torch.compile(fn, dynamic=False)(*args) + torch.cuda.synchronize() + finally: + ( + inductor_config.reorder_for_compute_comm_overlap, + inductor_config.reorder_for_compute_comm_overlap_passes, + inductor_config.force_disable_caches, + ) = prev + return cap + + +# user @triton.jit kernel wrapped as a torch custom op (an extern snode). +@triton.jit +def _add_mul_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + x = tl.load(x_ptr + offs, mask=mask) + y = tl.load(y_ptr + offs, mask=mask) + tl.store(out_ptr + offs, x * y + x, mask=mask) + + +@torch.library.custom_op("profiling_test::tri_add_mul", mutates_args=()) +def _tri_add_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x) + n = x.numel() + _add_mul_kernel[(triton.cdiv(n, 1024),)](x, y, out, n, BLOCK=1024) + return out + + +@_tri_add_mul.register_fake +def _(x, y): + return torch.empty_like(x) + + +def _assert_within(est, truth, lo=0.5, hi=2.0): + assert est is not None and est > 0, "op was not measured (est is None/0)" + assert truth > 0 + ratio = est / truth + assert lo <= ratio <= hi, f"estimate {est/1e3:.1f}us vs truth {truth/1e3:.1f}us = {ratio:.2f}x (want {lo}-{hi}x)" + + +@requires_cuda +def test_accuracy_matmul(pg_1rank): + """Extern matmul: estimator (aten replay) within 0.5-2x of CUDA-event truth.""" + dev = torch.cuda.current_device() + a = torch.randn(4096, 4096, device=dev, dtype=torch.bfloat16) + b = torch.randn(4096, 4096, device=dev, dtype=torch.bfloat16) + cap = _compile_and_capture(lambda a, b: a @ b, a, b) + _assert_within(cap.best("extern"), _cuda_time_ns(lambda: torch.mm(a, b))) + + +@requires_cuda +def test_accuracy_elementwise_fusion(pg_1rank): + """Fused elementwise chain -> ONE FusedSchedulerNode; estimator + (benchmark_fused_nodes) within 0.5-2x of the compiled kernel's CUDA-event time.""" + dev = torch.cuda.current_device() + x = torch.randn(8192, 8192, device=dev, dtype=torch.bfloat16) + + def deep_pointwise(x): + y = x + for _ in range(32): + y = torch.sin(y) * 1.001 + torch.cos(y) + return y + + cap = _compile_and_capture(deep_pointwise, x) + est = cap.best("fused") + # ground truth: the whole compiled graph is ~one fused kernel. + torch._dynamo.reset() + cfn = torch.compile(deep_pointwise, dynamic=False) + cfn(x) + _assert_within(est, _cuda_time_ns(lambda: cfn(x))) + + +@requires_cuda +def test_accuracy_custom_triton(pg_1rank): + """User @triton.jit op (extern): estimator replay within 0.5-2x of truth.""" + dev = torch.cuda.current_device() + x = torch.randn(8192, 8192, device=dev, dtype=torch.float32) + y = torch.randn(8192, 8192, device=dev, dtype=torch.float32) + cap = _compile_and_capture(lambda x, y: _tri_add_mul(x, y) + 1.0, x, y) + _assert_within(cap.best("extern"), _cuda_time_ns(lambda: _tri_add_mul(x, y))) + + +# =========================================================================== +# MULTI-CARD: profile_sync collective measurement accuracy. Driven via a +# torchrun subprocess helper (the collective measurement needs >=2 real ranks). +# The helper measures a real all-gather with _measure_collective_op and compares +# to an independent CUDA-event timing, asserting the estimate is within tolerance, +# and checks warm_and_sync reconciles rank-lockstep without deadlock. +# =========================================================================== +import shutil # noqa: E402 +import subprocess # noqa: E402 +from pathlib import Path # noqa: E402 + +_COLL_HELPER = Path(__file__).parent / "fsdp_overlap_helper" / "estimator_collective_helper.py" + + +@requires_cuda +@pytest.mark.skipif(shutil.which("torchrun") is None, reason="requires torchrun") +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") +def test_collective_profile_accuracy_multi_rank(): + env = os.environ.copy() + env["MAGI_LOGGING_LEVEL"] = "warning" + p = subprocess.run( + ["torchrun", "--nproc_per_node=2", "--master_port=29681", str(_COLL_HELPER)], + env=env, + capture_output=True, + text=True, + timeout=600, + ) + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" + assert "COLL_ACCURATE ok=True" in p.stdout, out[-3000:] + assert "COLL_WARMSYNC" in p.stdout and "ok=True" in p.stdout, out[-3000:] + assert "COLL_PASS" in p.stdout, out[-3000:] diff --git a/tests/feature_tests/test_profiling_registry.py b/tests/feature_tests/test_profiling_registry.py new file mode 100644 index 0000000..42587b2 --- /dev/null +++ b/tests/feature_tests/test_profiling_registry.py @@ -0,0 +1,82 @@ +# Copyright (c) 2025 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. + +"""Unit tests for the profiling benchmark-input registry +(``magi_compiler.profiling.benchmark_inputs``). + +Pure-CPU: the registry is a plain module-level dict/set, no torch/GPU/distributed. +""" + +import pytest + +from magi_compiler.profiling import benchmark_inputs as bi +from magi_compiler.profiling import get_benchmark_inputs_hook, op_has_internal_collective, register_benchmark_inputs + + +@pytest.fixture +def clean_registry(): + """Snapshot + restore the global registry so tests don't leak into each other.""" + hooks = dict(bi._BENCHMARK_INPUT_HOOKS) + coll = set(bi._INTERNAL_COLLECTIVE_OPS) + yield + bi._BENCHMARK_INPUT_HOOKS.clear() + bi._BENCHMARK_INPUT_HOOKS.update(hooks) + bi._INTERNAL_COLLECTIVE_OPS.clear() + bi._INTERNAL_COLLECTIVE_OPS.update(coll) + + +def test_unregistered_returns_none(clean_registry): + assert get_benchmark_inputs_hook("ns::never_registered") is None + assert op_has_internal_collective("ns::never_registered") is False + + +def test_register_and_get(clean_registry): + def hook(fx_node, realize): + return None + + register_benchmark_inputs("ns::op_a", hook) + assert get_benchmark_inputs_hook("ns::op_a") is hook + # not flagged as internal-collective by default + assert op_has_internal_collective("ns::op_a") is False + + +def test_internal_collective_flag(clean_registry): + register_benchmark_inputs("ns::coll_op", lambda n, r: None, has_internal_collective=True) + assert op_has_internal_collective("ns::coll_op") is True + assert get_benchmark_inputs_hook("ns::coll_op") is not None + + # a hook registered WITHOUT the flag must not be marked + register_benchmark_inputs("ns::plain_op", lambda n, r: None) + assert op_has_internal_collective("ns::plain_op") is False + + +def test_register_overrides_previous(clean_registry): + def hook1(n, r): + return "1" + + def hook2(n, r): + return "2" + + register_benchmark_inputs("ns::op_b", hook1) + assert get_benchmark_inputs_hook("ns::op_b") is hook1 + register_benchmark_inputs("ns::op_b", hook2) + assert get_benchmark_inputs_hook("ns::op_b") is hook2 + + +def test_hook_is_callable_returning_none(clean_registry): + """A no-op hook (returns None -> fall back to generic realize) is valid.""" + register_benchmark_inputs("ns::noop", lambda fx_node, realize: None, has_internal_collective=True) + hook = get_benchmark_inputs_hook("ns::noop") + assert hook(object(), lambda x: x) is None + assert op_has_internal_collective("ns::noop") is True From 6ef1ea7d696e1e38fa321b6dd00d58480ebcdab7 Mon Sep 17 00:00:00 2001 From: wtr Date: Mon, 13 Jul 2026 15:18:13 +0800 Subject: [PATCH 3/7] [Fix] FSDP fullgraph overlap: correct cost model to eliminate exposed pre-MoE all-gathers --- magi_compiler/config.py | 30 ++++++++++++ magi_compiler/magi_backend/magi_backend.py | 33 ++++++------- magi_compiler/passes/fsdp_overlap/reorder.py | 49 ++++++++++++------- magi_compiler/profiling/runtime_estimator.py | 22 +++++++++ .../fsdp_overlap_helper/e2e_helper.py | 2 +- 5 files changed, 97 insertions(+), 39 deletions(-) diff --git a/magi_compiler/config.py b/magi_compiler/config.py index 7c12027..201d0ff 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -310,6 +310,22 @@ class CompileConfig(BaseSettings): "(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, @@ -318,6 +334,20 @@ class CompileConfig(BaseSettings): "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( diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 2ea092f..503de1e 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -552,17 +552,20 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: 3. Install the latest-safe-launch reorder pass, REPLACING PyTorch's builtin raise_comms/sink_waits, and enable reorder_for_compute_comm_overlap. """ - pass - - from magi_compiler.passes.fsdp_overlap import FsdpOverlapReorder, lower_and_bucket_full_graph - from magi_compiler.profiling import ProfilingRuntimeEstimator + 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 + # The model's true subgraph-boundary ops delimit bucketing regions even # though disable_graph_split emptied fx_split_ops. - boundary_ops = resolve_defined_ops(self.compile_config.splitting_ops or []) + 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, @@ -598,9 +601,6 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: # -> WILL deadlock multi-rank; only for world_size==1. # Default: world_size>1 -> "profile_sync" (accurate + safe); ==1 -> "profile". # Back-compat: MAGI_COMPILE_FSDP_OVERLAP_ANALYTICAL_COST=1 forces "analytical". - import os as _os - - import torch.distributed as _dist # DEFAULT: multi-rank -> "profile_sync"; single-rank -> "profile". # profile_sync gives REAL measured costs AND a rank-identical schedule, which is @@ -616,18 +616,10 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: # but less accurate. MAGI_COMPILE_FSDP_OVERLAP_COST_MODE=analytical|profile_sync| # profile overrides; back-compat MAGI_COMPILE_FSDP_OVERLAP_ANALYTICAL_COST=1 # forces analytical. - _world = _dist.get_world_size() if (_dist.is_available() and _dist.is_initialized()) else 1 - _mode = _os.environ.get("MAGI_COMPILE_FSDP_OVERLAP_COST_MODE") - if _mode is None: - if _os.environ.get("MAGI_COMPILE_FSDP_OVERLAP_ANALYTICAL_COST") == "1": - _mode = "analytical" - else: - _mode = "profile_sync" if _world > 1 else "profile" - + _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 - magi_logger.info("FSDP fullgraph overlap: ANALYTICAL cost (world_size=%d)", _world) else: # Pass the estimator DIRECTLY to the reorder pass rather than installing it # globally at inductor_config.estimate_op_runtime -- the global hook is @@ -638,7 +630,6 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: self._fsdp_overlap_estimator = ProfilingRuntimeEstimator() self._fsdp_overlap_estimator._sync_across_ranks = _mode == "profile_sync" cost_fn = self._fsdp_overlap_estimator - magi_logger.info("FSDP fullgraph overlap: %s cost (world_size=%d)", _mode.upper(), _world) # Our reorder pass is the ONLY reorder pass. It moves each weight # all-gather LAUNCH earlier -- into the upstream compute of the PREVIOUS @@ -655,7 +646,11 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: # launch's earliest-legal position, else the WeakDep would pin the launch # right after the previous gather and block the move. See # scripts/demo/research_exposed_allgather_before_attn.md. - reorder = FsdpOverlapReorder(slack_ns=self.compile_config.fsdp_overlap_slack_ns, cost_fn=cost_fn) + 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] diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index 7a7dc4e..a18b979 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -114,8 +114,14 @@ def _is_multi_output(snode: BaseSchedulerNode) -> bool: class FsdpOverlapReorder: """Callable reorder pass (see module docstring).""" - def __init__(self, slack_ns: float = _DEFAULT_SLACK_NS, cost_fn=None) -> None: + def __init__(self, slack_ns: float = _DEFAULT_SLACK_NS, cost_fn=None, comm_contention_factor: float = 1.0) -> None: self.slack_ns = slack_ns + # Scales each gather's estimated comm when sizing its compute window + # (need = comm * factor + slack). The estimator measures collectives in + # ISOLATION; in-situ they contend with the very compute that hides them + # (~1.4-1.5x slower measured on 8xH100). See CompileConfig. + # fsdp_overlap_comm_contention_factor. + self.comm_contention_factor = comm_contention_factor # cost_fn: snode -> ns. Default uses Inductor's estimate_op_runtime hook, # which MagiBackend points at the profiling estimator. if cost_fn is None: @@ -136,6 +142,7 @@ def __deepcopy__(self, memo): # by reference (it is itself deepcopy-safe: see ProfilingRuntimeEstimator). new = FsdpOverlapReorder.__new__(FsdpOverlapReorder) new.slack_ns = self.slack_ns + new.comm_contention_factor = self.comm_contention_factor new._cost_fn = self._cost_fn new._cost_cache = {} memo[id(self)] = new @@ -268,7 +275,11 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: # Start just before the launch, but no later than where the previous # (later) gather already consumed compute down to. compute_idx = min(compute_idx, cur) - need = comm_runtime + self.slack_ns + # comm is measured in isolation but runs concurrent with the compute + # that hides it -> scale by the contention factor before adding slack. + # Under-reserving exposes the collective's tail (the wait stalls); + # over-reserving just launches earlier on an otherwise-idle comm stream. + need = comm_runtime * self.comm_contention_factor + self.slack_ns acc = 0.0 t = compute_idx while t > lower: @@ -285,15 +296,28 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: target = max(lower, t) targets[launch] = (target, group) compute_idx = target # next (earlier) gather resumes from actual placement + # Per-gather placement decision, the record that answers "why didn't + # this gather move earlier": + # cur = original program index of the launch + # target = where it was placed (== cur means NOT moved) + # lower = earliest LEGAL index (real-dep floor) it could move to + # fc_idx = first real consumer (the wait's user) + # comm = the gather's runtime it needs to hide + # acc_upstream = compute actually found in [target, cur] to hide it + # verdict = hidden (acc>=need) | COMPUTE-LIMITED (ran out of upstream + # compute before covering comm -- i.e. hit `lower` or the + # previous gather's placement first) if _debug_enabled(): magi_logger.debug( - "FSDP overlap: launch cur=%d -> target=%d fc=%d lower=%d | comm=%.1fus " "acc_upstream=%.1fus %s", + "FSDP overlap placement: launch cur=%d -> target=%d fc=%d lower=%d | " + "comm=%.1fus acc_upstream=%.1fus need=%.1fus %s", cur, target, fc_idx, lower, comm_runtime / 1e3, acc / 1e3, + need / 1e3, "hidden" if acc >= need else "COMPUTE-LIMITED", ) @@ -365,22 +389,9 @@ def _key(s): measured, cache_hits, ) - # Full op->time table at DEBUG, or to a file when MAGI_COMPILE_FSDP_OVERLAP_DUMP - # is set (so the estimate table can be captured WITHOUT global DEBUG logging, - # which can trip torch's PT2_COMPILE chromium-event assertion on some builds). - if hasattr(self._cost_fn, "summary"): - import os as _os - - if _debug_enabled(): - magi_logger.debug("FSDP overlap %s", self._cost_fn.summary()) - dump = _os.environ.get("MAGI_COMPILE_FSDP_OVERLAP_DUMP") - if dump: - try: - rank = _os.environ.get("RANK", "0") - with open(f"{dump}.rank{rank}", "a") as f: - f.write(self._cost_fn.summary() + "\n") - except Exception as exc: # noqa: BLE001 - magi_logger.warning("FSDP overlap: could not write cost dump: %s", exc) + # Full op->time table at DEBUG. + if _debug_enabled() and hasattr(self._cost_fn, "summary"): + magi_logger.debug("FSDP overlap %s", self._cost_fn.summary()) return order # -- group detection -------------------------------------------------- diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index af7ad81..99418d3 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -131,6 +131,16 @@ def _snode_label(snode: BaseSchedulerNode, max_shapes: int = 3) -> str: return f"{target}[{','.join(shapes)}]" if shapes else target +def _is_multi_output_unpack(snode: BaseSchedulerNode) -> bool: + """True for a ``MultiOutput`` snode -- the zero-cost getitem that unpacks one + output of a multi-output extern (FallbackKernel). It shares its origin fx + node with the parent extern, so it must never go through the structural-key + table (the key collides with the parent's and returns the parent's runtime).""" + from torch._inductor.ir import MultiOutput + + return type(getattr(snode, "node", None)) is MultiOutput + + def _structural_key(snode: BaseSchedulerNode) -> tuple | None: """A cache key that is identical for isomorphic kernels (same op set + same input shapes/dtypes) so repeated layers share one measurement. Returns None @@ -546,6 +556,18 @@ def __call__(self, snode: BaseSchedulerNode) -> float: if contains_wait(snode) and not contains_collective(snode): return _safe_analytical(snode) + # A MultiOutput unpack (getitem off a multi-output extern, e.g. attention's + # (out, lse) FallbackKernel) launches NO kernel -- it is 0-cost. It MUST be + # short-circuited BEFORE the op->time table: it shares its origin fx node + # with the parent extern, so ``_structural_key`` COLLIDES with the parent's + # key and the table would return the parent's full runtime (measured: the + # fa3 attention MultiOutput was costed 45.7ms; the FSDP-overlap reorder then + # believed a weight all-gather placed between the attention kernel and its + # unpack node was hidden by 45.7ms of "compute" -- one slot short of the + # real attention window -> the gather ran fully exposed before the MoE). + if _is_multi_output_unpack(snode): + return 0.0 + # COLLECTIVES: never benchmark a real collective HERE. The Inductor reorder # pass runs per-rank during independent, non-co-scheduled compilation, so # issuing an NCCL op in __call__ desyncs ranks -> watchdog hang (observed on diff --git a/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py b/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py index 3d550c3..2a0b02c 100644 --- a/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py +++ b/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py @@ -102,7 +102,6 @@ def main() -> None: dev = torch.cuda.current_device() torch.manual_seed(0) - os.environ["MAGI_COMPILE_FSDP_OVERLAP_COST_MODE"] = args.cost_mode # Surface the backend's INFO logs (redistribute lowering / bucketing / reorder # "repositioned N/M") to stderr so the pytest driver can assert the chain ran. # These come back through torch's logging even from standalone_compile. @@ -136,6 +135,7 @@ def _patch(cfg): cfg.disable_graph_split = True cfg.enable_fsdp_fullgraph_overlap = True cfg.fsdp_fullgraph_bucket_mode = "coalesced" + cfg.fsdp_overlap_cost_mode = args.cost_mode return cfg # dim 0 of the input (token count) is the dynamic dim. From b5ac999b055311c49d292429ff2a3515c7bd45a8 Mon Sep 17 00:00:00 2001 From: wtr Date: Mon, 13 Jul 2026 18:07:34 +0800 Subject: [PATCH 4/7] [Chore]: Improve code comments --- magi_compiler/config.py | 8 +-- magi_compiler/magi_backend/magi_backend.py | 61 +------------------ .../fsdp_overlap/redistribute_lowering.py | 13 ---- magi_compiler/passes/fsdp_overlap/reorder.py | 25 +------- magi_compiler/profiling/benchmark_inputs.py | 44 ++++++++++--- 5 files changed, 40 insertions(+), 111 deletions(-) diff --git a/magi_compiler/config.py b/magi_compiler/config.py index 201d0ff..daeb1b0 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -267,11 +267,7 @@ class CompileConfig(BaseSettings): 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. The boundary custom ops (e.g. " - "gaga4_mh_moe / attention) stay in the graph as opaque extern calls that Inductor emits " - "verbatim -- there is exactly one compiled submod instead of one per compute region between " - "boundaries. Useful for comparing whole-graph Inductor codegen against the piecewise path. " - "Overrides splitting_ops (they are ignored). Incompatible with cudagraph_mode=PIECEWISE." + "WHOLE graph to Inductor as a single piecewise submodule." ), ) enable_fsdp_fullgraph_overlap: bool = Field( @@ -284,8 +280,6 @@ class CompileConfig(BaseSettings): "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 " - "(the analytical roofline is unusable -- 0 for custom ops, 60x low for fused pointwise, 1500x high " - "for matmul; see scripts/demo/research_estimate_op_runtime_findings.md)." ), ) fsdp_fullgraph_bucket_mode: str = Field( diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 503de1e..4ba8ee7 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -561,8 +561,6 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: from magi_compiler.passes.fsdp_overlap import FsdpOverlapReorder, lower_and_bucket_full_graph from magi_compiler.profiling import ProfilingRuntimeEstimator - # The model's true subgraph-boundary ops delimit bucketing regions even - # though disable_graph_split emptied fx_split_ops. boundary_ops = resolve_defined_ops( [] if self.compile_config.disable_graph_split else [] or self.compile_config.splitting_ops ) @@ -580,72 +578,15 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: n_buckets, ) - # Cost model for the reorder's placement sweep. - # - # MULTI-RANK CORRECTNESS (critical): the reorder decides how far to hoist each - # weight all-gather from accumulated per-node costs. If those costs differ - # across ranks, the sweep places a different number of gathers before the eager - # CP all_to_all inside the attention op -> weight-PG vs CP-PG collectives - # interleave in rank-divergent order -> NCCL deadlock (verified on 8-GPU gaga4: - # per-rank profiling -> post-reorder graph `0037` differs across ranks -> hang - # at SeqNum=12). Costs must be rank-IDENTICAL. Three modes (env - # MAGI_COMPILE_FSDP_OVERLAP_COST_MODE = analytical | profile_sync | profile): - # * "analytical" : Inductor roofline (shapes+device only -> rank-identical). - # Deterministic, zero overhead, but less accurate. - # * "profile_sync" : REAL per-op profiling, re-measured in rank-lockstep - # (barrier + fixed iters) and MAX-reduced over gloo so the - # table is identical on every rank -- accurate AND safe. - # Requires the identical-graph precondition (guaranteed by - # the unconditional-pad lowering fix). - # * "profile" : plain per-rank profiling, NO sync. Rank-nondeterministic - # -> WILL deadlock multi-rank; only for world_size==1. - # Default: world_size>1 -> "profile_sync" (accurate + safe); ==1 -> "profile". - # Back-compat: MAGI_COMPILE_FSDP_OVERLAP_ANALYTICAL_COST=1 forces "analytical". - - # DEFAULT: multi-rank -> "profile_sync"; single-rank -> "profile". - # profile_sync gives REAL measured costs AND a rank-identical schedule, which is - # now SAFE because the model builder replicates uneven-Shard(0) params - # (_replicate_uneven_shard_params in simple_fsdp.py) so the compiled graph is - # structurally identical on every rank -> same structural-key set -> the - # rank-lockstep warm_and_sync max-reduce fully reconciles costs. (An earlier - # attempt without the replicate fix hung even with profile_sync, because trailing - # ranks had extra constant_pad_nd nodes from uneven shards -> divergent placement - # that cost-sync couldn't fix; replicating those tiny params removes the source.) - # "analytical" (Inductor roofline, shapes+device only) is the deadlock-free - # FALLBACK if a model still has per-rank graph divergence -- rank-deterministic - # but less accurate. MAGI_COMPILE_FSDP_OVERLAP_COST_MODE=analytical|profile_sync| - # profile overrides; back-compat MAGI_COMPILE_FSDP_OVERLAP_ANALYTICAL_COST=1 - # forces analytical. _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: - # Pass the estimator DIRECTLY to the reorder pass rather than installing it - # globally at inductor_config.estimate_op_runtime -- the global hook is - # consulted by Inductor's own scheduling/caching in ways that specialize - # dynamic shapes. In "profile_sync" the reorder detects warm_and_sync and - # reconciles costs across ranks; "profile" leaves them per-rank (single - # rank only). We set a flag the reorder reads to decide whether to sync. self._fsdp_overlap_estimator = ProfilingRuntimeEstimator() self._fsdp_overlap_estimator._sync_across_ranks = _mode == "profile_sync" cost_fn = self._fsdp_overlap_estimator - # Our reorder pass is the ONLY reorder pass. It moves each weight - # all-gather LAUNCH earlier -- into the upstream compute of the PREVIOUS - # region (a bounded distance~=1 prefetch) -- so the compute between the new - # launch position and the gather's (unmoved) wait hides the collective. - # We deliberately do NOT append Inductor's builtin ``sink_waits``: it defers - # every wait unconditionally, which lets every launch (they depend only on - # the param shards, so all are "ready" at graph start) float to the graph - # top -> all layers' weights resident at once -> OOM on real 8-way FSDP. - # The pass instead moves launches a bounded amount (only far enough that the - # accumulated upstream compute covers comm), so ~2 layers' weights are - # resident at a time. Crucially it ignores the artificial WeakDep ordering - # between collectives (FSDP gathers are independent) when computing each - # launch's earliest-legal position, else the WeakDep would pin the launch - # right after the previous gather and block the move. See - # scripts/demo/research_exposed_allgather_before_attn.md. reorder = FsdpOverlapReorder( slack_ns=self.compile_config.fsdp_overlap_slack_ns, cost_fn=cost_fn, @@ -731,7 +672,7 @@ def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[Spli # Step 5: visualize the split graph if envs.MAGI_ENABLE_FX_GRAPH_VIZ: - # save_fx_graph_visualization(split_gm.graph, sub_dir="after_split", filename="split_gm_root") + save_fx_graph_visualization(split_gm.graph, sub_dir="after_split", filename="split_gm_root") for item in piecewise_graphs: save_fx_graph_visualization(item.graph.graph, sub_dir="after_split", filename=item.submod_name) diff --git a/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py index cf419eb..8a5e262 100644 --- a/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py +++ b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py @@ -140,19 +140,6 @@ def lower_prim_redistribute_to_collectives(graph: fx.GraphModule) -> int: # Pad the local shard up to `chunk` rows when this rank owns fewer # (uneven Shard(0): trailing ranks get remainder/empty). # - # NOTE (multi-rank): this `if L < chunk` branch makes trailing ranks' graphs - # STRUCTURALLY DIFFERENT (extra constant_pad_nd nodes) from full-chunk ranks. - # That per-rank structure does NOT by itself break the FSDP-overlap reorder, - # AS LONG AS the reorder's cost model is rank-deterministic (analytical) so it - # places gathers identically regardless of the pad-node count -- verified on - # 8-GPU gaga4: analytical cost -> gather placement identical across ranks - # (4==4 before the 1st attention) and runs clean 4/4, even though rank0 has 0 - # pads and rank4 has 1170. (Emitting the pad unconditionally does NOT help: - # Inductor DCEs the zero-width pad on full-chunk ranks, so the structure - # diverges again post-lowering anyway.) PER-RANK PROFILING is the thing that - # breaks it -- its per-rank costs make the placement diverge (7 vs 5) and - # deadlock; MagiBackend therefore uses analytical cost for world_size > 1. - # # meta example_values MUST be created via FakeTensor ops (``local.new_empty``), # NOT ``torch.empty(..., device=cuda)``: a FakeTensor's ``.device`` is a real # ``cuda:N`` so torch.empty allocates a REAL full-size buffer per weight diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index a18b979..c82074f 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -189,14 +189,13 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: # This placement sweep MUST produce an IDENTICAL schedule on every rank, else # the weight all-gathers (weight PG) interleave with the eager CP all_to_all # (CP PG, inside the attention boundary op) in a rank-divergent order -> - # cross-PG NCCL collective-order mismatch -> deadlock (verified via flight - # recorder on 8-GPU gaga4: hang at SeqNum=12, rank0 races weight-PG ahead - # while peers block on the CP all_to_all). Rank-consistency needs BOTH: + # cross-PG NCCL collective-order mismatch -> deadlock + # Rank-consistency needs BOTH: # 1. IDENTICAL INPUT GRAPH per rank -- guaranteed by the FSDP redistribute # lowering now emitting 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 (different snode-list length -> - # divergent placement). Verified: stage `0008` byte-identical across ranks. + # divergent placement). # 2. RANK-DETERMINISTIC COSTS -- the sweep accumulates `_cost` to decide how # far to hoist each launch. Profiling benchmarks PER RANK (timing noise + # replays the attention op's internal CP all_to_all per rank), so with @@ -250,24 +249,6 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: lower = self._earliest_legal_index(group, order, index_of, buf_to_snode, op_to_snode) plans.append((launch, group, fc_idx, comm_runtime, lower)) - # ---- Sweep the compute pointer backward (UPSTREAM-of-launch scan) ---- - # For each gather (LATEST first) accumulate the compute that sits IMMEDIATELY - # BEFORE its launch, walking backward, until the accumulated compute covers - # comm; move the launch to just before that compute block. Because this pass - # moves ONLY the launch (never the wait) and lowering emits ``all_gather; - # wait`` back-to-back (wait at launch+1), the compute we walk past ends up in - # the [new_launch, wait] window and therefore actually overlaps the gather. - # Scanning UPSTREAM (from the launch) -- not downstream from the consumer -- - # is the fix: the old downstream scan counted compute that runs AFTER the - # (unmoved) wait and can never overlap, so single-weight gathers were judged - # "covered" and left glued to their wait -> fully exposed. - # - # A SINGLE compute pointer walks backward CONTINUOUSLY and is NEVER reset per - # gather (each gather resumes from where the previous, later gather stopped), - # so no two gathers claim the same compute -- this serializes the single NCCL - # stream and keeps collectives in their original relative order (targets only - # decrease). A gather with no upstream compute before `lower` (e.g. the - # graph's first gather) can't move and stays where it is (structural bubble). targets: dict = {} # launch -> target index (in original order space) compute_idx = len(order) # scan compute strictly below this for launch, group, fc_idx, comm_runtime, lower in reversed(plans): diff --git a/magi_compiler/profiling/benchmark_inputs.py b/magi_compiler/profiling/benchmark_inputs.py index bdf7a56..4e4bda6 100644 --- a/magi_compiler/profiling/benchmark_inputs.py +++ b/magi_compiler/profiling/benchmark_inputs.py @@ -16,15 +16,15 @@ Some opaque boundary custom ops cannot be replayed for timing from generic size-hinted tensors alone, because they carry VALUE-DEPENDENT metadata that must -be self-consistent (not just right-shaped) or they raise -- e.g. -``gaga4_fa3_with_sink_cp`` takes a ``cp_split_sizes`` tensor whose values must sum -to the sequence length and equal the per-rank seqlen, else its internal CP -all_to_all asserts and the op falls back to a 0 cost estimate. +be self-consistent (not just right-shaped) or they raise -- e.g. a context-parallel +attention op taking a split-sizes tensor whose values must sum to the sequence +length, else its internal all_to_all asserts and the op falls back to a 0 cost +estimate. -The owning code (which knows the op's semantics -- e.g. athena's gaga4 model) -registers a hook here that builds a VALID, RANK-DETERMINISTIC set of replay inputs. -MagiCompiler stays free of model-specific knowledge: ``_measure_extern`` just looks -up the op by name and, if a hook exists, uses it instead of the generic realizer. +The owning code (the model package that defines the custom op) registers a hook +here that builds a VALID, RANK-DETERMINISTIC set of replay inputs. MagiCompiler +stays free of model-specific knowledge: ``_measure_extern`` just looks up the op +by name and, if a hook exists, uses it instead of the generic realizer. A hook is ``fn(fx_node, realize) -> (args, kwargs) | None``: * ``fx_node`` -- the ``torch.fx.Node`` for the op (read its ``args``/``kwargs`` @@ -38,13 +38,39 @@ Hooks MUST produce rank-identical inputs (derive everything from shapes / static sizes, no per-rank timing or state) so the rank-lockstep ``warm_and_sync`` measure issues any internal collective in lockstep across ranks. + +Example -- a custom attention op ``mylib::attn_cp(q, k, v, split_sizes, scale)`` +whose internal all_to_all requires ``split_sizes`` values to sum to the sequence +length (a zero-filled tensor from the generic realizer would make it raise):: + + from magi_compiler.profiling import register_benchmark_inputs + + def _attn_cp_benchmark_inputs(fx_node, realize): + q_node, k_node, v_node, _split_node, scale = fx_node.args + q = realize(q_node) # plain tensor args: reuse the generic realizer + k = realize(k_node) + v = realize(v_node) + # Rebuild the VALUE-dependent arg from SHAPE hints only (rank-identical): + # a uniform split summing to the per-rank seqlen. + seq, group_size = q.shape[0], 4 + split_sizes = torch.full((group_size,), seq // group_size, dtype=torch.int32, device=q.device) + split_sizes[-1] += seq - int(split_sizes.sum()) + return (q, k, v, split_sizes, float(scale)), {} + + register_benchmark_inputs( + "mylib::attn_cp", _attn_cp_benchmark_inputs, has_internal_collective=True + ) + +Register at import time of the module that defines the op (so the hook exists +before any compile). Returning ``None`` from the hook falls back to the generic +realizer for that call. """ from __future__ import annotations from typing import Callable -# op name (``torch.ops.ns.op`` overload string, e.g. "athena::gaga4_fa3_with_sink_cp") +# op name (``torch.ops.ns.op`` overload string, e.g. "mylib::attn_cp") # -> hook. Also records which ops issue an internal collective (need fixed-iter # lockstep replay) -- so MagiCompiler no longer hardcodes model-specific op names. _BENCHMARK_INPUT_HOOKS: dict[str, Callable] = {} From 341a97a413b4e0be400618e48ef14b711da4c5de Mon Sep 17 00:00:00 2001 From: wtr Date: Tue, 14 Jul 2026 15:06:17 +0800 Subject: [PATCH 5/7] Add torchtitan in requirements.txt --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 4de0443..97e890c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,4 +5,5 @@ depyf graphviz pydantic-settings seaborn +torchtitan==0.2.0 triton==3.5.0 From f5ddfd2fb05ca785fc9428d6dee0e528892268a9 Mon Sep 17 00:00:00 2001 From: wtr Date: Thu, 16 Jul 2026 11:43:47 +0800 Subject: [PATCH 6/7] [Refactor] FSDP fullgraph overlap: whole-graph bucketing, nested FSDPConfig, cross-rank fail-fast --- magi_compiler/config.py | 132 ++++----- magi_compiler/magi_backend/magi_backend.py | 53 +--- magi_compiler/passes/fsdp_overlap/__init__.py | 22 +- .../passes/fsdp_overlap/bucket_all_gather.py | 101 ++----- .../passes/fsdp_overlap/lower_and_bucket.py | 80 +----- magi_compiler/passes/fsdp_overlap/reorder.py | 272 +++++++----------- magi_compiler/profiling/__init__.py | 12 - magi_compiler/profiling/runtime_estimator.py | 82 ++++-- .../fsdp_overlap_helper/e2e_helper.py | 33 +-- .../estimator_collective_helper.py | 28 +- .../fsdp_overlap_helper/reorder_helper.py | 52 +++- .../feature_tests/test_fsdp_overlap_bucket.py | 96 +++---- tests/feature_tests/test_fsdp_overlap_e2e.py | 4 +- .../test_fsdp_overlap_reorder.py | 19 +- .../feature_tests/test_profiling_estimator.py | 2 + 15 files changed, 438 insertions(+), 550 deletions(-) diff --git a/magi_compiler/config.py b/magi_compiler/config.py index daeb1b0..0af9176 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -182,6 +182,64 @@ class OffloadConfig(BaseModel): bandwidth_safety_factor: float = Field(0.9, description="The safety factor for the H2D bandwidth.") +class FSDPConfig(BaseModel): + """Whole-graph FSDP weight all-gather / compute overlap (SimpleFSDP models + compiled with ``disable_graph_split=True``). + """ + + enable_fullgraph_overlap: bool = Field( + False, + description=( + "Lower SimpleFSDP weight prim_redistribute to explicit collectives, bucket them (bucket_mode), " + "and install the latest-safe-launch reorder pass that hoists each all-gather launch just far " + "enough upstream for compute to hide it. Requires disable_graph_split=True and cudagraph_mode=NONE." + ), + ) + bucket_mode: str = Field( + "none", + description=( + "'none' = one all_gather + wait per weight; 'coalesced' = one all_gather_into_tensor_coalesced " + "per bucket (single launch, N getitems/waits). Buckets are whole-graph, broken only by " + "program-order dtype changes and the bucket_size_mib cap." + ), + ) + bucket_size_mib: int = Field( + 0, + ge=0, + description=( + "Per-bucket cap on accumulated local-shard MiB for coalesced bucketing. 0 = no cap " + "(one bucket per (group, dtype) run)." + ), + ) + cost_mode: Literal["profile_sync", "analytical"] = Field( + "profile_sync", + description=( + "Cost model for the reorder placement; must be rank-identical multi-rank (else NCCL deadlock). " + "'profile_sync' (default): real per-op profiling, re-measured in rank-lockstep and max-reduced; " + "requires structurally identical per-rank graphs (verified at runtime, degrades safely on " + "mismatch). 'analytical': Inductor roofline -- rank-deterministic, less accurate, deadlock-free " + "fallback." + ), + ) + slack_ns: float = Field( + 5000.0, + ge=0.0, + description=( + "Extra headroom (ns) added to each collective's runtime when sizing its compute window, " + "absorbing estimator error + launch latency." + ), + ) + comm_contention_factor: float = Field( + 1.5, + ge=1.0, + description=( + "Multiplier on each collective's estimated runtime when sizing its compute window " + "(need = comm * factor + slack): collectives are measured in isolation but run concurrent " + "with the compute that hides them (~1.4-1.5x slower in-situ on 8xH100)." + ), + ) + + def _find_cutlass_root() -> str: """Return the CUTLASS source root, or empty string if not found.""" path = os.environ.get("MAGI_CUTLASS_ROOT", "/usr/local/cutlass") @@ -270,76 +328,12 @@ class CompileConfig(BaseSettings): "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, + # ---- Whole-graph FSDP overlap ---- + fsdp_config: FSDPConfig = Field( + FSDPConfig(), 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." + "Whole-graph FSDP weight all-gather / compute overlap (lowering + bucketing + reorder + cost " + "model). Fields reachable via MAGI_COMPILE_FSDP_CONFIG__ env vars." ), ) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 4ba8ee7..ff4b44d 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -546,51 +546,42 @@ 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. + optionally bucket them over the whole graph (no region partitioning; + buckets break only at dtype changes and the size cap). 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. """ + fsdp_cfg = self.compile_config.fsdp_config if not self.compile_config.disable_graph_split: - raise ValueError("enable_fsdp_fullgraph_overlap requires disable_graph_split=True") + raise ValueError("fsdp_config.enable_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") + raise ValueError("fsdp_config.enable_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, - ) + bucket_size_bytes = int(fsdp_cfg.bucket_size_mib) * 1024 * 1024 + n_buckets = lower_and_bucket_full_graph(graph, fsdp_cfg.bucket_mode, 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, + fsdp_cfg.bucket_mode, + fsdp_cfg.bucket_size_mib, n_buckets, ) - _mode = self.compile_config.fsdp_overlap_cost_mode - if _mode == "analytical": + if fsdp_cfg.cost_mode == "analytical": self._fsdp_overlap_estimator = None cost_fn = None # reorder pass defaults to Inductor's analytical estimate - else: + else: # "profile_sync" -- the default for BOTH single- and multi-rank self._fsdp_overlap_estimator = ProfilingRuntimeEstimator() - self._fsdp_overlap_estimator._sync_across_ranks = _mode == "profile_sync" + self._fsdp_overlap_estimator._sync_across_ranks = True 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, + slack_ns=fsdp_cfg.slack_ns, cost_fn=cost_fn, comm_contention_factor=fsdp_cfg.comm_contention_factor ) self.inductor_compile_config["reorder_for_compute_comm_overlap"] = True self.inductor_compile_config["reorder_for_compute_comm_overlap_passes"] = [reorder] @@ -598,11 +589,6 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: @observe_lifecycle("graph_split") def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[SplitItem]]: # 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") @@ -616,17 +602,8 @@ def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[Spli 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: + # Step 1.4: whole-graph FSDP overlap. + if self.compile_config.fsdp_config.enable_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. diff --git a/magi_compiler/passes/fsdp_overlap/__init__.py b/magi_compiler/passes/fsdp_overlap/__init__.py index c85747e..d507a8a 100644 --- a/magi_compiler/passes/fsdp_overlap/__init__.py +++ b/magi_compiler/passes/fsdp_overlap/__init__.py @@ -12,31 +12,13 @@ # 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 .bucket_all_gather import bucket_weight_all_gather_coalesced 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", + "bucket_weight_all_gather_coalesced", "lower_prim_redistribute_to_collectives", "lower_and_bucket_full_graph", "FsdpOverlapReorder", diff --git a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py index 0c8c962..c40b030 100644 --- a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py +++ b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py @@ -61,12 +61,7 @@ def _gathers_a_weight(node: fx.Node) -> bool: def _is_weight_all_gather(node: fx.Node) -> bool: - """A SimpleFSDP weight ``all_gather_into_tensor`` launch. - - Recognized either by the ``magi_fsdp_weight_ag`` tag (set by the redistribute - lowering pass on the real gaga4 graph) OR structurally, when the gathered - source traces back to a weight/param placeholder (covers untagged graphs such - as the demo).""" + """A SimpleFSDP weight ``all_gather_into_tensor`` launch.""" if node.op != "call_function" or node.target is not _ALL_GATHER: return False if node.meta.get("magi_fsdp_weight_ag"): @@ -84,10 +79,10 @@ def _local_shard_bytes(ag: fx.Node) -> int: return int(m.numel()) * int(m.element_size()) -def _split_region_by_dtype_and_size( +def _split_by_dtype_and_size( ag_nodes: list[fx.Node], node_index: dict[fx.Node, int], bucket_size_bytes: int ) -> list[list[fx.Node]]: - """Split a region's weight all-gathers into buckets by PROGRAM-ADJACENCY (plan B). + """Split a run of weight all-gathers into buckets by PROGRAM-ADJACENCY (plan B). Walk ``ag_nodes`` in program order and start a NEW bucket whenever: * the dtype changes (a different-dtype weight physically breaks the run -- so @@ -139,13 +134,7 @@ def _producer_chain(node: fx.Node) -> list[fx.Node]: return chain -def _coalesce_one_bucket( - graph: fx.GraphModule, - node_to_subgraph_id: dict[fx.Node, int], - node_index: dict[fx.Node, int], - sid: int, - ag_nodes: list[fx.Node], -) -> None: +def _coalesce_one_bucket(graph: fx.GraphModule, node_index: dict[fx.Node, int], ag_nodes: list[fx.Node]) -> None: """Merge ONE bucket of same-dtype weight all_gathers into a single ``all_gather_into_tensor_coalesced`` (launch + per-member getitem + per-member wait). ``ag_nodes`` must be same (group, dtype); caller guarantees len >= 2.""" @@ -164,105 +153,73 @@ def _coalesce_one_bucket( chain = [loc, *_producer_chain(loc)] if loc.op in ("call_function", "call_method") else _producer_chain(loc) for prod in sorted(chain, key=lambda n: node_index[n]): first_ag.prepend(prod) - node_to_subgraph_id[prod] = sid with graph.graph.inserting_before(first_ag): coalesced = graph.graph.call_function(_ALL_GATHER_COALESCED, (list(locals_), world, group_name)) coalesced.meta["example_value"] = list(ag_metas) coalesced.meta["magi_fsdp_weight_ag"] = True coalesced.meta["magi_fsdp_weight_ag_coalesced"] = True - node_to_subgraph_id[coalesced] = sid outs = [] for i, am in enumerate(ag_metas): gi = graph.graph.call_function(operator.getitem, (coalesced, i)) gi.meta["example_value"] = am - node_to_subgraph_id[gi] = sid outs.append(gi) for out_i, old_wait, am in zip(outs, waits, ag_metas): with graph.graph.inserting_before(old_wait): wait_i = graph.graph.call_function(_WAIT, (out_i,)) wait_i.meta["example_value"] = am - node_to_subgraph_id[wait_i] = node_to_subgraph_id.get(old_wait, sid) old_wait.replace_all_uses_with(wait_i) for ag_old, wait_old in zip(ag_nodes, waits): - node_to_subgraph_id.pop(wait_old, None) - node_to_subgraph_id.pop(ag_old, None) graph.graph.erase_node(wait_old) graph.graph.erase_node(ag_old) -def bucket_weight_all_gather_coalesced_per_region( - graph: fx.GraphModule, node_to_subgraph_id: dict[fx.Node, int], bucket_size_bytes: int = 0 -) -> int: - """Coalesce, per compute region, the SimpleFSDP weight all-gathers into ONE - ``all_gather_into_tensor_coalesced`` per ``(region_id, group_name, dtype)`` - (or into several buckets when ``bucket_size_bytes > 0``). - - When ``bucket_size_bytes > 0``, a region is further split into multiple buckets: - the weight all_gathers are walked in PROGRAM ORDER and a new bucket starts - whenever the dtype changes or the accumulated local-shard bytes would exceed the - cap (see :func:`_split_region_by_dtype_and_size`). So only same-dtype, - program-adjacent weights within the byte budget coalesce; each bucket is one - ``all_gather_into_tensor_coalesced``. ``bucket_size_bytes == 0`` (default) keeps - the original behavior (one bucket per (region, group, dtype), no size cap). - - This does NOT cat the shards into one buffer. ``all_gather_into_tensor_coalesced`` - fuses the N launches into a single NCCL group while returning one - *weight-major* output buffer per input, so each member is recovered with a - zero-copy ``operator.getitem`` -- no ``cat`` on the compute stream, no - ``split_with_sizes`` clone, and no transient ~2x memory spike. - - For a group of N weight gathers (each a single ``all_gather_into_tensor`` whose - input is the padded local shard ``(chunk_i, *rest_i)``) this builds:: - - # launch side (the coalesced launch + its getitems stay together): - coalesced = all_gather_into_tensor_coalesced([local_0, ..., local_{N-1}], W, group) - out_i = getitem(coalesced, i) # (W*chunk_i, *rest_i), weight-major - # use side (one wait per member, left at its consumer): - wait_i = wait_tensor(out_i) - - ``out_i`` has exactly the shape of the member's old ``all_gather`` output, so - each member's existing downstream (slice / real use) is simply re-pointed from - its old ``wait_tensor`` to ``wait_i``. The ``coalesced`` launch and its - ``getitem`` nodes stay together; the scheduler-level ``FsdpOverlapReorder`` pass - later moves the launch as one unit for compute/comm overlap. - - Runs AFTER redistribute lowering (called from ``lower_and_bucket_full_graph``). +def bucket_weight_all_gather_coalesced(graph: fx.GraphModule, bucket_size_bytes: int = 0) -> int: + """Coalesce the SimpleFSDP weight all-gathers over the WHOLE graph: per process + group, walk them in program order and cut a new bucket at every dtype change or + when the accumulated local-shard bytes would exceed ``bucket_size_bytes`` + (0 = no cap). Each bucket of >= 2 gathers becomes:: + + coalesced = all_gather_into_tensor_coalesced([local_0..local_{N-1}], W, group) + out_i = getitem(coalesced, i) # same shape as the member's old output + wait_i = wait_tensor(out_i) # one wait per member, left at its consumer + + ``all_gather_into_tensor_coalesced`` fuses N launches into one NCCL group but + returns one buffer per input, so members are recovered by zero-copy getitem -- + no cat/split on the compute stream, no transient memory spike. Downstream users + are re-pointed from each old wait to ``wait_i``; the launch + getitems stay + together so ``FsdpOverlapReorder`` later moves them as one unit. + + Runs after redistribute lowering (via ``lower_and_bucket_full_graph``). Returns the number of coalesced buckets created. """ node_index = {n: i for i, n in enumerate(graph.graph.nodes)} - # Group by (region, group_name) ONLY -- dtype is NOT part of the key. Program - # order + the dtype-change rule in _split_region_by_dtype_and_size decide the - # actual buckets, so a different-dtype weight physically breaks a same-dtype run - # (plan B / strict program-adjacency). With bucket_size_bytes==0 this reduces to - # one bucket per (region, group, dtype) run. - groups: dict[tuple[int, str], list[fx.Node]] = defaultdict(list) + # Key by group_name only; dtype breaks buckets positionally inside + # _split_by_dtype_and_size (strict program-adjacency). + groups: dict[str, list[fx.Node]] = defaultdict(list) for node in graph.graph.nodes: if not _is_weight_all_gather(node): continue - sid = node_to_subgraph_id.get(node) - if sid is None: - continue _, _world, group_name = node.args - groups[(sid, group_name)].append(node) + groups[group_name].append(node) buckets = 0 - for (sid, group_name), ag_nodes in groups.items(): - for sub in _split_region_by_dtype_and_size(ag_nodes, node_index, bucket_size_bytes): + for group_name, ag_nodes in groups.items(): + for sub in _split_by_dtype_and_size(ag_nodes, node_index, bucket_size_bytes): if len(sub) < 2: continue # single weight -> keep its own all_gather (nothing to coalesce) - _coalesce_one_bucket(graph, node_to_subgraph_id, node_index, sid, sub) + _coalesce_one_bucket(graph, node_index, sub) buckets += 1 if buckets: graph.graph.lint() graph.recompile() magi_logger.info( - "FSDP weight all-gather bucketing (coalesced): created %d coalesced buckets across submods " "(bucket_size_bytes=%d)", + "FSDP weight all-gather bucketing (coalesced): created %d coalesced buckets (bucket_size_bytes=%d)", buckets, bucket_size_bytes, ) diff --git a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py index b13f9ac..34123dc 100644 --- a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py +++ b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py @@ -14,88 +14,27 @@ from __future__ import annotations -"""Whole-graph FSDP weight all-gather lowering + bucketing. - -The piecewise pipeline (see ``magi_backend._split_graph``) lowers SimpleFSDP -weight ``prim_redistribute`` into explicit collectives and then buckets them -*per submod*, keyed by ``node_to_subgraph_id``. Under whole-graph compilation -(``disable_graph_split=True``) there is no split and therefore no -``node_to_subgraph_id`` -- the entire model is a single Inductor graph. - -This module reuses the exact same lowering/bucketing implementations, but drives -them with a **region-numbered** sid map computed the same way ``_split_graph`` -numbers subgraphs (increment at each boundary op) -- WITHOUT actually splitting. -That keeps bucketing granularity per compute-region: the weight gathers used by -one layer's compute (between two boundary ops, e.g. attention / MoE) coalesce -into one collective, but different regions do NOT collapse into a single -model-wide gather (which would force the whole unsharded model resident and OOM, -defeating FSDP). - -The scheduler-level :mod:`magi_compiler.passes.fsdp_overlap.reorder` pass then -places each (possibly coalesced) launch at its latest-safe position. Bucketing -here is intentionally graph-level (not scheduler-level) so the collective node -structure is fixed before Inductor lowering. -""" - -import torch import torch.fx as fx from magi_compiler.utils import magi_logger -from .bucket_all_gather import bucket_weight_all_gather_coalesced_per_region +from .bucket_all_gather import bucket_weight_all_gather_coalesced from .redistribute_lowering import lower_prim_redistribute_to_collectives -def _build_region_sid_map(graph: fx.GraphModule, boundary_ops: list["torch._ops.OpOverload"]) -> dict[fx.Node, int]: - """Number graph nodes into regions delimited by ``boundary_ops``. - - Mirrors ``_split_graph`` Step 2's ``node_to_subgraph_id`` numbering (compute - regions get even ids, boundary ops get their own odd id), but is used only to - group weight all-gathers for bucketing -- the graph is never split. Weights - within the same region share a sid, so bucketing coalesces them; weights in - different regions stay in separate buckets. - """ - resolved = set(boundary_ops) - sid = 0 - mapping: dict[fx.Node, int] = {} - for node in graph.graph.nodes: - if node.op in ("output", "placeholder"): - continue - is_boundary = node.op == "call_function" and ( - node.target in resolved or (hasattr(node.target, "default") and node.target.default in resolved) - ) - if is_boundary: - sid += 1 - mapping[node] = sid - sid += 1 - else: - mapping[node] = sid - return mapping - - -def lower_and_bucket_full_graph( - graph: fx.GraphModule, - bucket_mode: str, - boundary_ops: list["torch._ops.OpOverload"] | None = None, - bucket_size_bytes: int = 0, -) -> int: +def lower_and_bucket_full_graph(graph: fx.GraphModule, bucket_mode: str, bucket_size_bytes: int = 0) -> int: """Lower SimpleFSDP weight redistribute -> explicit collectives, then - optionally bucket them per compute-region across the WHOLE graph. + optionally bucket them across the WHOLE graph (no subgraph partitioning). ``bucket_mode``: * ``"none"`` -- lowering only (N individual all_gather + N waits). - * ``"coalesced"`` -- per region: one all_gather_into_tensor_coalesced + * ``"coalesced"`` -- one all_gather_into_tensor_coalesced per bucket (ONE launch, N getitems, N waits). - ``boundary_ops`` are the subgraph-boundary op overloads (the model's - ``splitting_ops``); they delimit bucketing regions. When None/empty, all - same-(group, dtype) gathers fall in one region -- only safe for tiny graphs. - - ``bucket_size_bytes`` (coalesced mode only): when > 0, further split each region - into buckets of at most this many local-shard bytes, breaking at dtype changes - and the byte cap in program order (see - ``bucket_weight_all_gather_coalesced_per_region``). 0 = no cap (one bucket per - region/group/dtype). + ``bucket_size_bytes`` (coalesced mode only): when > 0, split the gathers into + buckets of at most this many local-shard bytes, breaking at dtype changes and + the byte cap in program order (see ``bucket_weight_all_gather_coalesced``). + 0 = no cap (one bucket per (group, dtype) run). Returns the number of buckets created. """ @@ -106,9 +45,8 @@ def lower_and_bucket_full_graph( if bucket_mode == "none": return 0 - sid_map = _build_region_sid_map(graph, boundary_ops or []) if bucket_mode == "coalesced": - n = bucket_weight_all_gather_coalesced_per_region(graph, sid_map, bucket_size_bytes=bucket_size_bytes) + n = bucket_weight_all_gather_coalesced(graph, bucket_size_bytes=bucket_size_bytes) else: raise ValueError(f"Unknown bucket_mode={bucket_mode!r}; expected 'none' or 'coalesced'") diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index c82074f..70b5fb2 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -16,50 +16,27 @@ """Latest-safe-launch FSDP all-gather / compute overlap reorder pass. -Installed into ``torch._inductor.config.reorder_for_compute_comm_overlap_passes`` -as a callable ``list[BaseSchedulerNode] -> list[BaseSchedulerNode]``. Runs on the -WHOLE Inductor graph (MagiCompiler ``disable_graph_split=True``), replacing -PyTorch's builtin ``raise_comms``/``sink_waits``. - -Objective (opposite of ``raise_comms``, which schedules comms as EARLY as -possible): for each FSDP weight ``all_gather`` launch, find its wait / first real -consumer and place the launch at the LATEST position whose downstream compute -still hides the collective, i.e. the latest slot where:: +Installed as the only ``reorder_for_compute_comm_overlap_passes`` entry (replaces +``raise_comms``/``sink_waits``); runs on the whole Inductor graph +(``disable_graph_split=True``). For each FSDP weight all-gather launch, place it +at the LATEST position whose downstream compute still hides the collective:: sum(compute runtime between launch and first-consumer) >= comm_runtime + slack -If there is not enough upstream compute, fall back to as-early-as-legal (so we -overlap what we can; never worse than raise_comms). - -TWO-POINTER back-to-front sweep (like the offload HeuristicScheduler's reverse -walk for "latest start of transmission", offload/scheduler.py:319): - * comm pointer -- weight all-gathers in REVERSE original order (last first); - * compute pointer -- a SINGLE index that walks backward CONTINUOUSLY over the - whole graph and is NEVER reset per gather. -Each gather consumes a contiguous run of compute nodes (accumulating their cost) -until it covers its (scaled) comm; its launch target is that stopping point. The -NEXT (earlier) gather resumes the compute pointer from where it stopped -- NOT -from just before its own consumer -- so no two gathers claim the same compute -(this serializes the single NCCL stream) and the collectives keep their original -relative order automatically (targets only decrease). All moves are then applied -in one stable-sort rebuild and validated once (``_validate_full``). - -The Inductor driver (``comms.reorder_compute_and_comm_for_overlap``) does NOT -validate or repair the returned order and does NOT re-sink collectives, so the -returned list MUST be a valid topological order. We guarantee that by only ever -inserting a launch group inside ``[earliest-legal, first-consumer)`` and -asserting the move keeps every producer before / every consumer after the group -(else we abort the move and leave the launch in place). - -Handles two graph-level forms produced by -``fsdp_overlap.lower_and_bucket.lower_and_bucket_full_graph``: -* no-bucket : 1 all_gather -> 1 wait -> 1 consumer; -* coalesced : 1 packed all_gather_into_tensor_coalesced (MultiOutputLayout) - -> N MultiOutput members -> N waits -> N consumers. The launch is - ONE snode with N waits; the packed collective + its N MultiOutput - members move together as one contiguous block, before any wait. +Not enough upstream compute -> as-early-as-legal (never worse than raise_comms). + +Algorithm: two-pointer back-to-front sweep. Gathers are visited in reverse +program order; a single compute pointer walks backward continuously and is never +reset, so each gather claims a disjoint run of compute (serializing the single +NCCL stream) and targets only decrease. All moves are applied in one stable-sort +rebuild and validated once (``_validate_full``) -- the Inductor driver does NOT +repair the returned order, so it must be a valid topological order. + +Handles both lowering forms: plain all_gather (1 launch / 1 wait) and coalesced +(1 packed launch + N MultiOutput members moved together as one block + N waits). """ +import hashlib from collections import defaultdict import torch @@ -111,35 +88,54 @@ def _is_multi_output(snode: BaseSchedulerNode) -> bool: return type(node) is MultiOutput +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.""" + h = hashlib.sha256() + for s in order: + h.update(type(s).__name__.encode()) + for sub in getattr(s, "snodes", None) or (s,): + n = getattr(sub, "node", None) + if n is None: + continue + 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()) + except Exception: # noqa: BLE001 + pass + origins = getattr(n, "origins", None) + if origins: + h.update("|".join(sorted(str(getattr(o, "target", o)) for o in origins)).encode()) + return h.hexdigest() + + class FsdpOverlapReorder: - """Callable reorder pass (see module docstring).""" + """Callable reorder pass.""" def __init__(self, slack_ns: float = _DEFAULT_SLACK_NS, cost_fn=None, comm_contention_factor: float = 1.0) -> None: self.slack_ns = slack_ns - # Scales each gather's estimated comm when sizing its compute window - # (need = comm * factor + slack). The estimator measures collectives in - # ISOLATION; in-situ they contend with the very compute that hides them - # (~1.4-1.5x slower measured on 8xH100). See CompileConfig. - # fsdp_overlap_comm_contention_factor. + # need = comm * factor + slack: collectives are measured in isolation but + # run concurrent with the compute that hides them (~1.4-1.5x slower on + # 8xH100). See CompileConfig.fsdp_config.comm_contention_factor. self.comm_contention_factor = comm_contention_factor - # cost_fn: snode -> ns. Default uses Inductor's estimate_op_runtime hook, - # which MagiBackend points at the profiling estimator. + # cost_fn: snode -> ns (default: Inductor's estimate_op_runtime hook). if cost_fn is None: from torch._inductor.comms import estimate_op_runtime cost_fn = estimate_op_runtime self._cost_fn = cost_fn - # Per-call cost cache keyed by snode. Reset at the start of every - # __call__ (snodes are unique per compile). NEVER let it survive into a - # deepcopy: Inductor serializes config.reorder_for_compute_comm_overlap_passes - # into the fx-graph cache key via deepcopy, and snode keys hold FakeTensors - # whose data_ptr access raises. __deepcopy__ below returns a clean instance. + # Per-compile cost cache. Must never survive into a deepcopy: Inductor + # deepcopies this pass into the fx-graph cache key, and snode keys hold + # FakeTensors whose data_ptr access raises. self._cost_cache: dict[BaseSchedulerNode, float] = {} def __deepcopy__(self, memo): - # Return a fresh, cache-free instance so config serialization (fx-graph - # cache key) never deepcopies snode/FakeTensor state. cost_fn is shared - # by reference (it is itself deepcopy-safe: see ProfilingRuntimeEstimator). + # Fresh, cache-free instance (see _cost_cache note); cost_fn shared by + # reference -- it is itself deepcopy-safe. new = FsdpOverlapReorder.__new__(FsdpOverlapReorder) new.slack_ns = self.slack_ns new.comm_contention_factor = self.comm_contention_factor @@ -185,40 +181,41 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: index_of = {s: i for i, s in enumerate(order)} - # ---- MULTI-RANK DETERMINISM (why this schedule is rank-identical) ---- - # This placement sweep MUST produce an IDENTICAL schedule on every rank, else - # the weight all-gathers (weight PG) interleave with the eager CP all_to_all - # (CP PG, inside the attention boundary op) in a rank-divergent order -> - # cross-PG NCCL collective-order mismatch -> deadlock - # Rank-consistency needs BOTH: - # 1. IDENTICAL INPUT GRAPH per rank -- guaranteed by the FSDP redistribute - # lowering now emitting 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 (different snode-list length -> - # divergent placement). - # 2. RANK-DETERMINISTIC COSTS -- the sweep accumulates `_cost` to decide how - # far to hoist each launch. Profiling benchmarks PER RANK (timing noise + - # replays the attention op's internal CP all_to_all per rank), so with - # identical graphs the post-reorder `0037` STILL diverged (rank0=7 vs - # rank4=5 gathers before the 1st attention) and hung. Two ways to make - # costs rank-identical, both supported (MagiBackend selects via cost_fn): - # (a) ANALYTICAL cost_fn (roofline, shapes+device only) -- deterministic - # by construction, zero overhead, but less accurate; - # (b) SYNCHRONIZED PROFILING -- the estimator exposes `warm_and_sync`, - # which re-measures every op in rank-lockstep (barrier + fixed iters) - # and MAX-reduces over gloo, giving REAL costs that are still - # identical on every rank. Safe now that the graph is identical - # (same key set on all ranks -> symmetric barriers/all_gather). - # Below: if the cost_fn supports warm_and_sync, warm the table on all - # compute nodes then sync it; on failure, bail (leave graph unchanged = - # overlap off, no hang). Analytical cost_fn has no warm_and_sync -> skip. + # Fail-fast: the index-based sweep (and the lockstep profiling below) both + # require structurally IDENTICAL per-rank graphs, else the weight gathers + # interleave with other collectives in rank-divergent order -> NCCL + # deadlock. Verify with one symmetric all_gather of a graph digest; every + # rank sees the same result, so all ranks take the same branch. + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + from magi_compiler.profiling.runtime_estimator import _get_cost_sync_group + + fp = (_graph_fingerprint(order), len(order), len(launches)) + world = dist.get_world_size() + all_fp: list = [None] * world + dist.all_gather_object(all_fp, fp, group=_get_cost_sync_group()) + if any(f != all_fp[0] for f in all_fp[1:]): + magi_logger.warning( + "FSDP overlap reorder: per-rank graphs are NOT structurally identical " + "((digest, n_snodes, n_weight_gathers) per rank: %s). Reordering would " + "produce rank-divergent collective order -> NCCL deadlock; leaving the " + "graph unchanged (overlap OFF for this graph). Likely cause: uneven " + "Shard(0) params -- replicate them or use chunk-padded uniform shards.", + [(f[0][:12], f[1], f[2]) if f else None for f in all_fp], + ) + return order + + # profile_sync: warm the estimator table on every node, then re-measure in + # rank-lockstep (warm_and_sync) so costs are rank-identical. On failure, + # leave the graph unchanged (overlap off, no hang). if hasattr(self._cost_fn, "warm_and_sync") and getattr(self._cost_fn, "_sync_across_ranks", False): try: for s in order: if self._is_compute(s) or contains_collective(s): - self._cost(s) # populate estimator._table (per structural key) + self._cost(s) n_changed = self._cost_fn.warm_and_sync() - self._cost_cache = {} # drop pass-local cache -> re-read synced costs + self._cost_cache = {} # re-read synced costs magi_logger.info( "FSDP overlap reorder: rank-synchronized profiling done (%d cost entries reconciled)", n_changed ) @@ -226,19 +223,9 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: magi_logger.warning("FSDP overlap reorder: synchronized profiling failed (%s); leaving graph unchanged", exc) return order - # ---- TWO-POINTER back-to-front sweep (like offload HeuristicScheduler) ---- - # comm pointer: weight all-gathers in REVERSE original order (last first). - # compute pointer: a SINGLE index that walks backward CONTINUOUSLY across the - # whole graph and is NEVER reset per gather -- each gather consumes a - # contiguous run of compute nodes to cover its (scaled) comm, and the next - # (earlier) gather resumes from where the pointer stopped (NOT from just - # before the consumer again). This (a) serializes the single NCCL stream -- - # no two gathers claim the same compute -- and (b) keeps collectives in their - # original relative order automatically (targets only decrease). Mirrors - # offload/scheduler.py's reverse `i` walk for "latest start of transmission". + # ---- two-pointer back-to-front sweep (see module docstring) ---- launches_in_order = sorted(launches, key=lambda s: index_of[s]) # original program order - # Per-gather static facts (on the STABLE original order). plans = [] # (launch, group, fc_idx, comm_runtime, lower) for launch in launches_in_order: group = self._launch_group(launch, order, buf_to_snode, users) @@ -256,10 +243,6 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: # Start just before the launch, but no later than where the previous # (later) gather already consumed compute down to. compute_idx = min(compute_idx, cur) - # comm is measured in isolation but runs concurrent with the compute - # that hides it -> scale by the contention factor before adding slack. - # Under-reserving exposes the collective's tail (the wait stalls); - # over-reserving just launches earlier on an otherwise-idle comm stream. need = comm_runtime * self.comm_contention_factor + self.slack_ns acc = 0.0 t = compute_idx @@ -270,10 +253,8 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: t -= 1 if acc >= need: break - # `t` is the earliest position whose downstream compute (up to the launch) - # covers comm. target < cur moves the launch earlier; target == cur means - # no upstream compute was available (graph head / previous gather took it) - # -> leave in place. target >= lower always (producers stay before it). + # target == cur means no upstream compute left (graph head or previous + # gather claimed it); target >= lower keeps real producers before it. target = max(lower, t) targets[launch] = (target, group) compute_idx = target # next (earlier) gather resumes from actual placement @@ -302,25 +283,12 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: "hidden" if acc >= need else "COMPUTE-LIMITED", ) - # ---- CROSS-RANK DETERMINISM: enforce weight-gather relative order ---- - # CRITICAL for multi-rank correctness. All FSDP weight all-gathers run on - # the SAME process group, so NCCL matches the Nth all-gather call across - # ranks positionally -- the ranks MUST issue them in an identical relative - # order or they deadlock (observed: "Watchdog caught collective operation - # timeout ... SeqNum=7 COALESCED"). The reverse sweep's ``target = - # max(lower, t)`` can INVERT two gathers' order: a later gather may move far - # up (small target, derived from PER-RANK profiled compute costs), while an - # earlier gather is pinned by its real-dep floor ``lower`` to a LARGER target - # -> the earlier gather sorts after the later one. Because the targets - # depend on per-rank kernel timings (ProfilingRuntimeEstimator benchmarks on - # each rank independently), whether the inversion happens differs per rank -> - # divergent collective order -> hang. Fix: clamp targets to be - # NON-DECREASING in original program index (a graph property, identical on - # every rank). Walk launches in original order with a running max; a gather - # can never be placed before an earlier-issued gather. This preserves the - # original weight-gather subsequence on all ranks regardless of cost jitter, - # at the cost of occasionally not moving a launch as early as its compute - # window would allow (it clusters just after the prior gather instead). + # Clamp targets NON-DECREASING in original program order. NCCL matches the + # Nth call on a PG positionally across ranks, so the gathers' relative order + # must be rank-identical; `max(lower, t)` can invert two gathers and whether + # the inversion happens depends on per-rank cost jitter -> deadlock. The + # clamp pins the original subsequence at the cost of occasionally placing a + # launch later than its compute window would allow. running = -1 for launch in launches_in_order: if launch not in targets: @@ -331,14 +299,10 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: targets[launch] = (target, group) running = target - # ---- apply all moves in ONE rebuild (targets are in the ORIGINAL index - # space; applying moves incrementally would shift those indices). Assign a - # sort key per node: non-group nodes keep their original index; each launch - # group is inserted just before the node originally at `target` (key - # target-0.5), members ordered among themselves by original index. A stable - # sort then realizes every move at once while preserving all other nodes' - # relative order and each group's internal order. Collective order is kept - # because targets are monotonic (the two-pointer only decreases). ---- + # Apply all moves in ONE stable-sort rebuild (targets live in the original + # index space; incremental moves would shift them). Each launch group sorts + # to key target-0.5 (just before the node originally at `target`), members + # keep their internal order, everything else keeps its original index. group_members: dict = {} for launch, (target, group) in targets.items(): for m in group: @@ -451,22 +415,15 @@ def _is_transparent(self, snode: BaseSchedulerNode) -> bool: # -- repositioning ---------------------------------------------------- def _earliest_legal_index(self, group, order, index_of, buf_to_snode, op_to_snode) -> int: - """1 + max index of any REAL (data-dependency) producer the group needs. - - Uses ONLY non-fake buffer dependencies, walked transitively. We must NOT - use ``snode.ancestors``: that transitive-closure set is polluted by the - artificial ``WeakDep`` edges Inductor inserts between collectives to - serialize them onto one comm stream (verified: a single-weight qkv/o gather - lists the PREVIOUS layer's coalesced all-gather as an ancestor via a - ``WeakDep`` on its output buffer -- but it does NOT read that buffer). FSDP - weight all-gathers gather INDEPENDENT param shards; there is no real - gather->gather data dependency. Counting the WeakDep pinned ``lower`` right - after the previous collective, which wrongly forbade moving the launch up - past the (real, hideable) compute between the two gathers -- the exact reason - the attention gathers stayed exposed. A gather's only real producer is its - own weight-shard placeholder (+ any to_local/pad/cast chain), so real - ``lower`` is ~0; the launch is then free to move up into earlier compute. - """ + """1 + max index of any REAL (non-fake buffer) producer the group needs. + + Deliberately NOT ``snode.ancestors``: that set is polluted by the fake + ``WeakDep`` edges Inductor inserts between collectives for comm-stream + serialization. Weight gathers read independent param shards -- there is no + real gather->gather dependency -- so counting the WeakDep would pin the + launch right after the previous collective and forbid the very hoist this + pass exists for. A gather's only real producer is its weight-shard + placeholder (+ to_local/pad/cast chain), so real ``lower`` is ~0.""" group_set = set(group) lo = 0 for s in group: @@ -480,22 +437,13 @@ def _earliest_legal_index(self, group, order, index_of, buf_to_snode, op_to_snod return lo def _validate_full(self, new_order, op_to_snode, buf_to_snode, users) -> bool: - """Check the rebuilt order is a valid topological order w.r.t. REAL data - dependencies: every node's non-fake buffer producers precede it (the - Inductor driver does NOT repair the order, so a real-dep violation would - silently miscompile). O(nodes * deps). - - We deliberately do NOT validate against ``snode.ancestors``: that set is the - transitive closure of ALL deps including the artificial ``WeakDep`` ordering - edges Inductor inserts between collectives (see ``_earliest_legal_index``). - Our pass intentionally reorders a weight all-gather across such a WeakDep - (there is no real gather->gather data dependency), so an ancestors-based - check would false-reject that legal move and no-op the whole reorder. The - per-node real-buffer-dep check below is itself a complete topological - validation over the real data-dependency DAG: if every node's direct real - producers precede it, the order is a valid real-dep topological order. - WeakDep ordering is advisory (stream serialization / memory) and is NOT a - correctness constraint, so violating it is safe.""" + """Valid topological order w.r.t. REAL data deps: every node's non-fake + buffer producers precede it (the driver does not repair the order, so a + violation would silently miscompile). Checking direct producers per node + is a complete validation of the real-dep DAG. ``snode.ancestors`` is NOT + used -- it includes the fake WeakDep edges this pass intentionally crosses + (see ``_earliest_legal_index``); an ancestors check would false-reject + every legal hoist. WeakDep is advisory, not a correctness constraint.""" pos = {s: i for i, s in enumerate(new_order)} for s in new_order: sp = pos[s] diff --git a/magi_compiler/profiling/__init__.py b/magi_compiler/profiling/__init__.py index 45c7883..467b3bf 100644 --- a/magi_compiler/profiling/__init__.py +++ b/magi_compiler/profiling/__init__.py @@ -12,18 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Per-op runtime cost model (profiling) -- general infrastructure, not a pass. - -Provides: -* :class:`.runtime_estimator.ProfilingRuntimeEstimator` -- a ``snode -> nanoseconds`` - cost model that MEASURES each scheduler node's real kernel time (Triton via - ``benchmark_fused_nodes``, extern/custom via eager replay, collectives via a - rank-lockstep replay), replacing Inductor's unreliable analytical roofline. Any - pass that needs accurate per-op timing (e.g. the FSDP-overlap reorder) uses it. -* :mod:`.benchmark_inputs` -- a model-agnostic registry so the owning code (e.g. - athena's gaga4) can supply valid replay inputs / lockstep flags for opaque custom - ops that can't be replayed from generic size-hinted tensors alone. -""" from .benchmark_inputs import get_benchmark_inputs_hook, op_has_internal_collective, register_benchmark_inputs from .runtime_estimator import ProfilingRuntimeEstimator diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index 99418d3..2dde980 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -444,8 +444,7 @@ def warm_and_sync(self) -> int: schedule on every rank -- else weight-PG gathers interleave with the eager CP all_to_all in rank-divergent order -> deadlock). - PRECONDITION (caller guarantees): the graph is structurally IDENTICAL on all - ranks (the unconditional-pad lowering fix ensures this). Therefore every rank + PRECONDITION: the graph is structurally IDENTICAL on all ranks, so every rank has exactly the same set of ``_structural_key`` keys already populated in the table by the reorder pass's warm-up loop, so: * the key iteration order (sorted) is identical on every rank; @@ -453,12 +452,23 @@ def warm_and_sync(self) -> int: the same step, so the barrier-wrapped fixed-iteration replay issues the op's internal CP all_to_all in lockstep (no NCCL count mismatch); * the max-reduce over gloo is symmetric. - - Steps: for each key in sorted order, barrier -> (re)measure locally with FIXED - iters -> barrier; then all_gather_object the {key: ns} maps and take the MAX. - Fixed-iteration timing (``_measure_one``) is mandatory -- a duration-adaptive - benchmark would run a per-rank iteration count and desync the internal - collective. Returns the number of table entries whose cost changed.""" + The precondition is VERIFIED (not assumed): the key sets are all_gather'd and + compared BEFORE the barrier loop. On any mismatch (per-rank shapes from + uneven Shard(0) padding, a ``_structural_key``/``_collective_spec`` failure on + one rank, ...) entering the loop would deadlock -- rank A's Nth barrier pairs + with rank B's all_gather (gloo count mismatch), or both ranks measure a + DIFFERENT collective-containing op at the same step (NCCL mismatch inside + ``_measure_one``). Instead we warn and DEGRADE to the analytical estimate for + every entry this compile measured per-rank: analytical is a pure shape+device + function, so the resulting cost table is rank-deterministic and the reorder + stays deadlock-free (same guarantee as fsdp_config.cost_mode=analytical). + + Steps: verify key sets match across ranks; then for each key in sorted order, + barrier -> (re)measure locally with FIXED iters -> barrier; then + all_gather_object the {key: ns} maps and take the MAX. Fixed-iteration timing + (``_measure_one``) is mandatory -- a duration-adaptive benchmark would run a + per-rank iteration count and desync the internal collective. Returns the + number of table entries whose cost changed.""" import torch.distributed as dist if not (dist.is_available() and dist.is_initialized()): @@ -473,6 +483,43 @@ def warm_and_sync(self) -> int: # Re-measure each in a rank-uniform (sorted) order under barriers so any # internal collective is issued in lockstep across ranks. keys = sorted(self._table.keys(), key=repr) + + # FAIL-FAST precondition check: the barrier loop below is only lockstep-safe + # if every rank iterates the SAME key sequence. all_gather_object is a single + # symmetric collective (safe regardless of key sets), and every rank sees the + # same gathered result, so all ranks take the same branch. + key_reprs = [repr(k) for k in keys] + all_key_reprs: list = [None] * world + dist.all_gather_object(all_key_reprs, key_reprs, group=group) + if any(kr != all_key_reprs[0] for kr in all_key_reprs[1:]): + ref = set(all_key_reprs[0] or []) + mine = set(key_reprs) + missing = sorted(ref - mine)[:3] + extra = sorted(mine - ref)[:3] + magi_logger.warning( + "warm_and_sync: cross-rank profiling key sets DIFFER (counts per rank: %s; " + "this rank vs rank0 -- missing %d e.g. %s, extra %d e.g. %s). The per-rank " + "graphs are not structurally identical, so rank-lockstep measurement would " + "deadlock. Falling back to the ANALYTICAL cost estimate for this graph " + "(rank-deterministic, less accurate). Consider fsdp_config.cost_mode=analytical.", + [len(kr or []) for kr in all_key_reprs], + len(ref - mine), + missing, + len(mine - ref), + extra, + ) + n = 0 + for k, snode in self._key_snode.items(): + e = self._table.get(k) + if e is None: + continue + ns = _safe_analytical(snode) + if ns != e.ns: + n += 1 + e.ns = ns + e.measured = False + self._key_snode.clear() + return n local_ns: dict = {} for k in keys: snode = self._key_snode.get(k) @@ -556,28 +603,9 @@ def __call__(self, snode: BaseSchedulerNode) -> float: if contains_wait(snode) and not contains_collective(snode): return _safe_analytical(snode) - # A MultiOutput unpack (getitem off a multi-output extern, e.g. attention's - # (out, lse) FallbackKernel) launches NO kernel -- it is 0-cost. It MUST be - # short-circuited BEFORE the op->time table: it shares its origin fx node - # with the parent extern, so ``_structural_key`` COLLIDES with the parent's - # key and the table would return the parent's full runtime (measured: the - # fa3 attention MultiOutput was costed 45.7ms; the FSDP-overlap reorder then - # believed a weight all-gather placed between the attention kernel and its - # unpack node was hidden by 45.7ms of "compute" -- one slot short of the - # real attention window -> the gather ran fully exposed before the MoE). if _is_multi_output_unpack(snode): return 0.0 - # COLLECTIVES: never benchmark a real collective HERE. The Inductor reorder - # pass runs per-rank during independent, non-co-scheduled compilation, so - # issuing an NCCL op in __call__ desyncs ranks -> watchdog hang (observed on - # 8-GPU gaga4: rank0 raced to FSDP seqNum~194 while peers never matched). - # SEED the cost with Inductor's static analytical NCCL estimate (a pure - # function of shapes+device -> rank-deterministic). In profile_sync mode we - # ALSO stash the snode so the rank-lockstep warm_and_sync() later OVERRIDES - # this seed with a REAL measured time (the only multi-rank-safe place to touch - # NCCL) -- which is where the accuracy comes from. Collectives enter the same - # op->time table as compute so warm_and_sync can find them. if contains_collective(snode): cnode = _leaf_collective(snode) spec = _collective_spec(cnode) if cnode is not None else None diff --git a/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py b/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py index 2a0b02c..bc82efb 100644 --- a/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py +++ b/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py @@ -17,13 +17,14 @@ driver (tests/feature_tests/test_fsdp_overlap_e2e.py). Chain under test (magi_backend._apply_fsdp_fullgraph_overlap): - redistribute lowering -> per-region bucketing (coalesced) -> FsdpOverlapReorder + redistribute lowering -> whole-graph bucketing (coalesced) -> FsdpOverlapReorder Model: two Linear layers whose params are Shard(0) DTensors (via torchtitan -``data_parallel(mode="fully_shard")`` -- the same mechanism gaga4 uses), with an +``data_parallel(mode="fully_shard")``), with an opaque ``@magi_register_custom_op(is_subgraph_boundary=True)`` op between them so the -graph has a real subgraph boundary (delimits bucketing regions) and the weight -``prim_redistribute`` nodes the lowering pass rewrites. +graph contains an opaque extern call alongside the +weight ``prim_redistribute`` nodes the lowering pass rewrites. Bucketing is +whole-graph: the boundary op does NOT split buckets. Run: torchrun --nproc_per_node=N tests/feature_tests/fsdp_overlap_helper/e2e_helper.py [--cost-mode analytical|profile_sync] @@ -49,18 +50,18 @@ from magi_compiler.config import CompileMode, CudaGraphMode -# An opaque boundary op: forces a real subgraph boundary in the graph (delimits the -# bucketing regions). Elementwise so it is trivially correct. +# An opaque boundary op: stays in the graph as an opaque extern call +# Elementwise so it is trivially correct. @magi_register_custom_op(name="fsdp_e2e::boundary_gelu", is_subgraph_boundary=True) def boundary_gelu(x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.gelu(x) class Block(nn.Module): - """One transformer-ish block: two Linears with an opaque subgraph-boundary op - between them. Multiple blocks give the graph MULTIPLE boundary-delimited regions, - each with several weight all-gathers -- so the bucketing pass coalesces per region - and the reorder pass has real cross-region compute to hide gathers behind.""" + """One transformer-ish block: two Linears with an opaque extern op between them. + Multiple blocks give the graph many Shard(0) weight all-gathers interleaved with + compute -- so the bucketing pass coalesces them whole-graph and the reorder pass + has real compute to hide gathers behind.""" def __init__(self, hidden: int): super().__init__() @@ -75,8 +76,8 @@ def forward(self, x): class TinyModel(nn.Module): - """A multi-layer stack of Blocks (default 4) -> many Shard(0) weights across - several regions, exercising the full lower->bucket->reorder chain at scale.""" + """A multi-layer stack of Blocks (default 4) -> many Shard(0) weights, + exercising the full lower->bucket->reorder chain at scale.""" def __init__(self, hidden: int, n_layers: int = 4): super().__init__() @@ -90,7 +91,7 @@ def forward(self, x): def main() -> None: ap = argparse.ArgumentParser() - ap.add_argument("--cost-mode", default="analytical", choices=["analytical", "profile_sync", "profile"]) + ap.add_argument("--cost-mode", default="analytical", choices=["analytical", "profile_sync"]) ap.add_argument("--hidden", type=int, default=256) ap.add_argument("--n-layers", type=int, default=4) args = ap.parse_args() @@ -133,9 +134,9 @@ def _patch(cfg): cfg.compile_mode = CompileMode.MAGI_COMPILE cfg.cudagraph_mode = CudaGraphMode.NONE cfg.disable_graph_split = True - cfg.enable_fsdp_fullgraph_overlap = True - cfg.fsdp_fullgraph_bucket_mode = "coalesced" - cfg.fsdp_overlap_cost_mode = args.cost_mode + cfg.fsdp_config.enable_fullgraph_overlap = True + cfg.fsdp_config.bucket_mode = "coalesced" + cfg.fsdp_config.cost_mode = args.cost_mode return cfg # dim 0 of the input (token count) is the dynamic dim. diff --git a/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py b/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py index c4b53df..45e2a85 100644 --- a/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py +++ b/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py @@ -23,7 +23,11 @@ 3. INDEPENDENTLY time the same all_gather with CUDA events (rank-lockstep); 4. assert the estimate is within a tolerance band of the independent measurement. Also exercises ``ProfilingRuntimeEstimator.warm_and_sync`` (the profile_sync entry) to -confirm it runs rank-lockstep without deadlock and reconciles a collective entry. +confirm it runs rank-lockstep without deadlock and reconciles a collective entry, and +the KEY-SET MISMATCH fail-fast: when one rank's table has an extra key (simulating a +per-rank structural divergence), warm_and_sync must NOT enter the barrier loop (which +would deadlock) -- it must detect the mismatch on every rank, warn, and degrade every +stashed entry to the analytical estimate (measured=False). Run: torchrun --nproc_per_node=2 .../estimator_collective_helper.py @@ -31,6 +35,7 @@ COLL_MEASURED est_us= real_us= ratio= COLL_ACCURATE ok= COLL_WARMSYNC reconciled= ok= + COLL_MISMATCH ok= COLL_PASS / COLL_FAIL """ @@ -44,7 +49,7 @@ from torch._inductor.utils import contains_collective from magi_compiler.profiling import ProfilingRuntimeEstimator -from magi_compiler.profiling.runtime_estimator import _measure_collective_op +from magi_compiler.profiling.runtime_estimator import ProfileEntry, _measure_collective_op _AG = torch.ops._c10d_functional.all_gather_into_tensor.default _WAIT = torch.ops._c10d_functional.wait_tensor.default @@ -142,8 +147,24 @@ def main() -> None: coll_entries = [e for e in est.table.values() if e.kind == "collective"] warmsync_ok = len(coll_entries) >= 1 and all(e.measured for e in coll_entries) + # 5. key-set MISMATCH fail-fast: rank 1 injects an extra table entry so the + # cross-rank key sets differ. warm_and_sync must detect this on EVERY rank + # (symmetric all_gather_object), skip the per-key barrier loop entirely (which + # would deadlock on the count mismatch), and degrade this compile's entries to + # the analytical estimate (measured=False). Completing at all proves no hang. + est2 = ProfilingRuntimeEstimator() + est2._sync_across_ranks = True + est2(coll_snode) # both ranks: seed the shared collective entry + if rank == 1: + fake_key = ("mismatch_only_on_rank1",) + est2._table[fake_key] = ProfileEntry(ns=1.0, kind="extern", label="fake", measured=True) + est2._key_snode[fake_key] = coll_snode + est2.warm_and_sync() + mismatch_ok = all(not e.measured for e in est2.table.values()) and not est2._key_snode + dist.barrier() + # gather agreement across ranks - ok_local = accurate and warmsync_ok + ok_local = accurate and warmsync_ok and mismatch_ok t = torch.tensor([1 if ok_local else 0], device=dev) dist.all_reduce(t) all_ok = int(t.item()) == world @@ -152,6 +173,7 @@ def main() -> None: print(f"COLL_MEASURED est_us={est_ns/1e3:.1f} real_us={real_ns/1e3:.1f} ratio={ratio:.2f}", flush=True) print(f"COLL_ACCURATE ok={accurate}", flush=True) print(f"COLL_WARMSYNC reconciled={n_reconciled} ok={warmsync_ok}", flush=True) + print(f"COLL_MISMATCH ok={mismatch_ok}", flush=True) print("COLL_PASS" if all_ok else "COLL_FAIL", flush=True) rc = 0 if all_ok else 1 else: diff --git a/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py b/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py index 34d27e7..3ade6b6 100644 --- a/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py +++ b/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py @@ -23,17 +23,26 @@ out = y + gathered_use # consumer after the compute and wrap the pass so we can assert it RAN and returned a valid schedule. +With ``--mismatch`` (needs >=2 ranks): rank 1's fn gets an EXTRA compute op so the +per-rank graphs are structurally DIFFERENT. The reorder pass's cross-rank +graph-fingerprint fail-fast must fire on EVERY rank (symmetric all_gather), warn, +and leave the schedule unchanged -- completing at all proves the check itself does +not desync. + Run: torchrun --nproc_per_node=1 tests/feature_tests/fsdp_overlap_helper/reorder_helper.py + torchrun --nproc_per_node=2 ... reorder_helper.py --mismatch Markers (rank 0): REORDER_CALLED gathers= REORDER_OK moved= (pass returned; N launches repositioned) REORDER_FINITE ok= (compiled output finite + matches eager) + REORDER_MISMATCH unchanged= (--mismatch only: schedule left untouched) REORDER_PASS / REORDER_FAIL """ from __future__ import annotations +import argparse import os import torch @@ -45,6 +54,9 @@ def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--mismatch", action="store_true", help="rank1 compiles a structurally different graph") + args = ap.parse_args() dist.init_process_group("cpu:gloo,cuda:nccl") rank = dist.get_rank() world = dist.get_world_size() @@ -60,20 +72,39 @@ def main() -> None: w0 = torch.randn(H, H, device=dev, dtype=torch.bfloat16) shard = torch.randn(H, H, device=dev, dtype=torch.bfloat16) + extra_op = args.mismatch and rank == 1 # structural per-rank divergence on demand + def fn(x, w0, shard): y = (x @ w0).relu() # upstream compute the gather can hide behind + if extra_op: + y = y.sin() # rank1-only node -> graphs differ across ranks g = _WAIT(_AG(shard, world, grp)) # weight all-gather + wait gathered = g.reshape(world * H, H)[:H] # use the gathered weight return y @ gathered - # instrument the pass: count how many times it runs and how many launches move. - calls = {"n": 0, "gathers": 0, "moved": 0} + # instrument the pass: count how many times it runs, how many launches move, + # and whether the returned schedule is identical to the input (fail-fast path). + calls = {"n": 0, "gathers": 0, "moved": 0, "unchanged": True, "warned_mismatch": False} orig_call = FsdpOverlapReorder.__call__ + # magi_logger output from inside an Inductor compile does not reliably reach the + # subprocess streams; intercept the warning call itself to detect the fail-fast. + orig_warning = _ro.magi_logger.warning + + def spy_warning(msg, *a, **kw): + if "NOT structurally identical" in str(msg): + calls["warned_mismatch"] = True + print(f"REORDER_WARNED rank={rank}", flush=True) + return orig_warning(msg, *a, **kw) + + _ro.magi_logger.warning = spy_warning + def spy(self, snodes): calls["n"] += 1 calls["gathers"] += sum(1 for s in snodes if _ro._is_weight_gather(s)) + before = list(snodes) out = orig_call(self, snodes) + calls["unchanged"] = len(out) == len(before) and all(a is b for a, b in zip(out, before)) return out FsdpOverlapReorder.__call__ = spy @@ -97,18 +128,29 @@ def spy(self, snodes): inductor_config.reorder_for_compute_comm_overlap_passes = prev_passes inductor_config.force_disable_caches = prev_cache FsdpOverlapReorder.__call__ = orig_call + _ro.magi_logger.warning = orig_warning finite = bool(torch.isfinite(out).all().item()) rel = ((out.float() - eager.float()).norm() / (eager.float().norm() + 1e-6)).item() numeric_ok = finite and rel < 5e-2 + # In --mismatch mode the fail-fast must leave the schedule untouched on EVERY + # rank; agree across ranks before printing. + ok_local = calls["n"] > 0 and calls["gathers"] >= 1 and numeric_ok + if args.mismatch: + ok_local = ok_local and calls["unchanged"] and calls["warned_mismatch"] + t = torch.tensor([1 if ok_local else 0], device=dev) + dist.all_reduce(t) + all_ok = int(t.item()) == world + if rank == 0: print(f"REORDER_CALLED gathers={calls['gathers']}", flush=True) print(f"REORDER_OK ran={calls['n'] > 0}", flush=True) print(f"REORDER_FINITE ok={numeric_ok} rel={rel:.5f}", flush=True) - ok = calls["n"] > 0 and calls["gathers"] >= 1 and numeric_ok - print("REORDER_PASS" if ok else "REORDER_FAIL", flush=True) - rc = 0 if ok else 1 + if args.mismatch: + print(f"REORDER_MISMATCH unchanged={calls['unchanged']}", flush=True) + print("REORDER_PASS" if all_ok else "REORDER_FAIL", flush=True) + rc = 0 if all_ok else 1 else: rc = 0 diff --git a/tests/feature_tests/test_fsdp_overlap_bucket.py b/tests/feature_tests/test_fsdp_overlap_bucket.py index 19e4eb8..4d3d00a 100644 --- a/tests/feature_tests/test_fsdp_overlap_bucket.py +++ b/tests/feature_tests/test_fsdp_overlap_bucket.py @@ -19,6 +19,9 @@ Pure-CPU: the bucketing pass operates on ``all_gather_into_tensor`` fx nodes tagged ``magi_fsdp_weight_ag`` (it never runs the ops), so we drive it with SYNTHETIC fx graphs built with meta-tensor ``example_value``s -- no DTensor / distributed / GPU. + +Bucketing is WHOLE-GRAPH: there is no region / subgraph partitioning. Buckets +break only at program-order dtype changes and the optional byte cap. """ from collections import Counter @@ -27,7 +30,7 @@ import torch import torch.fx as fx -from magi_compiler.passes.fsdp_overlap import bucket_weight_all_gather_coalesced_per_region, lower_and_bucket_full_graph +from magi_compiler.passes.fsdp_overlap import bucket_weight_all_gather_coalesced, lower_and_bucket_full_graph _AG = torch.ops._c10d_functional.all_gather_into_tensor.default _AG_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default @@ -40,14 +43,13 @@ def _build_ag_graph(specs, world=2, group="grp0"): """Build an fx graph of independent weight all-gathers. - ``specs``: list of dicts, each ``{shape, dtype, boundary_after?}``. For each spec + ``specs``: list of dicts, each ``{shape, dtype, compute_before?}``. For each spec we emit weight_shard placeholder -> all_gather_into_tensor (tagged - magi_fsdp_weight_ag) -> wait_tensor. A spec with ``boundary`` inserts an opaque - boundary op (aten.relu here) between gathers so ``_build_region_sid_map`` splits - regions. ALL placeholders are declared first (as in a real traced graph) so the - hoisted coalesced launch stays topologically valid. - - Returns (gm, boundary_op_or_None). + magi_fsdp_weight_ag) -> wait_tensor. A spec with ``compute_before`` inserts an + opaque compute op (aten.relu here) between gathers, which must NOT break the + bucket (whole-graph bucketing has no region boundaries). ALL placeholders are + declared first (as in a real traced graph) so the hoisted coalesced launch stays + topologically valid. """ g = fx.Graph() locs = [] @@ -58,8 +60,8 @@ def _build_ag_graph(specs, world=2, group="grp0"): outs = [] for i, (s, loc) in enumerate(zip(specs, locs)): - if s.get("boundary_before"): - # opaque boundary op consuming the previous wait (keeps it in the graph) + if s.get("compute_before"): + # opaque compute op consuming the previous wait (keeps it in the graph) b = g.call_function(torch.ops.aten.relu.default, (outs[-1],)) if outs else None if b is not None: b.meta["example_value"] = outs[-1].meta["example_value"] @@ -86,13 +88,12 @@ def _n(gm, target) -> int: # --------------------------------------------------------------------------- -# bucket_weight_all_gather_coalesced_per_region +# bucket_weight_all_gather_coalesced # --------------------------------------------------------------------------- -def test_coalesce_same_region_merges_to_one(): - """N same-(region,group,dtype) gathers -> ONE coalesced + N getitems + N waits.""" +def test_coalesce_merges_to_one(): + """N same-(group,dtype) gathers -> ONE coalesced + N getitems + N waits.""" gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 3) - sid_map = {n: 0 for n in gm.graph.nodes} - n_buckets = bucket_weight_all_gather_coalesced_per_region(gm, sid_map) + n_buckets = bucket_weight_all_gather_coalesced(gm) assert n_buckets == 1 assert _n(gm, _AG_COALESCED) == 1 @@ -105,8 +106,7 @@ def test_coalesce_same_region_merges_to_one(): def test_single_gather_not_coalesced(): """A lone weight gather has nothing to coalesce -> untouched, 0 buckets.""" gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}]) - sid_map = {n: 0 for n in gm.graph.nodes} - n_buckets = bucket_weight_all_gather_coalesced_per_region(gm, sid_map) + n_buckets = bucket_weight_all_gather_coalesced(gm) assert n_buckets == 0 assert _n(gm, _AG) == 1 assert _n(gm, _AG_COALESCED) == 0 @@ -123,44 +123,44 @@ def test_dtype_change_breaks_bucket(): {"shape": (4, 8), "dtype": torch.bfloat16}, ] ) - sid_map = {n: 0 for n in gm.graph.nodes} - n_buckets = bucket_weight_all_gather_coalesced_per_region(gm, sid_map) + n_buckets = bucket_weight_all_gather_coalesced(gm) assert n_buckets == 1 # only the {bf16,bf16} run assert _n(gm, _AG_COALESCED) == 1 assert _n(gm, _AG) == 2 # the fp32 and the trailing bf16 stay individual def test_bucket_size_bytes_caps_run(): - """Same dtype, but bucket_size_bytes forces a split. 4x8 bf16 shard = 512 B each; - cap at 512 B means each gather is its own bucket -> no coalescing (0 buckets>=2).""" + """Same dtype, but bucket_size_bytes forces a split. 4x8 bf16 shard = 64 B each; + cap at 64 B means each gather is its own bucket -> no coalescing (0 buckets>=2).""" gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 4) - sid_map = {n: 0 for n in gm.graph.nodes} # 4*8*2 = 64 B local shard each; cap at 64 B => each starts a new bucket -> singletons - n_buckets = bucket_weight_all_gather_coalesced_per_region(gm, sid_map, bucket_size_bytes=64) + n_buckets = bucket_weight_all_gather_coalesced(gm, bucket_size_bytes=64) assert n_buckets == 0 assert _n(gm, _AG) == 4 # cap at 128 B => 2 shards per bucket => 2 buckets of 2 gm2 = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 4) - sid2 = {n: 0 for n in gm2.graph.nodes} - n2 = bucket_weight_all_gather_coalesced_per_region(gm2, sid2, bucket_size_bytes=128) + n2 = bucket_weight_all_gather_coalesced(gm2, bucket_size_bytes=128) assert n2 == 2 assert _n(gm2, _AG_COALESCED) == 2 -def test_different_regions_not_merged(): - """Gathers in different sids (regions) never share a bucket.""" - gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 4) - ags = [n for n in gm.graph.nodes if n.op == "call_function" and n.target is _AG] - # put first 2 gathers (+their locals/waits) in region 0, last 2 in region 2 - sid_map = {} - for n in gm.graph.nodes: - sid_map[n] = 0 - for ag in ags[2:]: - sid_map[ag] = 2 - n_buckets = bucket_weight_all_gather_coalesced_per_region(gm, sid_map) - assert n_buckets == 2 # one per region - assert _n(gm, _AG_COALESCED) == 2 +def test_compute_between_gathers_does_not_break_bucket(): + """Whole-graph bucketing: an interleaved compute op between gathers does NOT + split them into separate buckets (no region boundaries).""" + gm = _build_ag_graph( + [ + {"shape": (4, 8), "dtype": torch.bfloat16}, + {"shape": (4, 8), "dtype": torch.bfloat16}, + {"shape": (4, 8), "dtype": torch.bfloat16, "compute_before": True}, + {"shape": (4, 8), "dtype": torch.bfloat16}, + ] + ) + n_buckets = bucket_weight_all_gather_coalesced(gm) + assert n_buckets == 1 # all 4 gathers in ONE bucket despite the relu in between + assert _n(gm, _AG_COALESCED) == 1 + assert _n(gm, _AG) == 0 + gm.graph.lint() # --------------------------------------------------------------------------- @@ -176,25 +176,19 @@ def test_lower_and_bucket_mode_none_returns_zero(): def test_lower_and_bucket_mode_coalesced(): - """mode 'coalesced' with no boundary_ops -> all gathers in one region -> 1 bucket.""" + """mode 'coalesced' -> all same-(group,dtype) gathers in one whole-graph bucket.""" gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 3) n = lower_and_bucket_full_graph(gm, "coalesced") assert n == 1 assert _n(gm, _AG_COALESCED) == 1 -def test_lower_and_bucket_boundary_ops_split_regions(): - """A boundary op between gathers splits them into separate buckets.""" - gm = _build_ag_graph( - [ - {"shape": (4, 8), "dtype": torch.bfloat16}, - {"shape": (4, 8), "dtype": torch.bfloat16}, - {"shape": (4, 8), "dtype": torch.bfloat16, "boundary_before": True}, - {"shape": (4, 8), "dtype": torch.bfloat16}, - ] - ) - n = lower_and_bucket_full_graph(gm, "coalesced", boundary_ops=[torch.ops.aten.relu.default]) - assert n == 2 # {g0,g1} before boundary, {g2,g3} after +def test_lower_and_bucket_size_cap(): + """bucket_size_bytes flows through the entry point: cap 128 B on 4x64 B shards + -> 2 buckets of 2.""" + gm = _build_ag_graph([{"shape": (4, 8), "dtype": torch.bfloat16}] * 4) + n = lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=128) + assert n == 2 assert _n(gm, _AG_COALESCED) == 2 diff --git a/tests/feature_tests/test_fsdp_overlap_e2e.py b/tests/feature_tests/test_fsdp_overlap_e2e.py index fd6ac0f..954c47c 100644 --- a/tests/feature_tests/test_fsdp_overlap_e2e.py +++ b/tests/feature_tests/test_fsdp_overlap_e2e.py @@ -13,7 +13,7 @@ # limitations under the License. """End-to-end FSDP-overlap chain test: a small SimpleFSDP-sharded model through -``magi_compile`` with ``enable_fsdp_fullgraph_overlap`` -- exercising the whole +``magi_compile`` with ``fsdp_config.enable_fullgraph_overlap`` -- exercising the whole chain (redistribute lowering -> bucketing -> FsdpOverlapReorder) in one real compile. Driven via a ``torchrun`` subprocess helper (fsdp_overlap_helper/e2e_helper.py); we @@ -88,7 +88,7 @@ def test_e2e_multi_rank_full_chain(): @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") def test_e2e_multi_rank_profile_sync(): """world=2 with profile_sync cost mode: exercises warm_and_sync's rank-lockstep - real measurement path end to end (the gaga4 default).""" + real measurement path end to end.""" p = _run(2, "profile_sync", "29643") out = p.stdout + p.stderr assert p.returncode == 0, f"helper failed:\n{out[-4000:]}" diff --git a/tests/feature_tests/test_fsdp_overlap_reorder.py b/tests/feature_tests/test_fsdp_overlap_reorder.py index 42a8adb..5e92142 100644 --- a/tests/feature_tests/test_fsdp_overlap_reorder.py +++ b/tests/feature_tests/test_fsdp_overlap_reorder.py @@ -34,11 +34,11 @@ requires_torchrun = pytest.mark.skipif(shutil.which("torchrun") is None, reason="requires torchrun") -def _run(nproc: int) -> subprocess.CompletedProcess: +def _run(nproc: int, *extra: str, port: str = "29631") -> subprocess.CompletedProcess: env = os.environ.copy() env["MAGI_LOGGING_LEVEL"] = env.get("MAGI_LOGGING_LEVEL", "info") return subprocess.run( - ["torchrun", f"--nproc_per_node={nproc}", "--master_port=29631", str(_HELPER)], + ["torchrun", f"--nproc_per_node={nproc}", f"--master_port={port}", str(_HELPER), *extra], env=env, capture_output=True, text=True, @@ -68,3 +68,18 @@ def test_reorder_multi_rank(): out = p.stdout + p.stderr assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" assert "REORDER_PASS" in p.stdout, out[-3000:] + + +@requires_cuda +@requires_torchrun +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") +def test_reorder_graph_mismatch_fail_fast(): + """world=2 with rank1 compiling a structurally DIFFERENT graph: the cross-rank + graph-fingerprint check must fire on both ranks (warning), leave the schedule + unchanged, and complete without deadlock.""" + p = _run(2, "--mismatch", port="29632") + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" + assert "REORDER_MISMATCH unchanged=True" in p.stdout, out[-3000:] + assert "REORDER_WARNED" in p.stdout, out[-3000:] # the fail-fast warning fired + assert "REORDER_PASS" in p.stdout, out[-3000:] diff --git a/tests/feature_tests/test_profiling_estimator.py b/tests/feature_tests/test_profiling_estimator.py index 8de54db..14fa02f 100644 --- a/tests/feature_tests/test_profiling_estimator.py +++ b/tests/feature_tests/test_profiling_estimator.py @@ -370,4 +370,6 @@ def test_collective_profile_accuracy_multi_rank(): assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" assert "COLL_ACCURATE ok=True" in p.stdout, out[-3000:] assert "COLL_WARMSYNC" in p.stdout and "ok=True" in p.stdout, out[-3000:] + # key-set mismatch fail-fast: no deadlock + degraded to analytical on all ranks + assert "COLL_MISMATCH ok=True" in p.stdout, out[-3000:] assert "COLL_PASS" in p.stdout, out[-3000:] From 410e734bb389b7800f4cb1fdd2c6a57db31c367a Mon Sep 17 00:00:00 2001 From: wtr Date: Thu, 16 Jul 2026 13:14:04 +0800 Subject: [PATCH 7/7] chore --- magi_compiler/profiling/benchmark_inputs.py | 79 ++--- magi_compiler/profiling/runtime_estimator.py | 302 ++++++------------- 2 files changed, 116 insertions(+), 265 deletions(-) diff --git a/magi_compiler/profiling/benchmark_inputs.py b/magi_compiler/profiling/benchmark_inputs.py index 4e4bda6..5f7bcae 100644 --- a/magi_compiler/profiling/benchmark_inputs.py +++ b/magi_compiler/profiling/benchmark_inputs.py @@ -14,76 +14,39 @@ """Per-op benchmark-input hooks for the profiling runtime estimator. -Some opaque boundary custom ops cannot be replayed for timing from generic -size-hinted tensors alone, because they carry VALUE-DEPENDENT metadata that must -be self-consistent (not just right-shaped) or they raise -- e.g. a context-parallel -attention op taking a split-sizes tensor whose values must sum to the sequence -length, else its internal all_to_all asserts and the op falls back to a 0 cost -estimate. - -The owning code (the model package that defines the custom op) registers a hook -here that builds a VALID, RANK-DETERMINISTIC set of replay inputs. MagiCompiler -stays free of model-specific knowledge: ``_measure_extern`` just looks up the op -by name and, if a hook exists, uses it instead of the generic realizer. - -A hook is ``fn(fx_node, realize) -> (args, kwargs) | None``: - * ``fx_node`` -- the ``torch.fx.Node`` for the op (read its ``args``/``kwargs`` - and their ``meta['val']`` for shapes); - * ``realize`` -- MagiCompiler's default arg realizer (fx.Node/SymInt/container - -> concrete tensor/int), so a hook can reuse it for the plain - tensor args and only special-case the metadata ones; - * returns ``(args, kwargs)`` of real objects to call the op with, or ``None`` to - fall back to the generic path. - -Hooks MUST produce rank-identical inputs (derive everything from shapes / static -sizes, no per-rank timing or state) so the rank-lockstep ``warm_and_sync`` measure -issues any internal collective in lockstep across ranks. - -Example -- a custom attention op ``mylib::attn_cp(q, k, v, split_sizes, scale)`` -whose internal all_to_all requires ``split_sizes`` values to sum to the sequence -length (a zero-filled tensor from the generic realizer would make it raise):: - - from magi_compiler.profiling import register_benchmark_inputs - - def _attn_cp_benchmark_inputs(fx_node, realize): - q_node, k_node, v_node, _split_node, scale = fx_node.args - q = realize(q_node) # plain tensor args: reuse the generic realizer - k = realize(k_node) - v = realize(v_node) - # Rebuild the VALUE-dependent arg from SHAPE hints only (rank-identical): - # a uniform split summing to the per-rank seqlen. - seq, group_size = q.shape[0], 4 - split_sizes = torch.full((group_size,), seq // group_size, dtype=torch.int32, device=q.device) - split_sizes[-1] += seq - int(split_sizes.sum()) - return (q, k, v, split_sizes, float(scale)), {} - - register_benchmark_inputs( - "mylib::attn_cp", _attn_cp_benchmark_inputs, has_internal_collective=True - ) - -Register at import time of the module that defines the op (so the hook exists -before any compile). Returning ``None`` from the hook falls back to the generic -realizer for that call. +Some custom ops cannot be replayed from generic size-hinted tensors: they carry +VALUE-DEPENDENT metadata that must be self-consistent or they raise (e.g. a CP +attention op whose split-sizes must sum to the sequence length) -- and would fall +back to a 0 cost. The model package that defines such an op registers a hook +here that builds valid replay inputs; MagiCompiler stays free of model-specific +op knowledge. + +Hook: ``fn(fx_node, realize) -> (args, kwargs) | None`` -- ``fx_node`` is the op's +``torch.fx.Node`` (shapes in ``meta['val']``), ``realize`` is the generic arg +realizer to reuse for plain tensor args; return None to fall back to the generic +path. Register at import time of the op-defining module, e.g.:: + + register_benchmark_inputs("mylib::attn_cp", _attn_cp_inputs, has_internal_collective=True) + +Hooks MUST produce rank-identical inputs (derive everything from shapes, no +per-rank state) so the rank-lockstep ``warm_and_sync`` measurement issues any +internal collective in lockstep. """ from __future__ import annotations from typing import Callable -# op name (``torch.ops.ns.op`` overload string, e.g. "mylib::attn_cp") -# -> hook. Also records which ops issue an internal collective (need fixed-iter -# lockstep replay) -- so MagiCompiler no longer hardcodes model-specific op names. +# op name (OpOverload string, e.g. "mylib::attn_cp") -> hook; plus the set of ops +# that issue an internal collective (need fixed-iter lockstep replay). _BENCHMARK_INPUT_HOOKS: dict[str, Callable] = {} _INTERNAL_COLLECTIVE_OPS: set[str] = set() def register_benchmark_inputs(op_name: str, fn: Callable, *, has_internal_collective: bool = False) -> None: """Register a replay-input builder for ``op_name`` (see module docstring). - - ``has_internal_collective``: mark this op as issuing an internal collective, so - the estimator replays it with a FIXED iteration count under barriers (keeps every - rank in lockstep -- a duration-adaptive count would desync the internal NCCL op). - """ + ``has_internal_collective``: replay with a fixed iteration count under barriers + (an adaptive count would desync the internal NCCL op across ranks).""" _BENCHMARK_INPUT_HOOKS[op_name] = fn if has_internal_collective: _INTERNAL_COLLECTIVE_OPS.add(op_name) diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index 2dde980..b10be06 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -14,56 +14,26 @@ from __future__ import annotations -"""Profiling-based ``estimate_op_runtime`` replacement. - -Inductor's default ``BaseSchedulerNode.get_estimated_runtime`` is a pure -analytical roofline. It is unreliable for exactly the nodes our FSDP overlap -reorder pass must size: - -* fused pointwise/reduction kernels -> ``estimate_flops()`` is None, so the - estimate degrades to ``bytes / dram_bw`` and IGNORES the compute fused into the - kernel (measured 60x under-estimate for a 32-deep fusion); -* matmul -> the device TFLOPS table is wrong on this box (0.5 vs ~990), giving a - ~1500x over-estimate; -* custom ops (Triton / flash-attn) -> ``ExternKernelSchedulerNode`` with - ``MultiOutputLayout`` -> dtype is None -> the estimate is silently 0. - -Full write-up + demos: ``scripts/demo/research_estimate_op_runtime_findings.md``. - -Instead we MEASURE each scheduler node's real kernel time: -* Triton (pointwise/reduction/template) snodes -> ``scheduler.benchmark_fused_nodes``, - which codegens the SAME fused kernel production emits and ``do_bench``es it - (verified 0.995-1.15x of the compiled-graph kernel); -* extern snodes (matmul / custom op) -> replay the aten/custom op on real inputs - rebuilt from the fx fake-tensor meta, timed with the Inductor benchmarker. - -COLLECTIVES are NOT benchmarked here -- they always use the analytical -``estimate_nccl_collective_runtime``. Benchmarking a real collective during compile -is not multi-rank-safe: the reorder pass runs per-rank during independent, -non-co-scheduled compilation, so issuing an NCCL op desyncs ranks -> watchdog hang -(seen on 8-GPU gaga4: rank0 raced to FSDP seqNum~194 while peers never matched). -For real measured comm, calibrate OUTSIDE compile at a synchronized point. How far -the analytical estimate is from reality is measured by -``scripts/demo/verify_collective_estimate.py`` (a standalone, properly-synchronized -harness -- not the compile path). - -Op -> time table ----------------- -The estimator maintains ``self._table: dict[key -> ProfileEntry]`` -- the -persistent op->time structure for COMPUTE nodes. The KEY is the op's STRUCTURAL -identity ``(target op, tuple[(input shape, dtype)])`` (``_structural_key``), NOT the -snode's unique name (``buf0``/``op42`` are unique per node and would defeat reuse). -Each distinct key is benchmarked ONCE on first encounter; every later isomorphic -snode (same op + same input shapes, e.g. the same matmul in every layer) reuses the -entry (``reuse_count++``). A 40-layer model thus measures O(#distinct kernels), not -O(#nodes). ``ProfileEntry`` carries ``(ns, kind, label, measured, reuse_count)``; -``estimator.summary()`` prints the whole table (see it at DEBUG log level). - -Passed to the reorder pass as ``cost_fn`` (a callable ``snode -> nanoseconds``). -The extern measurement path is ShapeEnv-isolated (real tensors from -size_hints, eager call) so they run on the dynamic base compile; only benchmark_fused_nodes -(fused Triton) would specialize the dynamic dim, so that path stays analytical while -free symbols exist. +"""Profiling-based ``estimate_op_runtime`` replacement (the reorder pass's cost_fn). + +Inductor's analytical roofline is unreliable for exactly the nodes the FSDP +overlap reorder must size: fused pointwise (~60x under), matmul (~1500x over on +this box), custom ops (silently 0). So we MEASURE: +fused Triton snodes via ``scheduler.benchmark_fused_nodes``, extern snodes +(matmul / custom op) by replaying the aten op on inputs rebuilt from fx meta. + +Collectives are never benchmarked in ``__call__`` (per-rank compile-time NCCL +desyncs ranks -> hang); they are seeded with the analytical estimate and +re-measured for real in the rank-lockstep ``warm_and_sync``. + +The op->time table (``self._table``) is keyed by STRUCTURAL identity +(op + input shapes/dtypes, ``_structural_key``) -- not the per-node name -- so +isomorphic ops across layers share one measurement: O(#distinct kernels), not +O(#nodes). ``summary()`` dumps the table (DEBUG). + +Extern measurement is ShapeEnv-isolated so it is safe on the dynamic base +compile; ``benchmark_fused_nodes`` would specialize the dynamic dim, so fused +Triton stays analytical while free symbols exist. """ import dataclasses @@ -79,10 +49,8 @@ from .benchmark_inputs import get_benchmark_inputs_hook, op_has_internal_collective -# Dedicated GLOO (CPU) group for exchanging profile metadata across ranks (built -# once, cached). A CPU/gloo group keeps the cost sync off the NCCL process groups -# the forward uses, so it can never interleave with / desync the weight-gather or -# CP collectives. +# Dedicated GLOO (CPU) group for the cost sync, built once -- keeps it off the +# NCCL process groups the forward uses (cannot desync weight-gather / CP comms). _COST_SYNC_GROUP = "uninit" @@ -112,10 +80,8 @@ class ProfileEntry: def _snode_label(snode: BaseSchedulerNode, max_shapes: int = 3) -> str: - """Short human-readable identity of an snode for the profile table / logs: - the op target plus its first few input shapes. (The snode's own NAME, e.g. - 'op123', is deliberately NOT the cache key -- it is unique per node and would - defeat cross-layer reuse; the key is the structural identity below.)""" + """Human-readable identity for the profile table: op target + first few input + shapes (for logs only; the cache key is ``_structural_key``).""" node = getattr(snode, "node", None) origin = node.get_origin_node() if (node is not None and hasattr(node, "get_origin_node")) else None target = str(getattr(origin, "target", type(node).__name__ if node is not None else "?")) @@ -132,10 +98,9 @@ def _snode_label(snode: BaseSchedulerNode, max_shapes: int = 3) -> str: def _is_multi_output_unpack(snode: BaseSchedulerNode) -> bool: - """True for a ``MultiOutput`` snode -- the zero-cost getitem that unpacks one - output of a multi-output extern (FallbackKernel). It shares its origin fx - node with the parent extern, so it must never go through the structural-key - table (the key collides with the parent's and returns the parent's runtime).""" + """Zero-cost MultiOutput getitem. Must never hit the structural-key table: + it shares its origin fx node with the parent extern, so the key collides and + would return the parent's full runtime.""" from torch._inductor.ir import MultiOutput return type(getattr(snode, "node", None)) is MultiOutput @@ -191,15 +156,10 @@ def _concrete_size(s, fallback: int = 1) -> int: def _realize_arg(v): - """Turn an fx arg into a concrete replay input (rank-deterministic). - - fx.Node(tensor) -> right-shaped tensor from size-hints; fx.Node/bare SymInt -> - concrete int; list/tuple/dict -> recursively realized PLAIN container. The - container de-immutabilization matters: FX stores list args as - torch.fx.immutable_collections.immutable_list, and the custom-op C++ arg parser - requires a plain List[int] for a ``SymInt[]`` arg -- an immutable_list of still- - symbolic / nested elements is rejected ("Expected List[int] ... found - immutable_list"), making the op fall back to a 0 cost.""" + """fx arg -> concrete replay input: Node(tensor) -> right-shaped tensor from + size-hints; SymInt -> concrete int; containers -> recursively realized PLAIN + list/tuple/dict. Plain matters: the custom-op C++ parser rejects an fx + immutable_list where ``SymInt[]`` expects List[int] (op would cost 0).""" if isinstance(v, torch.fx.Node): ev = v.meta.get("val") if isinstance(ev, torch.Tensor): @@ -223,16 +183,13 @@ def _realize_arg(v): def _measure_extern(snode: ExternKernelSchedulerNode, fixed_iters: bool = False) -> float: """Time an extern (matmul / custom-op) snode by replaying its aten op. - ``fixed_iters``: when True, time a CONSTANT number of iterations with CUDA events - instead of the duration-adaptive ``benchmark_gpu``. REQUIRED for ops that issue - an INTERNAL collective (e.g. the CP ``all_to_all`` inside gaga4_fa3_with_sink_cp): - the adaptive benchmarker runs a rank-dependent iteration count, so different ranks - would issue different numbers of that internal collective -> NCCL count mismatch -> - deadlock. A fixed count keeps every rank in lockstep (mirrors _measure_collective_op). + ``fixed_iters=True``: constant iteration count with CUDA events instead of the + duration-adaptive benchmarker. Required for ops with an INTERNAL collective + (CP all_to_all inside attention/MoE): adaptive iteration counts differ per rank + -> NCCL count mismatch -> deadlock. - Ops whose replay needs VALUE-CONSISTENT metadata (not just right-shaped tensors) - can register a hook via ``benchmark_inputs.register_benchmark_inputs`` that builds - a valid ``(args, kwargs)``; otherwise args come from the generic ``_realize_arg``.""" + Ops whose replay needs value-consistent metadata register a hook via + ``benchmark_inputs.register_benchmark_inputs``; otherwise ``_realize_arg``.""" fx_node = snode.node.get_origin_node() if fx_node is None: return 0.0 @@ -246,17 +203,11 @@ def _measure_extern(snode: ExternKernelSchedulerNode, fixed_iters: bool = False) args = tuple(_realize_arg(a) for a in fx_node.args) kwargs = {k: _realize_arg(v) for k, v in fx_node.kwargs.items()} - # Replay EAGERLY, decoupled from the enclosing compile: - # * torch._dynamo.disable(): _measure_extern runs while the OUTER graph is being - # compiled (Dynamo/Inductor active). A custom boundary op whose impl contains - # torch.compile'd regions (e.g. gaga4_fa3_with_sink_cp's sink-correction path) - # would otherwise RE-ENTER Dynamo on these concrete tensors, re-trace with - # DYNAMIC shapes (symbolic s*), and blow up ("Dynamo failed to run FX node ... - # broadcast" / "can't pickle cyclic objects") -> the op fell back to a 0 cost. - # We want the EAGER kernel time, so disable Dynamo for the replay. - # * no_grad: the compiled forward runs under inference_mode; some ops branch on - # torch.is_grad_enabled() (fa3 takes a training-only flash-attn wrapper path when - # grad is on, incompatible with the installed flash-attn). Match inference. + # Replay eagerly, decoupled from the enclosing compile: + # * dynamo.disable: an op whose impl contains torch.compile'd regions would + # re-enter Dynamo mid-compile and blow up; we want the eager kernel time. + # * no_grad: match the inference forward -- some ops branch on + # torch.is_grad_enabled() into incompatible paths. @torch._dynamo.disable def _call(): return target(*args, **kwargs) @@ -296,10 +247,9 @@ def _op_name(target) -> str: def _extern_has_internal_collective(snode: BaseSchedulerNode) -> bool: - """True for opaque boundary ops that issue collectives internally (CP attention / - MoE), so we must measure them with fixed iterations under a barrier. The set of - such ops is declared by the owning code via ``register_benchmark_inputs(..., - has_internal_collective=True)`` -- MagiCompiler holds no model-specific op names.""" + """Ops that issue collectives internally (declared via + ``register_benchmark_inputs(..., has_internal_collective=True)``) must be + measured with fixed iterations under a barrier.""" node = getattr(snode, "node", None) origin = node.get_origin_node() if (node is not None and hasattr(node, "get_origin_node")) else None target = getattr(origin, "target", None) if origin is not None else None @@ -377,10 +327,8 @@ def fn(): def fn(): _WAIT(_AG(ins[0], group_size, group_name)) - # FIXED iteration counts across ALL ranks. A duration-based benchmarker - # (e.g. benchmark_gpu shrinks benchmark_iters by per-rank estimated runtime) - # would issue a DIFFERENT number of collectives on different ranks -> - # NCCL count mismatch -> deadlock. A constant loop keeps every rank in lockstep. + # Fixed iteration count on all ranks -- an adaptive benchmarker would issue + # different numbers of collectives per rank -> NCCL count mismatch -> deadlock. _WARMUP, _ITERS = 3, 10 for _ in range(_WARMUP): fn() @@ -396,37 +344,24 @@ def fn(): class ProfilingRuntimeEstimator: - """Callable ``snode -> ns`` for ``config.estimate_op_runtime``. - - Measures compute nodes with real kernels (memoized); defers collectives and - waits to the analytical ``get_estimated_runtime`` (the reorder pass sizes - comm separately via ``estimate_nccl_collective_runtime`` or a calibrated - value). Never raises -- on any failure returns the analytical estimate (or - 0), so it degrades to today's behaviour rather than breaking compilation. - """ + """Callable ``snode -> ns`` (see module docstring). Never raises -- any + measurement failure falls back to the analytical estimate.""" def __init__(self) -> None: - # The op -> time table. Key is the STRUCTURAL identity of the op - # (target op + input shapes/dtypes for compute, or (coll, group, group_size, - # in-shapes) for collectives) -- NOT the snode's unique name, so isomorphic - # ops across repeated layers share ONE measurement. Each distinct key is - # profiled once on first encounter; later encounters reuse the entry. + # op -> time table, keyed by structural identity (see module docstring). self._table: dict[tuple, ProfileEntry] = {} self.n_measured = 0 self.n_cache_hits = 0 - # Set True by MagiBackend for multi-rank runs: the reorder pass then calls - # warm_and_sync() to reconcile costs across ranks (rank-identical schedule). + # True (profile_sync): the reorder pass calls warm_and_sync() to reconcile + # costs across ranks. self._sync_across_ranks = False - # Transient {structural_key -> representative snode}, used ONLY by - # warm_and_sync() to re-measure in rank-lockstep order. Kept OFF the - # ProfileEntry (which Inductor pickles into the fx-graph cache key) because - # snodes hold FakeTensors that are unpicklable / cyclic. Never serialized. + # Transient {key -> representative snode} for warm_and_sync re-measurement. + # Kept OFF ProfileEntry: snodes hold unpicklable FakeTensors and the entry + # is pickled into the fx-graph cache key. self._key_snode: dict = {} def __deepcopy__(self, memo): - # Shared by reference from the reorder pass; when config serialization - # deepcopies the pass list, return a clean instance (the memoized - # measurements are transient and hold no tensors, but avoid copying them). + # Config serialization deepcopies the pass list; return a clean instance. new = ProfilingRuntimeEstimator() new._sync_across_ranks = self._sync_across_ranks memo[id(self)] = new @@ -438,37 +373,18 @@ def table(self) -> "dict[tuple, ProfileEntry]": return self._table def warm_and_sync(self) -> int: - """Rank-LOCKSTEP profiling of every distinct op, so multi-rank runs get REAL - measured costs (more accurate than the analytical roofline) AND a rank- - identical cost table (required so the FSDP-overlap reorder produces the same - schedule on every rank -- else weight-PG gathers interleave with the eager CP - all_to_all in rank-divergent order -> deadlock). - - PRECONDITION: the graph is structurally IDENTICAL on all ranks, so every rank - has exactly the same set of ``_structural_key`` keys already populated in the - table by the reorder pass's warm-up loop, so: - * the key iteration order (sorted) is identical on every rank; - * for a collective-containing op (attention/MoE), every rank measures it at - the same step, so the barrier-wrapped fixed-iteration replay issues the - op's internal CP all_to_all in lockstep (no NCCL count mismatch); - * the max-reduce over gloo is symmetric. - The precondition is VERIFIED (not assumed): the key sets are all_gather'd and - compared BEFORE the barrier loop. On any mismatch (per-rank shapes from - uneven Shard(0) padding, a ``_structural_key``/``_collective_spec`` failure on - one rank, ...) entering the loop would deadlock -- rank A's Nth barrier pairs - with rank B's all_gather (gloo count mismatch), or both ranks measure a - DIFFERENT collective-containing op at the same step (NCCL mismatch inside - ``_measure_one``). Instead we warn and DEGRADE to the analytical estimate for - every entry this compile measured per-rank: analytical is a pure shape+device - function, so the resulting cost table is rank-deterministic and the reorder - stays deadlock-free (same guarantee as fsdp_config.cost_mode=analytical). - - Steps: verify key sets match across ranks; then for each key in sorted order, - barrier -> (re)measure locally with FIXED iters -> barrier; then - all_gather_object the {key: ns} maps and take the MAX. Fixed-iteration timing - (``_measure_one``) is mandatory -- a duration-adaptive benchmark would run a - per-rank iteration count and desync the internal collective. Returns the - number of table entries whose cost changed.""" + """Rank-lockstep re-measurement of every table entry, giving REAL costs + that are also rank-identical (required for a rank-identical reorder + schedule). Steps: verify key sets match across ranks; per sorted key, + barrier -> measure with FIXED iters -> barrier; finally all_gather the + {key: ns} maps and take the max. Returns #entries whose cost changed. + + The key-set check is a hard precondition: with divergent key sequences the + barrier loop deadlocks (barrier/all_gather count mismatch on gloo, or two + ranks lockstep-measuring DIFFERENT internal-collective ops -> NCCL + mismatch). On mismatch we warn and degrade this compile's entries to the + analytical estimate -- rank-deterministic, same guarantee as + fsdp_config.cost_mode=analytical.""" import torch.distributed as dist if not (dist.is_available() and dist.is_initialized()): @@ -478,16 +394,11 @@ def warm_and_sync(self) -> int: return 0 group = _get_cost_sync_group() - # The reorder warm-up already populated self._table (one entry per distinct - # structural key) and self._key_snode (a representative snode per key). - # Re-measure each in a rank-uniform (sorted) order under barriers so any - # internal collective is issued in lockstep across ranks. keys = sorted(self._table.keys(), key=repr) - # FAIL-FAST precondition check: the barrier loop below is only lockstep-safe - # if every rank iterates the SAME key sequence. all_gather_object is a single - # symmetric collective (safe regardless of key sets), and every rank sees the - # same gathered result, so all ranks take the same branch. + # Fail-fast key-set check (see docstring). all_gather_object is a single + # symmetric collective and every rank sees the same result, so all ranks + # take the same branch. key_reprs = [repr(k) for k in keys] all_key_reprs: list = [None] * world dist.all_gather_object(all_key_reprs, key_reprs, group=group) @@ -530,9 +441,7 @@ def warm_and_sync(self) -> int: local_ns[k] = self._table[k].ns # no cached snode -> keep prior measurement dist.barrier(group=group) - # keys re-measured on THIS rank (had a representative snode) -- gather across - # ranks so a key measured on any rank is flagged measured (the graph is - # identical, so this is the same set everywhere; the union is just defensive). + # Union of measured keys across ranks -> flag entries measured=True. measured_here = set(self._key_snode.keys()) gathered: list = [None] * world @@ -559,10 +468,8 @@ def warm_and_sync(self) -> int: return n def _measure_one(self, snode: BaseSchedulerNode) -> float: - """Measure a single snode with FIXED iterations (lockstep-safe), never raising - -- falls back to the analytical estimate. Collective-containing externs - (attention / MoE) use fixed-iter replay so every rank issues the same number - of internal collectives.""" + """Lockstep-safe single measurement (fixed iters for anything containing a + collective); never raises -- falls back to the analytical estimate.""" try: if contains_collective(snode): return _measure_collective_op(snode) @@ -578,11 +485,9 @@ def _measure_one(self, snode: BaseSchedulerNode) -> float: return _safe_analytical(snode) def summary(self) -> str: - """One line per distinct profiled op: kind, label, per-call time, #calls, - aggregate time (per-call * #calls). Each op also gets a machine-parseable - ``ESTLINE`` tag (kind|label|per_call_us|calls|total_us|measured) so a run's - estimates can be diffed against the real nsys kernel trace -- see - scripts/demo/compare_estimate_vs_nsys.py.""" + """One line per distinct op + a machine-parseable ``ESTLINE`` tag + (kind|label|per_call_us|calls|total_us|measured) for diffing against an + nsys trace.""" lines = [] for e in sorted(self._table.values(), key=lambda e: -e.ns * (e.reuse_count + 1)): calls = e.reuse_count + 1 # first encounter + reuses @@ -626,17 +531,13 @@ def __call__(self, snode: BaseSchedulerNode) -> float: is_extern = isinstance(snode, ExternKernelSchedulerNode) - # Benchmark compute nodes on the dynamic graph without specializing the - # dynamic dim. Extern (matmul / custom op) is ShapeEnv-isolated (replays the - # aten op on size-hinted real tensors) so it is safe even with free symbols. - # Fused Triton pointwise/reduction goes through benchmark_fused_nodes which - # re-enters Inductor codegen bound to the live ShapeEnv and WOULD specialize, - # so that path stays analytical while the graph is dynamic. Matmul/attention/ - # MoE (the dominant costs, and the source of the bogus roofline) are extern. + # Extern replay is ShapeEnv-isolated -> safe with free symbols; fused + # Triton (benchmark_fused_nodes) would specialize the dynamic dim, so it + # stays analytical while the graph is dynamic. if not is_extern and _graph_has_free_symbols(): return _safe_analytical(snode) - # --- op -> time table: profile a distinct key ONCE, reuse afterwards --- + # op -> time table: profile a distinct key once, reuse afterwards. key = _structural_key(snode) if key is not None: entry = self._table.get(key) @@ -645,9 +546,8 @@ def __call__(self, snode: BaseSchedulerNode) -> float: self.n_cache_hits += 1 return entry.ns - # First encounter of this key -> measure it. Measuring must NEVER break - # compilation: any failure (unbenchmarkable extern, FakeTensor deepcopy in - # the benchmark harness, ...) falls back to the analytical estimate. + # First encounter -> measure; any failure falls back to analytical + # (measuring must never break compilation). measured = True try: ns = self._measure_extern_safe(snode) if is_extern else self._measure(snode) @@ -660,9 +560,7 @@ def __call__(self, snode: BaseSchedulerNode) -> float: kind = "extern" if is_extern else "compute" label = _snode_label(snode) self._table[key] = ProfileEntry(ns=ns, kind=kind, label=label, measured=measured) - # Remember a representative snode (transient, unpicklable) so - # warm_and_sync() can re-measure this key in rank-lockstep order. - if self._sync_across_ranks: + if self._sync_across_ranks: # stash a representative snode for warm_and_sync self._key_snode[key] = snode magi_logger.debug( "profile[%s] %s -> %.2fus%s", kind, label, ns / 1e3, "" if measured else " (analytical fallback)" @@ -676,11 +574,9 @@ def _measure_extern_safe(self, snode: BaseSchedulerNode) -> float: return ns def _measure(self, snode: BaseSchedulerNode) -> float: - # Benchmarking runs real kernels at concrete (hinted) shapes. Under - # dynamic-shape compilation that would add ``Eq(sym, hint)`` guards and - # replacements into the ShapeEnv and SPECIALIZE the dynamic dim (breaking - # the compile for other seq lens). Suppress guard creation AND snapshot / - # restore the ShapeEnv's mutable specialization state so nothing leaks. + # Benchmarking at hinted concrete shapes must not leak Eq(sym, hint) + # guards/replacements into the live ShapeEnv (would specialize the dynamic + # dim): suppress guards + snapshot/restore the mutable state. with _shapeenv_sandbox(), _suppress_guards(): return self._measure_inner(snode) @@ -707,10 +603,8 @@ def _safe_analytical(snode: BaseSchedulerNode) -> float: def _graph_has_free_symbols() -> bool: - """True if the graph being compiled still has symbolic (dynamic) shapes. - - Benchmarking real kernels is only safe when everything is concrete; a - dynamic-shape compile has free symbols and must use the analytical estimate.""" + """True if the compile still has dynamic (symbolic) shapes -- any symbol with + a non-singleton range that is not yet a constant replacement.""" try: shape_env = V.graph.sizevars.shape_env except Exception: # noqa: BLE001 @@ -718,9 +612,6 @@ def _graph_has_free_symbols() -> bool: if shape_env is None: return False try: - # A fully-static graph has no unbacked/free symbols with a non-singleton - # range. If ANY symbol has a range wider than one value and is not yet a - # constant replacement, the graph is dynamic and must not be benchmarked. replacements = getattr(shape_env, "replacements", {}) for sym, vr in shape_env.var_to_range.items(): if sym in replacements: @@ -740,9 +631,8 @@ def _graph_has_free_symbols() -> bool: def _suppress_guards(): - """Context manager that suppresses ShapeEnv guard creation while we run real - benchmark kernels at hinted shapes, so measurement never specializes a - dynamic dim. Falls back to a no-op when no ShapeEnv is active.""" + """Suppress ShapeEnv guard creation during benchmarking (no-op without a + live ShapeEnv).""" from contextlib import nullcontext try: @@ -754,8 +644,8 @@ def _suppress_guards(): return nullcontext() -# ShapeEnv mutable fields that benchmarking could pollute with a `s -> hint` -# specialization; snapshot/restore them so a measurement never persists a guard. +# ShapeEnv mutable fields a benchmark could pollute with an `s -> hint` +# specialization; snapshotted/restored by _shapeenv_sandbox. _SHAPEENV_STATE_FIELDS = ( "guards", "axioms", @@ -769,10 +659,8 @@ def _suppress_guards(): class _shapeenv_sandbox: - """Snapshot the live ShapeEnv's specialization state on enter, restore it on - exit. Shallow-copies the mutable dict/list fields (their *contents* are - replaced wholesale by the compile, never mutated in place after we restore), - so benchmarking a kernel at a hinted concrete shape cannot leak an + """Snapshot the live ShapeEnv's specialization state on enter, restore on + exit, so a benchmark at hinted concrete shapes cannot leak an ``Eq(sym, hint)`` replacement/guard into the real compile.""" def __init__(self) -> None: