diff --git a/magi_compiler/config.py b/magi_compiler/config.py index 44fbd03..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") @@ -263,6 +321,21 @@ class CompileConfig(BaseSettings): "Each sub-graph between two splitting ops is compiled independently by Inductor." ), ) + disable_graph_split: bool = Field( + False, + description=( + "Skip FX-level splitting at the custom subgraph-boundary ops (splitting_ops) and hand the " + "WHOLE graph to Inductor as a single piecewise submodule." + ), + ) + # ---- Whole-graph FSDP overlap ---- + fsdp_config: FSDPConfig = Field( + FSDPConfig(), + description=( + "Whole-graph FSDP weight all-gather / compute overlap (lowering + bucketing + reorder + cost " + "model). Fields reachable via MAGI_COMPILE_FSDP_CONFIG__ env vars." + ), + ) # ---- 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..ff4b44d 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -542,14 +542,70 @@ def _init_cache(self) -> str: self.local_magi_cache_path.mkdir(parents=True, exist_ok=True) self.compiler_manager.initialize_cache(self.local_magi_cache_path) + def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: + """Whole-graph FSDP all-gather / compute overlap (disable_graph_split path). + + 1. Lower SimpleFSDP weight prim_redistribute -> explicit collectives and + optionally bucket them 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("fsdp_config.enable_fullgraph_overlap requires disable_graph_split=True") + + if self.compile_config.cudagraph_mode != CudaGraphMode.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 + + 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", + fsdp_cfg.bucket_mode, + fsdp_cfg.bucket_size_mib, + n_buckets, + ) + + if fsdp_cfg.cost_mode == "analytical": + self._fsdp_overlap_estimator = None + cost_fn = None # reorder pass defaults to Inductor's analytical estimate + else: # "profile_sync" -- the default for BOTH single- and multi-rank + self._fsdp_overlap_estimator = ProfilingRuntimeEstimator() + self._fsdp_overlap_estimator._sync_across_ranks = True + cost_fn = self._fsdp_overlap_estimator + + reorder = FsdpOverlapReorder( + 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] + @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. + if self.compile_config.disable_graph_split: + if self.compile_config.cudagraph_mode == CudaGraphMode.PIECEWISE: + raise ValueError("disable_graph_split is incompatible with cudagraph_mode=PIECEWISE") + fx_split_ops = [] + magi_logger.info( + "disable_graph_split=True: skipping FX-level graph split; compiling the whole graph as one submod" + ) + else: + fx_split_ops = self.compile_config.splitting_ops or [] resolved_ops: list[torch._ops.OpOverload] = resolve_defined_ops(fx_split_ops) magi_logger.info(f"Setting up FX-level graph split with ops: {fx_split_ops=}") magi_logger.info(f"Resolved splitting ops for FX-level graph split: {resolved_ops=}") + # Step 1.4: whole-graph FSDP overlap. + 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. subgraph_id = 0 node_to_subgraph_id = {} diff --git a/magi_compiler/passes/fsdp_overlap/__init__.py b/magi_compiler/passes/fsdp_overlap/__init__.py new file mode 100644 index 0000000..d507a8a --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/__init__.py @@ -0,0 +1,25 @@ +# 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 .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", + "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..c40b030 --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py @@ -0,0 +1,226 @@ +# 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.""" + 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_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 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 + 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_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.""" + 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) + + 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 + + outs = [] + for i, am in enumerate(ag_metas): + gi = graph.graph.call_function(operator.getitem, (coalesced, i)) + gi.meta["example_value"] = am + 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 + old_wait.replace_all_uses_with(wait_i) + + for ag_old, wait_old in zip(ag_nodes, waits): + graph.graph.erase_node(wait_old) + graph.graph.erase_node(ag_old) + + +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)} + + # 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 + _, _world, group_name = node.args + groups[group_name].append(node) + + buckets = 0 + 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_index, sub) + buckets += 1 + + if buckets: + graph.graph.lint() + graph.recompile() + magi_logger.info( + "FSDP weight all-gather bucketing (coalesced): created %d coalesced buckets (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..34123dc --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py @@ -0,0 +1,54 @@ +# 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 torch.fx as fx + +from magi_compiler.utils import magi_logger + +from .bucket_all_gather import bucket_weight_all_gather_coalesced +from .redistribute_lowering import lower_prim_redistribute_to_collectives + + +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 across the WHOLE graph (no subgraph partitioning). + + ``bucket_mode``: + * ``"none"`` -- lowering only (N individual all_gather + N waits). + * ``"coalesced"`` -- one all_gather_into_tensor_coalesced per bucket + (ONE launch, N getitems, N waits). + + ``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. + """ + 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 + + if bucket_mode == "coalesced": + 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'") + + magi_logger.info("Whole-graph FSDP bucketing (%s): created %d buckets", bucket_mode, n) + return n diff --git a/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py new file mode 100644 index 0000000..8a5e262 --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py @@ -0,0 +1,183 @@ +# 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 + +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). + # + # 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)) + 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/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py new file mode 100644 index 0000000..70b5fb2 --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -0,0 +1,467 @@ +# 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 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 + +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 +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 + + +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.""" + + def __init__(self, slack_ns: float = _DEFAULT_SLACK_NS, cost_fn=None, comm_contention_factor: float = 1.0) -> None: + self.slack_ns = slack_ns + # 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: 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-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): + # 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 + 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)} + + # 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) + n_changed = self._cost_fn.warm_and_sync() + self._cost_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 (see module docstring) ---- + launches_in_order = sorted(launches, key=lambda s: index_of[s]) # original program 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)) + + 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.comm_contention_factor + 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 + # 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 + # 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 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", + ) + + # 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: + continue + target, group = targets[launch] + if target < running: + target = running + targets[launch] = (target, group) + running = target + + # 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: + 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. + if _debug_enabled() and hasattr(self._cost_fn, "summary"): + magi_logger.debug("FSDP overlap %s", self._cost_fn.summary()) + 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 (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: + 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: + """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] + 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/profiling/__init__.py b/magi_compiler/profiling/__init__.py new file mode 100644 index 0000000..467b3bf --- /dev/null +++ b/magi_compiler/profiling/__init__.py @@ -0,0 +1,19 @@ +# 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 .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..5f7bcae --- /dev/null +++ b/magi_compiler/profiling/benchmark_inputs.py @@ -0,0 +1,60 @@ +# 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 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 (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``: 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) + + +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..b10be06 --- /dev/null +++ b/magi_compiler/profiling/runtime_estimator.py @@ -0,0 +1,696 @@ +# 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 (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 +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 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" + + +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: + """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 "?")) + 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 _is_multi_output_unpack(snode: BaseSchedulerNode) -> bool: + """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 + + +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): + """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): + 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=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 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 + 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: + # * 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) + + 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: + """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 + 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 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() + 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`` (see module docstring). Never raises -- any + measurement failure falls back to the analytical estimate.""" + + def __init__(self) -> None: + # op -> time table, keyed by structural identity (see module docstring). + self._table: dict[tuple, ProfileEntry] = {} + self.n_measured = 0 + self.n_cache_hits = 0 + # True (profile_sync): the reorder pass calls warm_and_sync() to reconcile + # costs across ranks. + self._sync_across_ranks = False + # 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): + # Config serialization deepcopies the pass list; return a clean instance. + 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 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()): + return 0 + world = dist.get_world_size() + if world <= 1: + return 0 + group = _get_cost_sync_group() + + keys = sorted(self._table.keys(), key=repr) + + # 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) + 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) + 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) + + # Union of measured keys across ranks -> flag entries measured=True. + 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: + """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) + 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 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 + 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) + + if _is_multi_output_unpack(snode): + return 0.0 + + 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) + + # 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. + 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 -> 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) + 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) + 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)" + ) + 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 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) + + 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 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 + return False + if shape_env is None: + return False + try: + 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(): + """Suppress ShapeEnv guard creation during benchmarking (no-op without a + live ShapeEnv).""" + 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 a benchmark could pollute with an `s -> hint` +# specialization; snapshotted/restored by _shapeenv_sandbox. +_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 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: + 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/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/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 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..bc82efb --- /dev/null +++ b/tests/feature_tests/fsdp_overlap_helper/e2e_helper.py @@ -0,0 +1,173 @@ +# 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 -> whole-graph bucketing (coalesced) -> FsdpOverlapReorder + +Model: two Linear layers whose params are Shard(0) DTensors (via torchtitan +``data_parallel(mode="fully_shard")``), with an +opaque ``@magi_register_custom_op(is_subgraph_boundary=True)`` op between them so the +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] + +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: 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 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__() + 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, + 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"]) + 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) + + # 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.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. + 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..45e2a85 --- /dev/null +++ b/tests/feature_tests/fsdp_overlap_helper/estimator_collective_helper.py @@ -0,0 +1,188 @@ +# 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, 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 + +Markers (rank 0): + COLL_MEASURED est_us= real_us= ratio= + COLL_ACCURATE ok= + COLL_WARMSYNC reconciled= ok= + COLL_MISMATCH 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 ProfileEntry, _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) + + # 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 and mismatch_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(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: + 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..3ade6b6 --- /dev/null +++ b/tests/feature_tests/fsdp_overlap_helper/reorder_helper.py @@ -0,0 +1,163 @@ +# 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. + +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 +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: + 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() + 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) + + 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, 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 + + 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 + _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) + 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 + + dist.barrier() + dist.destroy_process_group() + raise SystemExit(rc) + + +if __name__ == "__main__": + main() 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..4d3d00a --- /dev/null +++ b/tests/feature_tests/test_fsdp_overlap_bucket.py @@ -0,0 +1,198 @@ +# 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. + +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 + +import pytest +import torch +import torch.fx as fx + +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 +_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, compute_before?}``. For each spec + we emit weight_shard placeholder -> all_gather_into_tensor (tagged + 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 = [] + 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("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"] + 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 +# --------------------------------------------------------------------------- +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) + n_buckets = bucket_weight_all_gather_coalesced(gm) + + 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}]) + n_buckets = bucket_weight_all_gather_coalesced(gm) + 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}, + ] + ) + 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 = 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) + # 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(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) + n2 = bucket_weight_all_gather_coalesced(gm2, bucket_size_bytes=128) + assert n2 == 2 + assert _n(gm2, _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() + + +# --------------------------------------------------------------------------- +# 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' -> 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_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 + + +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..954c47c --- /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 ``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 +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.""" + 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..5e92142 --- /dev/null +++ b/tests/feature_tests/test_fsdp_overlap_reorder.py @@ -0,0 +1,85 @@ +# 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, *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}", f"--master_port={port}", str(_HELPER), *extra], + 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:] + + +@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 new file mode 100644 index 0000000..14fa02f --- /dev/null +++ b/tests/feature_tests/test_profiling_estimator.py @@ -0,0 +1,375 @@ +# 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:] + # 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:] 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