diff --git a/magi_compiler/config.py b/magi_compiler/config.py index f833cdf..b81fd37 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -256,7 +256,7 @@ class FSDPConfig(BaseModel): ), ) comm_overlap_window_scale: float = Field( - 1.5, + 1.0, ge=1.0, description=( "Multiplier on each collective's estimated runtime when sizing its compute window " @@ -264,6 +264,21 @@ class FSDPConfig(BaseModel): "with the compute that hides them (~1.4-1.5x slower in-situ on 8xH100)." ), ) + transport: Literal["nccl", "copy_engine"] = Field( + "nccl", + description=( + "Weight all-gather path. 'nccl': SM kernels. 'copy_engine': bind weights into " + "symmetric memory and gather with peer copy-engine reads (0 SM, no per-step barrier). " + "Needs NVLink on the FSDP mesh dim and static weights; unbound weights stay on NCCL." + ), + ) + symm_min_shard_mib: int = Field( + 0, + ge=0, + description=( + "Minimum local-shard MiB to bind for copy_engine. Smaller shards stay on NCCL. " "0 = bind every eligible weight." + ), + ) def _find_cutlass_root() -> str: diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 8a487d2..4e2c08a 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -592,12 +592,13 @@ 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: + def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule, example_inputs) -> 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). + 1. Lower SimpleFSDP weight prim_redistribute -> explicit collectives, bind + the gathered weights into symmetric memory when the transport is the copy + engine, 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 @@ -615,9 +616,17 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: 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) + n_buckets = lower_and_bucket_full_graph( + graph, + fsdp_cfg.bucket_mode, + bucket_size_bytes=bucket_size_bytes, + transport=fsdp_cfg.transport, + example_inputs=example_inputs, + min_shard_bytes=int(fsdp_cfg.symm_min_shard_mib) * 1024 * 1024, + ) magi_logger.info( - "FSDP fullgraph overlap: bucket_mode=%s bucket_size=%d MiB created %d buckets", + "FSDP fullgraph overlap: transport=%s bucket_mode=%s bucket_size=%d MiB created %d buckets", + fsdp_cfg.transport, fsdp_cfg.bucket_mode, fsdp_cfg.bucket_size_mib, n_buckets, @@ -640,7 +649,7 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: 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]]: + def _split_graph(self, graph: fx.GraphModule, example_inputs) -> tuple[fx.GraphModule, list[SplitItem]]: # Step 1: resolve the splitting ops. if self.compile_config.disable_graph_split: assert ( @@ -658,7 +667,7 @@ def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[Spli # Step 1.4: whole-graph FSDP overlap. if self.compile_config.fsdp_config.enable_fsdp: - self._apply_fsdp_fullgraph_overlap(graph) + self._apply_fsdp_fullgraph_overlap(graph, example_inputs) # Step 2: split graph by ops, we split graph based on resolved_ops, which becomes the partitioned single graph. subgraph_id = 0 @@ -719,7 +728,7 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> MagiSerializableFun self.full_graph_pass_manager(graph) - split_gm, piecewise_graphs = self._split_graph(graph) + split_gm, piecewise_graphs = self._split_graph(graph, example_inputs) submod_names_to_compile = [item.submod_name for item in piecewise_graphs if not item.is_splitting_graph] compilation_counter.num_piecewise_graphs_seen += len(piecewise_graphs) diff --git a/magi_compiler/passes/fsdp_overlap/__init__.py b/magi_compiler/passes/fsdp_overlap/__init__.py index d507a8a..245ca4f 100644 --- a/magi_compiler/passes/fsdp_overlap/__init__.py +++ b/magi_compiler/passes/fsdp_overlap/__init__.py @@ -13,13 +13,24 @@ # limitations under the License. from .bucket_all_gather import bucket_weight_all_gather_coalesced +from .copy_engine import bind_weights_for_copy_engine, copy_engine_weight_candidates, rewrite_weight_ag_to_copy_engine from .lower_and_bucket import lower_and_bucket_full_graph +from .node_meta import CE_BOUND, UNEVEN_SHARD, WEIGHT_AG, is_ce_bound, is_uneven_shard, is_weight_ag from .redistribute_lowering import lower_prim_redistribute_to_collectives from .reorder import FsdpOverlapReorder __all__ = [ + "bind_weights_for_copy_engine", "bucket_weight_all_gather_coalesced", + "copy_engine_weight_candidates", "lower_prim_redistribute_to_collectives", "lower_and_bucket_full_graph", + "rewrite_weight_ag_to_copy_engine", "FsdpOverlapReorder", + "CE_BOUND", + "UNEVEN_SHARD", + "WEIGHT_AG", + "is_ce_bound", + "is_uneven_shard", + "is_weight_ag", ] diff --git a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py index c4ef6f6..35ea330 100644 --- a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py +++ b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py @@ -22,6 +22,8 @@ from magi_compiler.utils import magi_logger +from .node_meta import is_ce_bound, is_uneven_shard, is_weight_ag, mark_ce_bound, mark_weight_ag + _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 @@ -64,19 +66,18 @@ 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"): + if is_weight_ag(node): 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") + """Bytes one rank contributes to a weight all_gather""" + m = ag.meta.get("example_value") if m is None: return 0 - return int(m.numel()) * int(m.element_size()) + _loc, world, _group = ag.args + return int(m.numel()) * int(m.element_size()) // int(world) def _split_by_dtype_and_size( @@ -157,8 +158,9 @@ def _coalesce_one_bucket(graph: fx.GraphModule, node_index: dict[fx.Node, int], 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 + mark_weight_ag(coalesced, uneven=any(is_uneven_shard(ag) for ag in ag_nodes)) + if all(is_ce_bound(ag) for ag in ag_nodes): + mark_ce_bound(coalesced) outs = [] for i, am in enumerate(ag_metas): @@ -177,7 +179,7 @@ def _coalesce_one_bucket(graph: fx.GraphModule, node_index: dict[fx.Node, int], graph.graph.erase_node(ag_old) -def bucket_weight_all_gather_coalesced(graph: fx.GraphModule, bucket_size_bytes: int = 0) -> int: +def bucket_weight_all_gather_coalesced(graph: fx.GraphModule, bucket_size_bytes: int = 0, split_by=None) -> 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`` @@ -193,22 +195,26 @@ def bucket_weight_all_gather_coalesced(graph: fx.GraphModule, bucket_size_bytes: are re-pointed from each old wait to ``wait_i``; the launch + getitems stay together so ``FsdpOverlapReorder`` later moves them as one unit. + ``split_by(node)`` adds a second bucket key, used in copy-engine mode to keep + bound and unbound gathers apart -- a bucket is one submission, so it cannot span + two transports. + 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) + # Key by (group_name, transport class); dtype breaks buckets positionally + # inside _split_by_dtype_and_size (strict program-adjacency). + groups: dict[tuple, 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) + groups[(group_name, split_by(node) if split_by is not None else None)].append(node) buckets = 0 - for group_name, ag_nodes in groups.items(): + for ag_nodes in groups.values(): 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) diff --git a/magi_compiler/passes/fsdp_overlap/copy_engine.py b/magi_compiler/passes/fsdp_overlap/copy_engine.py new file mode 100644 index 0000000..7be2467 --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/copy_engine.py @@ -0,0 +1,108 @@ +# 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. + +"""Graph-side copy-engine transport: bind eligible weights, then retarget tagged gathers. + +``bind_weights_for_copy_engine`` runs between lowering and bucketing so a bucket +stays all-bound or all-unbound. ``rewrite_weight_ag_to_copy_engine`` retargets +whatever carries ``node_meta.CE_BOUND``. Operators live in ``symm_mem/``. +""" + +from __future__ import annotations + +from typing import Any, Sequence + +import torch +import torch.fx as fx + +from magi_compiler.utils import magi_logger + +from .node_meta import is_ce_bound, is_uneven_shard, is_weight_ag, mark_ce_bound + +_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default +_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default + + +def _weight_holder(src) -> fx.Node | None: + """The placeholder/get_attr behind ``to_local(...)``, or None for anything else.""" + if not isinstance(src, fx.Node) or src.op != "call_method" or src.target != "to_local": + return None + owner = src.args[0] if src.args else None + return owner if isinstance(owner, fx.Node) and owner.op in ("placeholder", "get_attr") else None + + +def copy_engine_weight_candidates(graph: fx.GraphModule) -> list[tuple[fx.Node, fx.Node]]: + """Weight all-gathers the copy engine could serve, as ``(gather, holder)`` pairs. + + Holder must be exactly ``to_local(placeholder|get_attr)``: anything in between + is a temporary with no peer views. Binding attempts only these. + """ + candidates: list[tuple[fx.Node, fx.Node]] = [] + for node in graph.graph.nodes: + if node.op != "call_function" or node.target is not _ALL_GATHER: + continue + if not is_weight_ag(node) or is_uneven_shard(node): + continue + holder = _weight_holder(node.args[0] if node.args else None) + if holder is not None: + candidates.append((node, holder)) + return candidates + + +def bind_weights_for_copy_engine(graph: fx.GraphModule, example_inputs: Sequence[Any] | None, min_shard_bytes: int = 0) -> int: + """Move eligible weights into symmetric memory and tag the gathers that got served. + + Must run before bucketing. ``example_inputs`` are the live tensors Dynamo + captured; without them nothing binds. Returns how many gathers are now + copy-engine backed. + """ + from magi_compiler.symm_mem.bind import bind_graph_weights + + placeholders = graph.graph.find_nodes(op="placeholder") + served = bind_graph_weights( + graph=graph, + candidates=copy_engine_weight_candidates(graph), + # Built here so the NCCL path never pays for it. + placeholder_examples=dict(zip((n.name for n in placeholders), example_inputs or ())), + min_shard_bytes=min_shard_bytes, + ) + for node in served: + mark_ce_bound(node) + return len(served) + + +def rewrite_weight_ag_to_copy_engine(graph: fx.GraphModule) -> int: + """Retarget CE_BOUND weight gathers onto copy-engine ops. Returns how many.""" + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED + + rewritten = 0 + skipped = 0 + for node in graph.graph.nodes: + if node.op != "call_function" or node.target not in (_ALL_GATHER, _ALL_GATHER_COALESCED): + continue + if not is_weight_ag(node): + continue + if not is_ce_bound(node): + skipped += 1 + continue + node.target = CE_ALL_GATHER if node.target is _ALL_GATHER else CE_ALL_GATHER_COALESCED + rewritten += 1 + + if rewritten: + graph.graph.lint() + graph.recompile() + magi_logger.info( + "FSDP copy-engine rewrite: %d weight all-gather(s) retargeted, %d left on NCCL (weight not bound)", rewritten, skipped + ) + return rewritten diff --git a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py index 34123dc..4524f0c 100644 --- a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py +++ b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py @@ -14,15 +14,26 @@ from __future__ import annotations +from typing import Any, Sequence + import torch.fx as fx from magi_compiler.utils import magi_logger from .bucket_all_gather import bucket_weight_all_gather_coalesced +from .copy_engine import bind_weights_for_copy_engine, rewrite_weight_ag_to_copy_engine +from .node_meta import is_ce_bound 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: +def lower_and_bucket_full_graph( + graph: fx.GraphModule, + bucket_mode: str, + bucket_size_bytes: int = 0, + transport: str = "nccl", + example_inputs: Sequence[Any] | None = None, + min_shard_bytes: int = 0, +) -> int: """Lower SimpleFSDP weight redistribute -> explicit collectives, then optionally bucket them across the WHOLE graph (no subgraph partitioning). @@ -36,19 +47,31 @@ def lower_and_bucket_full_graph(graph: fx.GraphModule, bucket_mode: str, bucket_ the byte cap in program order (see ``bucket_weight_all_gather_coalesced``). 0 = no cap (one bucket per (group, dtype) run). + ``transport="copy_engine"`` wraps the bucketing in the two steps ``copy_engine`` + owns: binding right after lowering, the retarget at the very end. Bucketing + then keys off what binding served, so bound and unbound gathers are split into + separate buckets rather than the unbound ones being dropped from bucketing -- + losing the copy engine must not also lose coalescing. ``example_inputs`` and + ``min_shard_bytes`` are read only on this path. + 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 transport == "copy_engine": + bind_weights_for_copy_engine(graph, example_inputs, min_shard_bytes) + bucket_mode = (bucket_mode or "none").lower() + n = 0 if bucket_mode == "coalesced": - n = bucket_weight_all_gather_coalesced(graph, bucket_size_bytes=bucket_size_bytes) - else: + split_by = is_ce_bound if transport == "copy_engine" else None + n = bucket_weight_all_gather_coalesced(graph, bucket_size_bytes=bucket_size_bytes, split_by=split_by) + magi_logger.info("Whole-graph FSDP bucketing (%s): created %d buckets", bucket_mode, n) + elif bucket_mode not in ("none", ""): 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) + if transport == "copy_engine": + rewrite_weight_ag_to_copy_engine(graph) + return n diff --git a/magi_compiler/passes/fsdp_overlap/node_meta.py b/magi_compiler/passes/fsdp_overlap/node_meta.py new file mode 100644 index 0000000..ae5b54e --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/node_meta.py @@ -0,0 +1,57 @@ +# 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. + +"""Node-meta keys the FSDP overlap passes share. + +``node.meta`` is the only channel that survives bucketing, which rebuilds nodes. +Keys live here so a typo cannot silently drop a weight back to NCCL. +""" + +from __future__ import annotations + +import torch.fx as fx + +# All-gather that gathers a SimpleFSDP weight. +WEIGHT_AG = "magi_fsdp_weight_ag" + +# Weight gather whose Shard(0) does not divide across the mesh. +UNEVEN_SHARD = "magi_fsdp_uneven_shard" + +# Weight gather whose shard now lives in symmetric memory. Set by binding only. +CE_BOUND = "magi_ce_bound" + + +def is_weight_ag(node: fx.Node) -> bool: + return bool(node.meta.get(WEIGHT_AG)) + + +def is_uneven_shard(node: fx.Node) -> bool: + return bool(node.meta.get(UNEVEN_SHARD)) + + +def is_ce_bound(node: fx.Node) -> bool: + return bool(node.meta.get(CE_BOUND)) + + +def mark_weight_ag(node: fx.Node, *, uneven: bool) -> None: + """Tag a newly built all-gather as a weight gather. + + ``uneven`` is keyword-only: omitting it would default to the unsafe answer. + """ + node.meta[WEIGHT_AG] = True + node.meta[UNEVEN_SHARD] = uneven + + +def mark_ce_bound(node: fx.Node) -> None: + node.meta[CE_BOUND] = True diff --git a/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py index 8a5e262..5958d96 100644 --- a/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py +++ b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py @@ -14,6 +14,7 @@ from __future__ import annotations +import inspect import math import torch @@ -21,6 +22,8 @@ from magi_compiler.utils import magi_logger +from .node_meta import mark_weight_ag + # 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 @@ -33,6 +36,32 @@ def _is_prim(node: fx.Node, name: str) -> bool: return node.op == "call_function" and getattr(node.target, "__name__", None) == name +def _is_redistribute(node: fx.Node) -> bool: + """A SimpleFSDP ``param.redistribute(...)``, in either shape Dynamo emits for it. + + Through torch 2.9 the non-proxyable arguments were captured in an on-the-fly + function named ``prim_redistribute``; from 2.12 Dynamo traces the method call + itself and the placements and dtypes ride along in ``kwargs``. Matching only + the older shape leaves the newer graph silently un-lowered: no marked gather, + so the copy-engine rewrite and the bucketing find nothing to work with. + """ + return _is_prim(node, "prim_redistribute") or (node.op == "call_method" and node.target == "redistribute") + + +def _is_to_local(node: fx.Node) -> bool: + return _is_prim(node, "prim_to_local") or (node.op == "call_method" and node.target == "to_local") + + +def _forward_dtype(node: fx.Node): + """The dtype SimpleFSDP casts the shard to before gathering it, or None.""" + if node.op == "call_method": + return node.kwargs.get("forward_dtype") + try: + return inspect.getclosurevars(node.target).nonlocals.get("kwargs_as_value", {}).get("forward_dtype") + except Exception: # noqa: BLE001 - a prim without the expected closure is simply unannotated + return None + + 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 @@ -49,10 +78,10 @@ def _dtensor_meta(node: fx.Node): 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). + """Rewrite SimpleFSDP weight redistribute + 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 @@ -74,13 +103,13 @@ def lower_prim_redistribute_to_collectives(graph: fx.GraphModule) -> int: skipped = 0 for node in list(graph.graph.nodes): - if not _is_prim(node, "prim_redistribute"): + if not _is_redistribute(node): 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) + # to_local consumer (the node whose output the rest of the graph uses). + to_local = next((u for u in node.users if _is_to_local(u)), None) if to_local is None: skipped += 1 continue @@ -115,13 +144,7 @@ def lower_prim_redistribute_to_collectives(graph: fx.GraphModule) -> int: 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 + fwd_dtype = _forward_dtype(node) with graph.graph.inserting_before(node): # The weight placeholder is still a Shard(0) DTensor; the functional @@ -152,9 +175,7 @@ def lower_prim_redistribute_to_collectives(graph: fx.GraphModule) -> int: 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 + mark_weight_ag(ag, uneven=world * chunk != F) wait = graph.graph.call_function(_WAIT, (ag,)) wait.meta["example_value"] = local.new_empty((world * chunk, *local.shape[1:]), dtype=cur_dtype) @@ -176,8 +197,6 @@ def lower_prim_redistribute_to_collectives(graph: fx.GraphModule) -> int: graph.graph.lint() graph.recompile() magi_logger.info( - "FSDP redistribute lowering: lowered %d weight prim_redistribute -> explicit collectives (skipped %d)", - lowered, - skipped, + "FSDP redistribute lowering: lowered %d weight 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 index 3c80ab1..1c2624b 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -28,9 +28,13 @@ 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. +transfer stream) and targets only decrease. Compute is claimed a whole node at a +time, so the gather that stops in front of a long kernel leaves most of it unused; +that remainder carries to the next gather instead of being discarded, which is +what keeps a 95us collective from spending a 33ms attention. 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). @@ -51,7 +55,21 @@ _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) + + +def _ce_ag_ops(): + """Copy-engine gather ops, imported lazily so this pass stays importable + without a CUDA build.""" + try: + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED + + return tuple(op for op in (CE_ALL_GATHER, CE_ALL_GATHER_COALESCED) if op is not None) + except Exception: # noqa: BLE001 + return () + + +_CE_AG_OPS = _ce_ag_ops() +_WEIGHT_AG_OPS = tuple(op for op in (_AG, _AG_COALESCED, *_CE_AG_OPS) if op is not None) # Default extra headroom (ns) added to each collective's runtime when sizing the # compute window, absorbing estimator error + kernel-launch latency so the wait @@ -59,19 +77,43 @@ _DEFAULT_WINDOW_MARGIN_NS = 5_000.0 +def _is_ce_ag_ir(node) -> bool: + """ + ``magi::ce_all_gather`` lowers to an ordinary FallbackKernel, so + Inductor's ``is_collective`` does not recognize it. + """ + return getattr(node, "op_overload", None) in _CE_AG_OPS + + +def _is_ce_ag_coalesced(node) -> bool: + try: + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED + except Exception: # noqa: BLE001 + return False + return CE_ALL_GATHER_COALESCED is not None and getattr(node, "op_overload", None) is CE_ALL_GATHER_COALESCED + + +def _is_gather_ir(node) -> bool: + return node is not None and (is_collective(node) or _is_ce_ag_ir(node)) + + 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): + if _is_gather_ir(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): + if _is_gather_ir(cn): return cn return None +def _issues_transfer(snode: BaseSchedulerNode) -> bool: + return contains_collective(snode) or _leaf_collective_node(snode) is not 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 @@ -107,12 +149,13 @@ def _collective_kind_key(snode: BaseSchedulerNode) -> tuple: def _collective_skeleton(order: list[BaseSchedulerNode]) -> tuple[list[int], list[tuple]]: """The graph's collective skeleton: indices (ascending) and rank-comparable - kinds of every snode that ISSUES NCCL -- functional collectives plus custom ops - with an internal collective . This sequence is what must stay rank-identical; - the compute between two consecutive entries is rank-private.""" + kinds of every snode that issues a transfer -- functional NCCL collectives, + custom ops with an internal collective, and copy-engine / symmetric-memory + gathers. This sequence is what must stay rank-identical; the compute + between two consecutive entries is rank-private.""" from magi_compiler.profiling.runtime_estimator import snode_issues_collective - idx = [i for i, s in enumerate(order) if snode_issues_collective(s)] + idx = [i for i, s in enumerate(order) if snode_issues_collective(s) or _issues_transfer(s)] return idx, [_collective_kind_key(order[i]) for i in idx] @@ -212,7 +255,7 @@ def _cost(self, snode: BaseSchedulerNode) -> float: @staticmethod def _is_compute(snode: BaseSchedulerNode) -> bool: - return not contains_collective(snode) and not contains_wait(snode) + return not _issues_transfer(snode) and not contains_wait(snode) # -- main ------------------------------------------------------------- def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: @@ -247,7 +290,7 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: 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): + if self._is_compute(s) or _issues_transfer(s): self._cost(s) n_changed = self._cost_fn.warm_and_sync() self._cost_cache = {} # re-read synced costs @@ -279,25 +322,25 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: targets: dict = {} # launch -> target index (in original order space) compute_idx = len(order) # scan compute strictly below this + carry = 0.0 # runtime the previous gather left unspent in its boundary node 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_overlap_window_scale + self.comm_overlap_window_margin_ns - acc = 0.0 + acc = carry_in = carry t = compute_idx - while t > lower: + while acc < need and 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) + carry = max(0.0, acc - need) 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": @@ -306,18 +349,24 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: # 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 + # carry_in = capacity inherited from the later gather's boundary node + # (target==cur with a large carry_in means it was already + # covered and did not have to move at all) + # acc_upstream = carry_in plus the compute found in [target, cur] # 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) magi_logger.debug( - "FSDP overlap placement: launch cur=%d -> target=%d fc=%d lower=%d | " - "comm=%.1fus acc_upstream=%.1fus need=%.1fus %s", + "FSDP overlap placement: launch %s(%s) cur=%d -> target=%d fc=%d lower=%d | " + "comm=%.1fus carry_in=%.1fus acc_upstream=%.1fus need=%.1fus %s", + launch.get_name(), + getattr(_leaf_collective_node(launch), "op_overload", "?"), cur, target, fc_idx, lower, comm_runtime / 1e3, + carry_in / 1e3, acc / 1e3, need / 1e3, "hidden" if acc >= need else "COMPUTE-LIMITED", @@ -476,7 +525,7 @@ def slot_of(idx: int) -> int: target, group = targets[launch] slot_lo = max(lowers[launch], skel_idx[q - 1] + 1 if q > 0 else 0) slot_hi = skel_idx[q] if q < len(skel_idx) else index_of[launch] - new_target = min(max(target, slot_lo), slot_hi) + new_target = min(max(target, slot_lo), max(slot_hi, slot_lo)) targets[launch] = (new_target, group) magi_logger.debug( "FSDP overlap slot consensus: launch cur=%d slot=%d/%d (mine=%s) target %d -> %d [%d, %d]", @@ -500,27 +549,48 @@ def _launch_group(self, launch, order, buf_to_snode, users) -> list[BaseSchedule """ 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()) + produced = set(launch.get_buffer_names()) + if node is not None and (getattr(node, "op_overload", None) is _AG_COALESCED or _is_ce_ag_coalesced(node)): 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) + if _is_ce_ag_coalesced(node): + for s in order: + if s is launch or s in group or contains_wait(s) or not self._is_transparent(s): + continue + deps = [d for d in s.unmet_dependencies if not _is_fake_dep(d)] + if deps and all(d.name in produced for d in deps): + group.append(s) + elif _is_ce_ag_ir(node): + for s in order: + if s is launch or contains_wait(s) or not self._is_transparent(s): + continue + deps = [d for d in s.unmet_dependencies if not _is_fake_dep(d)] + if deps and all(d.name in produced for d in deps): + 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 + """The waits guarding this launch, reached through any alias layer. + + Searching only the launch's direct readers was enough while every gather + was an Inductor collective; a custom-op gather puts an alias snode between + the launch and its wait, and missing the wait silently drops the gather + from the placement plan altogether. + """ + stack = [b for s in group for b in s.get_buffer_names()] + waits: list[BaseSchedulerNode] = [] + seen: set = set() + while stack: + for u in users.get(stack.pop(), ()): if u in seen: continue seen.add(u) if contains_wait(u): waits.append(u) + elif self._is_transparent(u): + stack.extend(u.get_buffer_names()) return waits def _first_consumer_index(self, launch, group, order, users) -> int | None: diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index db3246a..a83f073 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -45,7 +45,7 @@ """ import dataclasses -from typing import Any +from typing import Any, Optional import torch from torch._inductor.runtime.benchmarking import benchmarker @@ -90,21 +90,81 @@ class ProfileEntry: reuse_count: int = 0 # how many later snodes reused this entry +def _fx_node_of(node) -> Optional[torch.fx.Node]: + """The fx node an IR node was lowered from. + + ``origin_node`` is left unset by several ExternKernel subclasses -- notably + ``UserDefinedTritonKernel``, whose lowering builds it with positional args + only. ``ExternKernel.__init__`` always records the node being lowered as + ``fx_node``, so fall back to that; without it every user-defined Triton + kernel is unreplayable and silently costs 0. + """ + if node is None: + return None + origin = node.get_origin_node() if hasattr(node, "get_origin_node") else None + if origin is not None: + return origin + fx_node = getattr(node, "fx_node", None) + return fx_node if isinstance(fx_node, torch.fx.Node) else None + + +def _iter_tensor_metas(value): + """Every FakeTensor meta reachable from an fx arg, recursing into containers. + + The user-defined Triton HOP passes its tensors inside a nested ``kwargs`` + dict, so a flat scan over ``args``/``kwargs`` sees no shapes at all. + """ + if isinstance(value, torch.fx.Node): + ev = value.meta.get("val") + if isinstance(ev, torch.Tensor): + yield ev + elif isinstance(value, (list, tuple)): + for item in value: + yield from _iter_tensor_metas(item) + elif isinstance(value, dict): + for item in value.values(): + yield from _iter_tensor_metas(item) + + +def _fx_target_name(fx_node: Optional[torch.fx.Node], node) -> str: + """Op identity string. All user-defined Triton kernels share one HOP target, + so qualify it with the kernel's own name -- otherwise hundreds of distinct + kernels collapse onto a single cache entry.""" + if fx_node is None: + return type(node).__name__ if node is not None else "?" + target = str(fx_node.target) + kernel_name = _triton_kernel_name(fx_node) + return f"{target}:{kernel_name}" if kernel_name else target + + +def _triton_kernel_name(fx_node: torch.fx.Node) -> Optional[str]: + """Name of the user-defined Triton kernel this node launches, if any.""" + kernel_idx = (fx_node.kwargs or {}).get("kernel_idx") + if not isinstance(kernel_idx, int): + return None + try: + from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table + + kernel = kernel_side_table.get_kernel(kernel_idx) + except Exception: # noqa: BLE001 + return str(kernel_idx) + fn = getattr(kernel, "fn", kernel) + return getattr(fn, "__name__", None) or str(kernel_idx) + + 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 "?")) + origin = _fx_node_of(node) + target = _fx_target_name(origin, node) 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 + for ev in _iter_tensor_metas((*origin.args, *getattr(origin, "kwargs", {}).values())): + 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 @@ -126,14 +186,12 @@ def _structural_key(snode: BaseSchedulerNode) -> tuple | None: 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__)) + origin = _fx_node_of(node) + target = _fx_target_name(origin, node) 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))) + for ev in _iter_tensor_metas((*origin.args, *origin.kwargs.values())): + shapes.append((tuple(_static(ev.shape)), str(ev.dtype))) parts.append((target, tuple(shapes))) return tuple(parts) @@ -216,19 +274,15 @@ def _realize_arg(v): 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. +def _extern_replay_fn(snode: ExternKernelSchedulerNode): + """A callable that runs this extern's aten op on rebuilt inputs, or None. Replay inputs: generic ``_realize_arg``, then an optional same-signature - hook (``materialize_inputs``) that rebuilds value-consistent metadata.""" - fx_node = snode.node.get_origin_node() + hook (``materialize_inputs``) that rebuilds value-consistent metadata. + """ + fx_node = _fx_node_of(snode.node) if fx_node is None: - return 0.0 + return None target = fx_node.target args = tuple(_realize_arg(a) for a in fx_node.args) @@ -248,22 +302,45 @@ 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): + return fn + + +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.""" + fn = _extern_replay_fn(snode) + if fn is None: + # Never report 0: an unreplayable op is unknown, not free, and a silent 0 + # makes the overlap pass treat a real kernel as a gap it can hoist across. + raise RuntimeError(f"{snode.get_name()}: no replayable fx node, cannot measure") + if fixed_iters: + return _time_fixed(fn) + fn() # warmup / correctness + return benchmarker.benchmark_gpu(fn) * 1e6 # ms -> ns + + +def _time_fixed(fn, warmup: int = 3, iters: int = 10) -> float: + """CUDA-event timing over a FIXED iteration count, in nanoseconds. + + Fixed, not adaptive: anything that issues a collective must issue the same + number of them on every rank, or the NCCL counts diverge and the ranks + deadlock inside what is supposed to be a measurement. + """ + 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): + for _ in range(iters): fn() end.record() torch.cuda.synchronize() - return (start.elapsed_time(end) / _ITERS) * 1e6 # ms/iter -> ns + return (start.elapsed_time(end) / iters) * 1e6 # ms/iter -> ns def _op_name(target) -> str: @@ -326,6 +403,103 @@ def _collective_spec(node): return op, group_name, group_size, specs +def _ce_ag_ops(): + """Copy-engine gather ops, or empty when the runtime is unavailable.""" + try: + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED + + return tuple(op for op in (CE_ALL_GATHER, CE_ALL_GATHER_COALESCED) if op is not None) + except Exception: # noqa: BLE001 + return () + + +def _leaf_ce_ag(snode: BaseSchedulerNode): + """The copy-engine gather IR node inside ``snode``, or None. + + It is an ordinary FallbackKernel, not a ``_CollectiveKernel``, so none of + Inductor's collective predicates see it. + """ + ops = _ce_ag_ops() + if not ops: + return None + for n in (getattr(snode, "node", None), *(getattr(c, "node", None) for c in getattr(snode, "snodes", []) or [])): + if n is not None and getattr(n, "op_overload", None) in ops: + return n + return None + + +def _is_ce_ag_coalesced_ir(node) -> bool: + try: + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED + except Exception: # noqa: BLE001 + return False + return CE_ALL_GATHER_COALESCED is not None and getattr(node, "op_overload", None) is CE_ALL_GATHER_COALESCED + + +def _ce_ag_spec(node): + """(shapes, dtype, group_size, group_name). Use ``constant_args``; ``get_origin_node()`` is unset and would cost 0.""" + args = getattr(node, "constant_args", None) + if not args or len(args) < 2: + return None + group_size, group_name = args[-2:] + ins = list(node.inputs) + if not ins: + return None + shapes = tuple(tuple(_concrete_size(s) for s in inp.layout.size) for inp in ins) + return shapes, ins[0].layout.dtype, int(group_size), str(group_name) + + +def _ce_ag_launch_wait(snode: BaseSchedulerNode): + """``(launch, wait)`` replaying a copy-engine gather, or None. + + Split in two rather than one fused closure so the cost model can time + ``wait(launch())`` as a unit. + """ + from magi_compiler.symm_mem import find_shard_by_layout + + node = _leaf_ce_ag(snode) + spec = _ce_ag_spec(node) if node is not None else None + if spec is None: + return None + shapes, dtype, group_size, group_name = spec + shards = [find_shard_by_layout(shape, dtype) for shape in shapes] + if any(s is None for s in shards): + magi_logger.warning( + "No registered symmetric shard with layout %s/%s; the copy-engine gather keeps its " + "analytical cost and its overlap window may be mis-sized", + shapes, + dtype, + ) + return None + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED + + if _is_ce_ag_coalesced_ir(node): + op = CE_ALL_GATHER_COALESCED + return (lambda: op(shards, group_size, group_name)), lambda outs: [_WAIT(o) for o in outs] + return (lambda: CE_ALL_GATHER(shards[0], group_size, group_name)), _WAIT + + +def _measure_ce_ag(snode: BaseSchedulerNode) -> float: + """Time ``wait(launch())``. Launch-only is ~3us CPU issue; copies run on a side stream.""" + pair = _ce_ag_launch_wait(snode) + if pair is None: + return 0.0 + launch, wait = pair + return _time_fixed(lambda: wait(launch())) + + +def _ce_ag_label(snode: BaseSchedulerNode) -> str: + node = _leaf_ce_ag(snode) + spec = _ce_ag_spec(node) if node is not None else None + if spec is None: + return _snode_label(snode) + shapes, _dtype, group_size, _gn = spec + shape0 = "x".join(str(x) for x in shapes[0]) + if len(shapes) > 1: + return f"ce_all_gather_coalesced(ws={group_size},n={len(shapes)},{shape0})" + return f"ce_all_gather(ws={group_size},{shape0})" + + def _collective_label(snode: BaseSchedulerNode) -> str: """Readable identity of a collective: op name, world size, #inputs + first shape.""" node = _leaf_collective(snode) @@ -337,43 +511,26 @@ def _collective_label(snode: BaseSchedulerNode) -> str: 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.""" +def _nccl_launch_wait(snode: BaseSchedulerNode): + """``(launch, wait)`` replaying a functional all-gather, or None.""" node = _leaf_collective(snode) - if node is None: - return 0.0 - spec = _collective_spec(node) + spec = _collective_spec(node) if node is not None else None if spec is None: - return 0.0 + return None 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: + return (lambda: _AG_COALESCED(ins, group_size, group_name), lambda outs: [_WAIT(o) for o in outs]) + return (lambda: _AG(ins[0], group_size, group_name)), _WAIT - 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 +def _measure_collective_op(snode: BaseSchedulerNode) -> float: + """Replay the functional all-gather (+wait) on real tensors and time it.""" + pair = _nccl_launch_wait(snode) + if pair is None: + return 0.0 + launch, wait = pair + return _time_fixed(lambda: wait(launch())) class ProfilingRuntimeEstimator: @@ -465,8 +622,10 @@ def warm_and_sync(self) -> int: snode = self._key_snode.get(k) dist.barrier(group=group) # snode is non-None on every rank by measurable_reprs construction. - local_ns[k] = self._measure_one(snode) - measured_here.add(k) + ns, ok = self._measure_one(snode) + local_ns[k] = ns + if ok: + measured_here.add(k) dist.barrier(group=group) gathered: list = [None] * world @@ -504,22 +663,25 @@ def warm_and_sync(self) -> int: self._key_snode.clear() # drop snode refs (unpicklable) once sync is done return n - def _measure_one(self, snode: BaseSchedulerNode) -> float: + def _measure_one(self, snode: BaseSchedulerNode) -> tuple[float, bool]: """Lockstep-safe single measurement (fixed iters for anything containing a - collective); never raises -- falls back to the analytical estimate.""" + collective). Never raises; returns ``(ns, measured)`` so the caller can + tell a real timing from the analytical fallback.""" try: + if _leaf_ce_ag(snode) is not None: + return _measure_ce_ag(snode), True if contains_collective(snode): - return _measure_collective_op(snode) + return _measure_collective_op(snode), True 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) + return ns, True + return self._measure(snode), True 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) + magi_logger.warning("warm/sync measure fell back to analytical for %s: %s", snode.get_name(), exc) + return _safe_analytical(snode), False def summary(self) -> str: """One line per distinct op + a machine-parseable ``ESTLINE`` tag @@ -548,6 +710,29 @@ def __call__(self, snode: BaseSchedulerNode) -> float: if _is_multi_output_unpack(snode): return 0.0 + if _leaf_ce_ag(snode) is not None: + node = _leaf_ce_ag(snode) + spec = _ce_ag_spec(node) + if spec is None: + return _safe_analytical(snode) + shapes, dtype, group_size, _gn = spec + ckey = ("ce_ag", group_size, shapes, str(dtype)) + 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) + self._table[ckey] = ProfileEntry(ns=ns, kind="ce_ag", label=_ce_ag_label(snode), measured=False) + if self._sync_across_ranks: + self._key_snode[ckey] = snode + else: + ns = _measure_ce_ag(snode) + self._table[ckey].ns = ns + self._table[ckey].measured = True + self.n_measured += 1 + return ns + if contains_collective(snode): cnode = _leaf_collective(snode) spec = _collective_spec(cnode) if cnode is not None else None diff --git a/magi_compiler/symm_mem/__init__.py b/magi_compiler/symm_mem/__init__.py new file mode 100644 index 0000000..3874358 --- /dev/null +++ b/magi_compiler/symm_mem/__init__.py @@ -0,0 +1,44 @@ +# 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. + +"""Symmetric-memory weight storage and the copy-engine all-gather built on it. + +``all_gather`` is deliberately not re-exported here: importing it defines the +``magi::ce_all_gather`` ops, and callers use the import itself as the probe for +whether the copy-engine transport is available. Import that module by path. +""" + +from .bind import bind_graph_weights, bind_parameters, group_name_of +from .symm_buffer import ( + ShardEntry, + SymmBuffer, + alloc_shard, + find_shard_by_layout, + lookup_shard, + registered_buffers, + reset_registry, +) + +__all__ = [ + "ShardEntry", + "SymmBuffer", + "alloc_shard", + "bind_graph_weights", + "bind_parameters", + "find_shard_by_layout", + "group_name_of", + "lookup_shard", + "registered_buffers", + "reset_registry", +] diff --git a/magi_compiler/symm_mem/all_gather.py b/magi_compiler/symm_mem/all_gather.py new file mode 100644 index 0000000..4fec639 --- /dev/null +++ b/magi_compiler/symm_mem/all_gather.py @@ -0,0 +1,233 @@ +# 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 ctypes +from functools import lru_cache + +import torch +import torch._C._distributed_c10d as _c10d + +from magi_compiler.utils import magi_logger + +from .symm_buffer import lookup_shard + +_LIB = torch.library.Library("magi", "FRAGMENT") +# Signatures mirror ``_c10d_functional::all_gather_into_tensor`` / ``_coalesced`` so +# the rewrite pass can retarget a node without rebuilding its args. That is also the +# only reason ``group_name`` is here: the copy engine reads peers from the SymmBuffer. +_SCHEMA = "ce_all_gather(Tensor local, int group_size, str group_name) -> Tensor" +_SCHEMA_COALESCED = "ce_all_gather_coalesced(Tensor[] shards, int group_size, str group_name) -> Tensor[]" + +# One gather: where it lands, the local shard, and every rank's view of that shard. +_Gather = tuple[torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...]] +# ``cudaMemcpyBatchAsync`` arguments frozen for one submission: dsts, srcs, sizes, count. +_Plan = tuple[ctypes.Array, ctypes.Array, ctypes.Array, int] + + +class _cudaMemLocation(ctypes.Structure): + _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)] + + +class _cudaMemcpyAttributes(ctypes.Structure): + _fields_ = [ + ("srcAccessOrder", ctypes.c_int), + ("srcLocHint", _cudaMemLocation), + ("dstLocHint", _cudaMemLocation), + ("flags", ctypes.c_uint), + ] + + +class BatchMemcpy: + """``cudaMemcpyBatchAsync``: one submission for a whole layer's copies. + + Runtime signature is 8 args with no ``failIdx`` (that's driver-only + ``cuMemcpyBatchAsync``), and it rejects the legacy null stream. + """ + + _SRC_ACCESS_ORDER_STREAM = 1 + + def __init__(self) -> None: + lib = ctypes.CDLL(None) + lib.cudaGetErrorString.restype = ctypes.c_char_p + self._strerror = lib.cudaGetErrorString + self._fn = lib.cudaMemcpyBatchAsync + self._fn.restype = ctypes.c_int + self._fn.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_size_t, + ctypes.POINTER(_cudaMemcpyAttributes), + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_size_t, + ctypes.c_void_p, + ] + self._attrs = (_cudaMemcpyAttributes * 1)() + self._attrs[0].srcAccessOrder = self._SRC_ACCESS_ORDER_STREAM + self._attr_idxs = (ctypes.c_size_t * 1)(0) + + @staticmethod + def plan(triples: list[tuple[int, int, int]]) -> _Plan: + """Freeze ``(dst_ptr, src_ptr, nbytes)`` for one submission. + + Raw pointers, not tensor views: a per-peer view is ~1.5us of CPU. + Valid only while both sides stay alive; rebuilt per call. + """ + n = len(triples) + return ( + (ctypes.c_void_p * n)(*[d for d, _s, _b in triples]), + (ctypes.c_void_p * n)(*[s for _d, s, _b in triples]), + (ctypes.c_size_t * n)(*[b for _d, _s, b in triples]), + n, + ) + + def run(self, plan: _Plan, stream: torch.cuda.Stream) -> None: + dsts, srcs, sizes, n = plan + rc = self._fn(dsts, srcs, sizes, n, self._attrs, self._attr_idxs, 1, ctypes.c_void_p(stream.cuda_stream)) + if rc: + raise RuntimeError(f"cudaMemcpyBatchAsync failed: {self._strerror(rc).decode()}") + + +@lru_cache(maxsize=1) +def _batcher() -> BatchMemcpy | None: + try: + return BatchMemcpy() + except (AttributeError, OSError) as exc: + magi_logger.warning("cudaMemcpyBatchAsync unavailable (%s); falling back to per-copy submission", exc) + return None + + +class _EventWork(_c10d.Work): + """c10d Work whose ``wait()`` is a stream wait on the copy-engine event.""" + + def __init__(self, event: torch.cuda.Event) -> None: + super().__init__() + self._event = event + + def wait(self, timeout=None) -> bool: # noqa: ARG002 - c10d's signature + torch.cuda.current_stream().wait_event(self._event) + return True + + +@lru_cache(maxsize=1) +def _copy_stream() -> torch.cuda.Stream: + """The one stream every copy-engine gather is submitted on.""" + return torch.cuda.Stream() + + +def _shard_peers(local: torch.Tensor, group_size: int) -> tuple[torch.Tensor, ...]: + """The registered peer views of a local shard, validated.""" + entry = lookup_shard(local.data_ptr()) + if entry is None: + raise RuntimeError( + "magi::ce_all_gather got a tensor that is not a registered symmetric-memory shard. " + "Only weights materialized through a SymmBuffer can be gathered by the copy engine; " + "the rewrite pass should have left this gather on NCCL." + ) + peers = entry.peer_views + if len(peers) != group_size: + raise RuntimeError(f"shard has {len(peers)} peers but the gather asks for group_size={group_size}") + return peers + + +def _copy_triples(gathers: list[_Gather]) -> list[tuple[int, int, int]]: + """Flatten to one ``(dst_ptr, src_ptr, nbytes)`` per (member, peer) pair.""" + triples: list[tuple[int, int, int]] = [] + for out, local, peers in gathers: + nbytes = local.numel() * local.element_size() # dest contiguous; rank r at r*nbytes + base = out.data_ptr() + triples.extend((base + r * nbytes, p.data_ptr(), nbytes) for r, p in enumerate(peers)) + return triples + + +def _copy_per_peer(gathers: list[_Gather]) -> None: + """Fallback for runtimes without ``cudaMemcpyBatchAsync``: one ``copy_`` per peer.""" + for out, local, peers in gathers: + rows = local.shape[0] + for r, p in enumerate(peers): + out[r * rows : (r + 1) * rows].copy_(p, non_blocking=True) + + +def _issue_gathers(gathers: list[_Gather]) -> torch.cuda.Event: + """Submit every gather as one batch; return the event that completes them all. + + Stream sync, submission and event are paid once for the whole call, not + once per member -- that fixed CPU cost otherwise inflates the overlap window. + """ + batcher = _batcher() + stream = _copy_stream() + stream.wait_stream(torch.cuda.current_stream()) # copies after compute-stream writes to the shards + with torch.cuda.stream(stream): + if batcher is not None: + batcher.run(BatchMemcpy.plan(_copy_triples(gathers)), stream) + else: + _copy_per_peer(gathers) + event = torch.cuda.Event() + event.record(stream) + return event + + +def _gather_dest(local: torch.Tensor, group_size: int) -> torch.Tensor: + """Destination for gathering ``local``: rank r's shard lands at row ``r * rows``.""" + return local.new_empty((local.shape[0] * group_size, *local.shape[1:])) + + +def _ce_all_gather(local: torch.Tensor, group_size: int, group_name: str) -> torch.Tensor: + peers = _shard_peers(local, group_size) + out = _gather_dest(local, group_size) + event = _issue_gathers([(out, local, peers)]) + _c10d._register_work(out, _EventWork(event)) + return out + + +def _ce_all_gather_meta(local: torch.Tensor, group_size: int, group_name: str) -> torch.Tensor: + return _gather_dest(local, group_size) + + +def _ce_all_gather_coalesced(shards: list[torch.Tensor], group_size: int, group_name: str) -> list[torch.Tensor]: + """One stream sync, one batch, one event for the whole bucket -- not one per member.""" + gathers: list[_Gather] = [] + for local in shards: + peers = _shard_peers(local, group_size) # validates before allocating + gathers.append((_gather_dest(local, group_size), local, peers)) + event = _issue_gathers(gathers) + outs = [out for out, _local, _peers in gathers] + # Registry takes ownership of each Work; members share the event, not the wrapper. + for out in outs: + _c10d._register_work(out, _EventWork(event)) + return outs + + +def _ce_all_gather_coalesced_meta(shards: list[torch.Tensor], group_size: int, group_name: str) -> list[torch.Tensor]: + return [_gather_dest(local, group_size) for local in shards] + + +def _register() -> None: + _LIB.define(_SCHEMA) + _LIB.impl("ce_all_gather", _ce_all_gather, "CUDA") + _LIB.impl("ce_all_gather", _ce_all_gather_meta, "Meta") + + _LIB.define(_SCHEMA_COALESCED) + _LIB.impl("ce_all_gather_coalesced", _ce_all_gather_coalesced, "CUDA") + _LIB.impl("ce_all_gather_coalesced", _ce_all_gather_coalesced_meta, "Meta") + + +_register() + +# Importing this module is what makes the ops exist, so these are always bound -- +# callers guard the import, not the value. +CE_ALL_GATHER = torch.ops.magi.ce_all_gather.default +CE_ALL_GATHER_COALESCED = torch.ops.magi.ce_all_gather_coalesced.default diff --git a/magi_compiler/symm_mem/bind.py b/magi_compiler/symm_mem/bind.py new file mode 100644 index 0000000..ab6f357 --- /dev/null +++ b/magi_compiler/symm_mem/bind.py @@ -0,0 +1,370 @@ +# 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. + +"""Move captured-graph weights into symmetric memory. + +Runs between lowering and bucketing. Downstream keys off what actually moved; +anything that cannot bind stays on NCCL. +""" + +from __future__ import annotations + +from collections import Counter, defaultdict +from dataclasses import dataclass +from typing import Any, Iterable, Mapping + +import torch +import torch.distributed as dist +import torch.fx as fx + +from magi_compiler.utils import magi_logger + +from .symm_buffer import lookup_shard, open_buffer, register_shard + +_AGREEMENT_GROUP: tuple[Any, Any] = (None, None) +"""``(the default process group it was built for, the gloo group)``.""" + + +@dataclass(frozen=True) +class BindCandidate: + """One weight ready to be moved. ``local`` is ``param._local_tensor``.""" + + local: torch.Tensor + group_name: str + gather: fx.Node | None = None + holder: fx.Node | None = None + + +def bind_graph_weights( + graph: fx.GraphModule, + candidates: Iterable[tuple[fx.Node, fx.Node]], + placeholder_examples: Mapping[str, Any], + min_shard_bytes: int = 0, +) -> set[fx.Node]: + """Move the weights behind ``candidates`` into symmetric memory. + + ``candidates`` is ``(gather, holder)``; ``placeholder_examples`` maps + placeholder name to the live tensor. Returns the gathers now copy-engine + backed -- the caller tags them. Failures are dropped, not raised. + """ + plan, skipped = _plan(graph, candidates, placeholder_examples, min_shard_bytes) + + # Before the empty check, so every rank reaches the collective. + if not _agree_across_ranks(plan): + magi_logger.warning( + "copy-engine binding: ranks disagree on which weights to bind; leaving every weight " + "gather on NCCL. A gather retargeted on only some ranks would never complete." + ) + return set() + + if not plan: + magi_logger.info("copy-engine binding: nothing to bind (%s)", _describe(skipped)) + return set() + + allocated = _move(plan) + + served = {c.gather for c in plan if c.gather is not None} + magi_logger.info( + "copy-engine binding: %d weight gather(s) bound (%d new allocation(s), %.1f MiB); %s", + len(served), + allocated, + sum(c.local.numel() * c.local.element_size() for c in plan) / 2**20, + _describe(skipped), + ) + return served + + +def bind_parameters(params: Iterable[Any], min_shard_bytes: int = 0) -> int: + """Move an explicit list of sharded parameters into symmetric memory. + + For callers with no graph (op benches, AOT compile with fake inputs). + Returns how many parameters are now bound. + """ + plan: list[BindCandidate] = [] + for p in params: + if _unbindable(p, min_shard_bytes) is not None: + continue + try: + plan.append(BindCandidate(local=p._local_tensor, group_name=group_name_of(p))) + except RuntimeError: + continue + + if not _agree_across_ranks(plan): + magi_logger.warning("copy-engine binding: ranks disagree on which parameters to bind; binding none") + return 0 + if not plan: + return 0 + + _move(plan) + return len(plan) + + +def _plan( + graph: fx.GraphModule, + candidates: Iterable[tuple[fx.Node, fx.Node]], + placeholder_examples: Mapping[str, Any], + min_shard_bytes: int, +) -> tuple[list[BindCandidate], Counter]: + """Pair each candidate gather with a live shard, dropping the ones that fail.""" + plan: list[BindCandidate] = [] + skipped: Counter = Counter() + + for gather, holder in candidates: + param = _resolve(graph, holder, placeholder_examples) + if param is None: + skipped["graph input has no live parameter behind it"] += 1 + continue + why = _unbindable(param, min_shard_bytes) + if why is not None: + skipped[why] += 1 + continue + try: + group_name = group_name_of(param) + except RuntimeError as exc: + skipped[str(exc)] += 1 + continue + plan.append(BindCandidate(local=param._local_tensor, group_name=group_name, gather=gather, holder=holder)) + + return plan, skipped + + +def _resolve(graph: fx.GraphModule, holder: fx.Node, placeholder_examples: Mapping[str, Any]) -> Any: + """The live object a weight-holding node stands for, or None.""" + if holder.op == "placeholder": + return placeholder_examples.get(holder.name) + if holder.op == "get_attr": + obj: Any = graph + for part in str(holder.target).split("."): + obj = getattr(obj, part, None) + if obj is None: + return None + return obj + return None + + +def _unbindable(param: Any, min_shard_bytes: int) -> str | None: + """Why ``param`` cannot back a copy-engine gather, or None if it can.""" + from torch._subclasses.fake_tensor import FakeTensor + from torch.distributed.tensor import DTensor, Shard + + if not isinstance(param, DTensor): + return "graph input is not a DTensor" + + local = param._local_tensor + if isinstance(local, FakeTensor) or local.is_meta: + return "graph input is a fake/meta tensor" + + placements = param.placements + if len(placements) != 1 or not isinstance(placements[0], Shard) or placements[0].dim != 0: + return f"placement {tuple(placements)} is not a single Shard(0)" + # Fixed-stride peer reads cannot express uneven rank shards. + if int(param.shape[0]) % int(param.device_mesh.size(0)): + return "Shard(0) does not divide evenly across the mesh" + + if local.device.type != "cuda": + return f"shard lives on {local.device.type}, not cuda" + if not local.is_contiguous(): + return "shard is not contiguous" + if local.numel() * local.element_size() < min_shard_bytes: + return "shard is below the size floor" + return None + + +def group_name_of(t) -> str: + """The process group a sharded parameter must rendezvous on. + + Never defaulted to WORLD: dense and expert weights may sit on different meshes. + """ + mesh = getattr(t, "device_mesh", None) + names = getattr(mesh, "_dim_group_names", None) if mesh is not None else None + if not names: + raise RuntimeError( + f"cannot resolve the process group of {type(t).__name__} (device_mesh={mesh!r}); " + "copy-engine binding needs the mesh dim the weight is sharded over" + ) + return names[0] + + +_WINDOW_BYTES = 4 << 30 +"""Cap on one symmetric window (4 GiB). + +The driver allows 128 windows per process, but a window is opened while every +shard it will absorb is still resident, so this is also what binding adds to +the peak. A shard bigger than this still gets a window to itself. +""" + + +def _windows(plan: list[BindCandidate]) -> list[list[BindCandidate]]: + """Split the plan into per-window lists, grouped by (group, dtype).""" + fresh: list[BindCandidate] = [] + seen: set[int] = set() + for c in plan: + ptr = c.local.data_ptr() + # Already symmetric: a tied weight earlier in this plan, or a previous compile. + if ptr in seen or lookup_shard(ptr) is not None: + continue + seen.add(ptr) + fresh.append(c) + + by_kind: dict[tuple[str, torch.dtype], list[BindCandidate]] = defaultdict(list) + for c in fresh: + by_kind[(c.group_name, c.local.dtype)].append(c) + + windows: list[list[BindCandidate]] = [] + for members in by_kind.values(): + current: list[BindCandidate] = [] + current_bytes = 0 + for c in members: + nbytes = c.local.numel() * c.local.element_size() + if current and current_bytes + nbytes > _WINDOW_BYTES: + windows.append(current) + current, current_bytes = [], 0 + current.append(c) + current_bytes += nbytes + if current: + windows.append(current) + return windows + + +_ALIGN_BYTES = 512 +"""Every shard starts on a multiple of this many bytes. + +512B is what the copy engine wants for peak throughput. It is a property of the +transport, not of the window, so it belongs here with the rest of the layout +decision. +""" + + +def _layout(members: list[BindCandidate]) -> tuple[list[int], int]: + """Each shard's element offset within its window, and the window's numel. + + Offsets and total come out of the same walk, so the window cannot disagree + with what is dispensed from it. Ranks agree on the layout because they agree + on the plan (``_agree_across_ranks``), which is what makes offset ``k`` the + same shard on every peer. + """ + # Offsets are element counts, and one window holds one dtype, so the byte + # alignment is a fixed stride for the whole window. + align = _ALIGN_BYTES // members[0].local.element_size() + offsets: list[int] = [] + total = 0 + for c in members: + offsets.append(total) + total += (c.local.numel() + align - 1) // align * align + return offsets, total + + +def _move(plan: list[BindCandidate]) -> int: + """Allocate, fill, repoint and publish every shard. Returns the new allocation count. + + Symmetric memory cannot reuse the caching allocator's blocks, so each window + ends with ``empty_cache``. An allocation failure is fatal: ranks have already + issued a matching prefix of rendezvous. + """ + windows = _windows(plan) + allocated = 0 + for i, members in enumerate(windows): + head = members[0] + offsets, window_numel = _layout(members) + try: + buffer = open_buffer(head.local.dtype, head.local.device, head.group_name, window_numel) + except RuntimeError: + free, total = torch.cuda.mem_get_info() + magi_logger.error( + "copy-engine binding: failed to open symmetric window %d of %d (%d shard(s), " + "%.1f MiB of %s); driver has %.1f of %.1f GiB free, torch reserves %.1f GiB of " + "which %.1f GiB is live", + i, + len(windows), + len(members), + sum(c.local.numel() * c.local.element_size() for c in members) / (1 << 20), + head.local.dtype, + free / (1 << 30), + total / (1 << 30), + torch.cuda.memory_reserved() / (1 << 30), + torch.cuda.memory_allocated() / (1 << 30), + ) + raise + + for c, offset in zip(members, offsets): + symm = buffer.get_tensor(offset, c.local.shape) + symm.copy_(c.local) + register_shard(symm, buffer, offset) + # In place: Dynamo already guarded these exact objects. + c.local.data = symm + allocated += 1 + + torch.cuda.empty_cache() + + if dist.is_available() and dist.is_initialized(): + torch.cuda.synchronize() + dist.barrier() + return allocated + + +def _agree_across_ranks(plan: list[BindCandidate]) -> bool: + """True if every rank arrived at the same plan, in the same order. + + ``group_name`` is not in the fingerprint: expert weights sit on a subgroup of + the world, so each rank correctly resolves a different group for the same weight. + """ + if not (dist.is_available() and dist.is_initialized()): + return True + + mine = [(c.holder.name if c.holder is not None else "", tuple(c.local.shape), str(c.local.dtype)) for c in plan] + group = _agreement_group() + gathered: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, mine, group=group) + + for rank, theirs in enumerate(gathered): + if theirs == mine: + continue + i = next((i for i, (a, b) in enumerate(zip(mine, theirs)) if a != b), min(len(mine), len(theirs))) + magi_logger.warning( + "copy-engine binding: this rank plans %d shard(s) and rank %d plans %d; they first differ at " + "index %d, where this rank has %s and rank %d has %s", + len(mine), + rank, + len(theirs), + i, + mine[i] if i < len(mine) else "", + rank, + theirs[i] if i < len(theirs) else "", + ) + return False + return True + + +def _agreement_group(): + """A CPU group for the plan check, rebuilt if the default process group changes. + + ``new_group`` is collective, so every rank must reach here even with an empty plan. + """ + global _AGREEMENT_GROUP + owner, group = _AGREEMENT_GROUP + if owner is dist.group.WORLD: + return group + try: + group = dist.new_group(backend="gloo") + except Exception as exc: # noqa: BLE001 + magi_logger.warning("copy-engine binding: gloo group unavailable (%s); using the default group", exc) + group = None + _AGREEMENT_GROUP = (dist.group.WORLD, group) + return group + + +def _describe(skipped: Counter) -> str: + if not skipped: + return "no candidate was skipped" + return "skipped: " + ", ".join(f"{n}x {why}" for why, n in skipped.most_common()) diff --git a/magi_compiler/symm_mem/symm_buffer.py b/magi_compiler/symm_mem/symm_buffer.py new file mode 100644 index 0000000..9bf8b40 --- /dev/null +++ b/magi_compiler/symm_mem/symm_buffer.py @@ -0,0 +1,170 @@ +# 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. + +"""Symmetric-memory windows for copy-engine weight all-gather. + +A window is a plain addressable region: open one, then read any slot in it by +``(offset, shape)``. Deciding what goes where -- and keeping that decision +identical on every rank -- belongs to the caller; for weights that is ``bind``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +class SymmBuffer: + """One symmetric-memory window, addressed by element offset. + + The driver caps windows at 128 per process regardless of size, so callers + pool many shards into one window. An offset names the same bytes on every + rank, which is what makes a peer read meaningful -- so a caller that hands + out slots must hand out the same ones on every rank. + """ + + def __init__( + self, dtype: torch.dtype, device: torch.device, group_name: str, numel: int, buf: torch.Tensor, handle + ) -> None: + """Wrap an already-rendezvous'd window. Callers want ``open``.""" + self.dtype = dtype + self.device = device + self.group_name = group_name + self.numel = numel + self.buf = buf + self.handle = handle + + @classmethod + def open(cls, dtype: torch.dtype, device: torch.device, group_name: str, numel: int) -> SymmBuffer: + """Allocate and rendezvous a window of ``numel`` elements. + + Collective: every rank must open the same windows in the same order. + """ + import torch.distributed._symmetric_memory as symm_mem + + symm_mem.enable_symm_mem_for_group(group_name) + buf = symm_mem.empty(numel, dtype=dtype, device=device) + handle = symm_mem.rendezvous(buf, group_name) + return cls(dtype, device, group_name, numel, buf, handle) + + def get_tensor(self, offset: int, shape: torch.Size | tuple[int, ...], *, rank: int | None = None) -> torch.Tensor: + """The slot at ``offset`` elements in, on ``rank`` (default: this rank). + + Not a slice of ``buf``: the returned tensor's storage starts at the slot. + Dynamo memoized ``storage_offset == 0`` for bound parameters, and a + mid-window slice would contradict the shape env. + """ + shape = tuple(int(s) for s in shape) + numel = 1 + for s in shape: + numel *= s + if offset < 0 or offset + numel > self.numel: + raise RuntimeError( + f"symmetric window overflow: {shape} at offset {offset} needs {offset + numel} " + f"elems of a {self.numel}-elem window" + ) + return self.handle.get_buffer(self.handle.rank if rank is None else rank, shape, self.dtype, offset) + + def peer_tensors(self, offset: int, shape: torch.Size | tuple[int, ...]) -> list[torch.Tensor]: + """``world_size`` views of one slot, one per rank. + + They borrow the window mapping; ``self.buf`` keeps it alive. + """ + return [self.get_tensor(offset, shape, rank=r) for r in range(self.handle.world_size)] + + @property + def nbytes(self) -> int: + return self.numel * self.dtype.itemsize + + def contains(self, t: torch.Tensor) -> bool: + base = self.buf.data_ptr() + return base <= t.data_ptr() < base + self.nbytes + + +@dataclass(frozen=True) +class ShardEntry: + """What a run-time gather needs to know about one local shard.""" + + buffer: SymmBuffer + offset: int + local: torch.Tensor + peer_views: tuple[torch.Tensor, ...] + + @property + def shape(self) -> tuple[int, ...]: + return tuple(self.local.shape) + + @property + def dtype(self) -> torch.dtype: + return self.local.dtype + + +_SHARD_REGISTRY: dict[int, ShardEntry] = {} +_BUFFERS: list[SymmBuffer] = [] + + +def open_buffer(dtype: torch.dtype, device: torch.device, group_name: str, numel: int) -> SymmBuffer: + """Open one window of ``numel`` elements, and keep it alive for the process. + + Whoever writes the window owes its peers a barrier before they read it; that + is the writer's business, not the registry's. + """ + buffer = SymmBuffer.open(dtype, device, group_name, numel) + _BUFFERS.append(buffer) + return buffer + + +def register_shard(local: torch.Tensor, buffer: SymmBuffer, offset: int) -> ShardEntry: + """Record a slot so the run-time gather can find its peer views.""" + entry = ShardEntry(buffer=buffer, offset=offset, local=local, peer_views=tuple(buffer.peer_tensors(offset, local.shape))) + _SHARD_REGISTRY[local.data_ptr()] = entry + return entry + + +def lookup_shard(data_ptr: int) -> ShardEntry | None: + """The registered shard starting at ``data_ptr``, or None if it is not one.""" + return _SHARD_REGISTRY.get(data_ptr) + + +def registered_buffers() -> list[SymmBuffer]: + return list(_BUFFERS) + + +def find_shard_by_layout(shape, dtype: torch.dtype) -> torch.Tensor | None: + """Any registered shard with this exact layout, for the runtime estimator's stand-in.""" + want = tuple(int(s) for s in shape) + for entry in _SHARD_REGISTRY.values(): + if entry.shape == want and entry.dtype == dtype: + return entry.local + return None + + +def reset_registry() -> None: + """Drop every window and shard. Tests only -- frees the symmetric allocations.""" + _SHARD_REGISTRY.clear() + _BUFFERS.clear() + + +def alloc_shard(shape, dtype: torch.dtype, device: torch.device, group_name: str) -> torch.Tensor: + """One shard in a window of its own. Tests and the cost model only -- binding a whole model must pool.""" + shape = tuple(int(s) for s in shape) + numel = 1 + for s in shape: + numel *= s + + buffer = open_buffer(dtype, device, group_name, numel) + shard = buffer.get_tensor(0, shape) + register_shard(shard, buffer, 0) + return shard diff --git a/tests/feature_tests/cache/test_cache_topology_isolation.py b/tests/feature_tests/cache/test_cache_topology_isolation.py index c7df040..70393f9 100644 --- a/tests/feature_tests/cache/test_cache_topology_isolation.py +++ b/tests/feature_tests/cache/test_cache_topology_isolation.py @@ -160,8 +160,9 @@ def forward(self, x): return x probe = Probe() + conf = SimpleNamespace(fsdp_config=SimpleNamespace(transport="nccl")) with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "cp8_dp1"}): - _magi_compile_bound_method(probe, {"x": 0}, SimpleNamespace(), "probe", method_name="forward") + _magi_compile_bound_method(probe, {"x": 0}, conf, "probe", method_name="forward") with patch.dict(os.environ, {"MAGI_COMPILE_TOPOLOGY_KEY": "cp4_dp2"}): with patch("magi_compiler._api._lazy_init_magi_state", side_effect=fake_init): diff --git a/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py index 7de0f86..4367926 100644 --- a/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py +++ b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py @@ -38,6 +38,16 @@ ``_negotiate_mode`` directly with synthetic per-rank inputs and assert it returns the expected mode for each rung of the ladder (identical / slot / pinned / abort). +With ``--copy-engine``: the same shape, but the gathers are +``magi::ce_all_gather`` reading a symmetric buffer. Two things are checked that +NCCL does not exercise. First, recognition: the gather is a plain fallback +kernel with an alias node between it and its wait, so the pass has to see through +that or it silently plans nothing. Second, slot safety: the gathers cycle +through a small set of RESIDENT destination buffers, and a launch hoisted above +the last read of the buffer it is about to overwrite would corrupt a weight in +flight. Inductor cannot infer that constraint -- the reuse is invisible in the +graph -- so ``REORDER_SLOTS`` asserts it directly on the emitted schedule. + Run: torchrun --nproc_per_node=1 tests/feature_tests/fsdp_overlap_helper/reorder_helper.py torchrun --nproc_per_node=2 ... reorder_helper.py --mismatch torchrun --nproc_per_node=2 ... reorder_helper.py --modes-only @@ -50,6 +60,7 @@ REORDER_MISMATCH local= (--mismatch only: divergent-graph path taken) REORDER_SLOT rank= (--mismatch only: SLOT-consensus mode chosen) REORDER_MODES ok= (--modes-only: the mode ladder returned as expected) + REORDER_SLOTS ok= (--copy-engine: no launch overwrites a live slot) REORDER_PASS / REORDER_FAIL """ @@ -58,9 +69,11 @@ import argparse import os -import torch -import torch._inductor.config as inductor_config -import torch.distributed as dist +os.environ.setdefault("TORCH_SYMM_MEM_DISABLE_MULTICAST", "1") + +import torch # noqa: E402 +import torch._inductor.config as inductor_config # noqa: E402 +import torch.distributed as dist # noqa: E402 from magi_compiler.passes.fsdp_overlap import FsdpOverlapReorder from magi_compiler.passes.fsdp_overlap import reorder as _ro @@ -110,6 +123,7 @@ def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--mismatch", action="store_true", help="rank1 compiles a structurally different graph") ap.add_argument("--modes-only", action="store_true", help="only self-check the mode ladder (gloo, no compile)") + ap.add_argument("--copy-engine", action="store_true", help="gather from a symmetric buffer instead of NCCL") args = ap.parse_args() if args.modes_only: @@ -138,6 +152,7 @@ def main() -> None: _WAIT = torch.ops._c10d_functional.wait_tensor.default H = 512 + N_CE_LAYERS = 3 w0 = torch.randn(H, H, device=dev, dtype=torch.bfloat16) shard = torch.randn(H, H, device=dev, dtype=torch.bfloat16) @@ -152,6 +167,43 @@ def fn(x, w0, shard): gathered = g.reshape(world * H, H)[:H] # use the gathered weight return y @ gathered + ce_shards: list = [] + if args.copy_engine: + from magi_compiler.symm_mem import alloc_shard + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER + + for i in range(N_CE_LAYERS): + s = alloc_shard((H, H), torch.bfloat16, torch.device("cuda", dev), grp) + s.normal_(0.0, H**-0.5).add_(0.01 * i) + ce_shards.append(s) + # A peer read is only legal once that peer has written its shard, and + # these shards are filled here, so the barrier belongs here too. + torch.cuda.synchronize() + dist.barrier() + + def fn(x, w0, shards): # noqa: F811 - deliberately replaces the NCCL variant + y = (x @ w0).relu() + y = _WAIT(_AR(y, "sum", grp)) + acc = None + for i, sh in enumerate(shards): + g = _WAIT(CE_ALL_GATHER(sh, world, grp)) + z = y @ g.reshape(world * H, H)[:H] + acc = z if acc is None else acc + z + return acc + + def ref_fn(x, w0, shards): + """Same arithmetic over NCCL: an independent answer, so the numeric + check is a real cross-transport comparison rather than the copy + engine grading its own homework.""" + y = (x @ w0).relu() + y = _WAIT(_AR(y, "sum", grp)) + acc = None + for sh in shards: + g = _WAIT(_AG(sh, world, grp)) + z = y @ g.reshape(world * H, H)[:H] + acc = z if acc is None else acc + z + return acc + # instrument the pass: count how many times it runs, how many launches move, # and whether the returned schedule is identical to the input (LOCAL path). calls = {"n": 0, "gathers": 0, "moved": 0, "unchanged": True, "warned_mismatch": False, "slot_mode": False} @@ -184,7 +236,16 @@ def spy(self, snodes): FsdpOverlapReorder.__call__ = spy - reorder = FsdpOverlapReorder(comm_overlap_window_margin_ns=5000.0) # default cost_fn (Inductor analytical) + def greedy_cost(snode) -> float: + """A gather nobody can hide: Inductor's analytical model prices a + fallback kernel at 0us, so without this the launches barely move.""" + from torch._inductor.comms import estimate_op_runtime + + if _ro._is_ce_ag_ir(_ro._leaf_collective_node(snode)): + return 1e7 # 10ms, far more than the whole graph's compute + return estimate_op_runtime(snode) + + reorder = FsdpOverlapReorder(comm_overlap_window_margin_ns=5000.0, cost_fn=greedy_cost if args.copy_engine else None) 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 @@ -194,9 +255,11 @@ def spy(self, snodes): try: torch._dynamo.reset() x = torch.randn(H, H, device=dev, dtype=torch.bfloat16) - eager = fn(x, w0, shard) + weights = ce_shards if args.copy_engine else shard + eager = (ref_fn if args.copy_engine else fn)(x, w0, weights) + torch.cuda.synchronize() compiled = torch.compile(fn, dynamic=False) - out = compiled(x, w0, shard) + out = compiled(x, w0, weights) torch.cuda.synchronize() finally: inductor_config.reorder_for_compute_comm_overlap = prev_flag @@ -221,6 +284,8 @@ def spy(self, snodes): ok_local = calls["n"] > 0 and calls["gathers"] >= 1 and numeric_ok and skeleton_ok if args.mismatch: ok_local = ok_local and calls["warned_mismatch"] and calls["slot_mode"] + if args.copy_engine: + ok_local = ok_local and calls["gathers"] >= N_CE_LAYERS t = torch.tensor([1 if ok_local else 0], device=dev) dist.all_reduce(t) all_ok = int(t.item()) == world diff --git a/tests/feature_tests/fsdp/fsdp_overlap_helper/uneven_shard_helper.py b/tests/feature_tests/fsdp/fsdp_overlap_helper/uneven_shard_helper.py new file mode 100644 index 0000000..13b8a23 --- /dev/null +++ b/tests/feature_tests/fsdp/fsdp_overlap_helper/uneven_shard_helper.py @@ -0,0 +1,269 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""torchrun entrypoint: every rank must make the SAME transport and bucketing +decision for an uneven ``Shard(0)`` weight. + +A collective is a joint operation: if one rank keeps a weight gather on NCCL while +its peers move it to the copy engine, the NCCL all-gather waits for peers that will +never arrive and the step hangs. Nothing in a single rank's graph reveals this, so +it cannot be tested in-process at one rank -- which is why it lives here. + +``Shard(0)`` gives ``ceil(F / world)`` rows to the leading ranks and the remainder +to the trailing ones, so with ``F % world != 0`` the ranks do not agree on the local +shard length. Two decisions were derived from that length and diverged because of +it: + + 1. **Transport and bucket membership.** The lowering pads the short shards up to a + full chunk, and copy-engine eligibility was decided by asking whether the + gather's input was the shard itself -- which a pad makes false. So the trailing + ranks kept the gather on NCCL and every other rank moved it to the copy engine. + 2. **Symmetric placement.** A shard's allocation is sized from the local numel, so + an uneven shard makes it rank-dependent -- and the gather copies one fixed-size + slab per peer. + +The byte cap on bucket size is checked here too, though it never diverged: it was +accounted from the gather's input, which the pad had already brought back to +``chunk``. It is asserted because that is a coincidence, not a design. + +A mixed graph (even and uneven weights interleaved, copy-engine transport) is the +production shape: one odd weight must not drag its even neighbours onto NCCL, and +the even ones must not drag the odd one onto the copy engine. + +Checked below, at ``world`` ranks, for an evenly and an unevenly divisible weight: +the full set of gather targets and bucket sizes is all-gathered and compared, and so +is the buffer window size. + +Driven by ``tests/feature_tests/test_uneven_shard_transport.py``. Run directly with:: + + torchrun --nproc_per_node=2 tests/feature_tests/fsdp_overlap_helper/uneven_shard_helper.py + +Markers (rank 0): + UNEVEN_TRANSPORT rows= agree= targets= + UNEVEN_NCCL_BUCKETS rows= agree= sizes= + UNEVEN_SYMM agree= in_buffer= skipped= + UNEVEN_MIXED agree= targets= sizes= + UNEVEN_PASS +""" + +from __future__ import annotations + +import argparse +import os +from collections import Counter +from typing import Sequence + +os.environ.setdefault("TORCH_SYMM_MEM_DISABLE_MULTICAST", "1") + +import torch # noqa: E402 +import torch.distributed as dist # noqa: E402 +import torch.fx as fx # noqa: E402 +import torch.nn as nn # noqa: E402 +from torch.distributed.device_mesh import init_device_mesh # noqa: E402 + +_NCCL_AG = torch.ops._c10d_functional.all_gather_into_tensor.default +_NCCL_AG_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default + + +def _build_graph(mesh, row_counts: Sequence[int], cols: int, dtype=torch.bfloat16) -> fx.GraphModule: + """One SimpleFSDP weight redistribute per entry of ``row_counts``. + + Every placeholder is declared before the first redistribute, as in a real traced + graph, so a hoisted coalesced launch stays topologically valid. The metas are real + DTensors on ``mesh``, which is the point: the local shard shape is whatever this + rank actually owns, and mixed even/uneven rows keep that per weight. + """ + from torch.distributed.tensor import Partial, Replicate, Shard, distribute_tensor + + g = fx.Graph() + weights = [] + replicated_metas = [] + for i, rows in enumerate(row_counts): + full = torch.zeros(rows, cols, device="cuda", dtype=dtype) + w = g.placeholder(f"layer_{i}_weight_parameter") + w.meta["example_value"] = distribute_tensor(full, mesh, [Shard(0)]) + weights.append(w) + replicated_metas.append(distribute_tensor(full, mesh, [Replicate()])) + + outs = [] + for w, replicated in zip(weights, replicated_metas): + rd = g.call_method("redistribute", (w,), {"placements": [Replicate()], "forward_dtype": None, "backward_dtype": None}) + rd.meta["example_value"] = replicated + tl = g.call_method("to_local", (rd,), {"grad_placements": [Partial()]}) + tl.meta["example_value"] = replicated._local_tensor + outs.append(tl) + g.output(tuple(outs)) + return fx.GraphModule(nn.Module(), g) + + +def _example_inputs(gm) -> list[object]: + """What Dynamo would hand the backend: the live weights, in placeholder order. + + Read off the placeholder metas here, since this graph is hand-built rather + than traced -- the metas are the real DTensors, so binding sees the same + rank-dependent local shapes it would in production. + """ + return [n.meta["example_value"] for n in gm.graph.find_nodes(op="placeholder")] + + +def _gather_targets(gm) -> tuple[dict[str, int], list[tuple[str, int]]]: + """Which transport each gather ended up on, and the size of each coalesced launch.""" + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED + + names = { + _NCCL_AG: "nccl", + _NCCL_AG_COALESCED: "nccl_coalesced", + CE_ALL_GATHER: "ce", + CE_ALL_GATHER_COALESCED: "ce_coalesced", + } + counts: Counter[str] = Counter() + sizes: list[tuple[str, int]] = [] + for node in gm.graph.nodes: + if node.op != "call_function": + continue + name = names.get(node.target) + if name is None: + continue + counts[name] += 1 + if name.endswith("_coalesced"): + sizes.append((name, len(node.args[0]))) + return dict(counts), sizes + + +def _agree(value) -> bool: + """True when every rank produced the same value.""" + seen = [None] * dist.get_world_size() + dist.all_gather_object(seen, value) + return all(v == seen[0] for v in seen) + + +def _check_transport(mesh, rows: int, *, n: int, cols: int, say) -> bool: + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.symm_mem import reset_registry + + reset_registry() # each check accounts for its own windows + gm = _build_graph(mesh, [rows] * n, cols) + lower_and_bucket_full_graph( + gm, "coalesced", bucket_size_bytes=0, transport="copy_engine", example_inputs=_example_inputs(gm) + ) + targets, sizes = _gather_targets(gm) + ok = _agree((targets, sizes)) + say(f"UNEVEN_TRANSPORT rows={rows} agree={ok} targets={targets} sizes={sizes}") + return ok + + +def _check_nccl_bucket_cap(mesh, rows: int, *, n: int, cols: int, cap: int, say) -> bool: + """The byte cap must cut buckets in the same place on every rank. + + Runs on the DEFAULT transport, where the uneven weight is bucketed rather than + excluded, so the cap is the only thing deciding membership. + """ + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + + gm = _build_graph(mesh, [rows] * n, cols) + lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=cap, transport="nccl") + _targets, sizes = _gather_targets(gm) + ok = _agree(sizes) + say(f"UNEVEN_NCCL_BUCKETS rows={rows} agree={ok} sizes={sizes}") + return ok + + +def _check_binding(mesh, *, even_rows: int, uneven_rows: int, cols: int, say) -> bool: + """An uneven shard must not be moved into symmetric memory. + + A shard's allocation is sized from the local numel, so an uneven weight makes it + rank-dependent, and the gather copies one fixed-size slab per peer: a rank whose + peers own fewer rows reads past the end of theirs. Nothing at run time + re-derives that, so it has to be excluded here. + """ + from torch.distributed.tensor import Shard, distribute_tensor + + from magi_compiler.symm_mem import bind_parameters, lookup_shard, registered_buffers, reset_registry + + reset_registry() + + def param(rows: int) -> nn.Parameter: + t = torch.zeros(rows, cols, device="cuda", dtype=torch.bfloat16) + return nn.Parameter(distribute_tensor(t, mesh, [Shard(0)])) + + named = [("even", param(even_rows)), ("uneven", param(uneven_rows))] + bind_parameters([p for _name, p in named]) + + window_bytes = sorted(w.nbytes for w in registered_buffers()) + bound = [name for name, p in named if lookup_shard(p._local_tensor.data_ptr())] + ok = _agree((window_bytes, bound)) and bound == ["even"] + say(f"UNEVEN_SYMM agree={ok} in_buffer={bound} window_bytes={window_bytes}") + return ok + + +def _check_mixed(mesh, *, even_rows: int, uneven_rows: int, cols: int, say) -> bool: + """Even and uneven weights in one copy-engine graph must split by transport. + + Interleaved so a program-order bucket would mix them: the even pair stays on + the copy engine as one coalesced launch, the uneven pair stays on NCCL as + another. All ranks must report the same split. + """ + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.symm_mem import reset_registry + + reset_registry() + gm = _build_graph(mesh, [even_rows, uneven_rows, even_rows, uneven_rows], cols) + lower_and_bucket_full_graph( + gm, "coalesced", bucket_size_bytes=0, transport="copy_engine", example_inputs=_example_inputs(gm) + ) + targets, sizes = _gather_targets(gm) + targets = dict(sorted(targets.items())) + sizes = sorted(sizes) + expected = ({"ce_coalesced": 1, "nccl_coalesced": 1}, [("ce_coalesced", 2), ("nccl_coalesced", 2)]) + ok = _agree((targets, sizes)) and (targets, sizes) == expected + say(f"UNEVEN_MIXED agree={ok} targets={targets} sizes={sizes}") + return ok + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--n-weights", type=int, default=4) + ap.add_argument("--cols", type=int, default=256) + ap.add_argument("--cap-bytes", type=int, default=2048, help="bucket byte cap that splits ranks if accounted locally") + args = ap.parse_args() + + rank = int(os.environ.get("RANK", "0")) + torch.cuda.set_device(rank % torch.cuda.device_count()) + dist.init_process_group("nccl") + world = dist.get_world_size() + assert world >= 2, "an uneven Shard(0) needs at least 2 ranks to be uneven" + mesh = init_device_mesh("cuda", (world,)) + + def say(msg: str) -> None: + if rank == 0: + print(msg, flush=True) + + # world+1 rows never divides by world (for world >= 2), so the last rank owns 1 row + # where the others own 2. world*2 rows is the even control. + uneven_rows, even_rows = world + 1, world * 2 + + ok = _check_transport(mesh, even_rows, n=args.n_weights, cols=args.cols, say=say) + ok &= _check_transport(mesh, uneven_rows, n=args.n_weights, cols=args.cols, say=say) + ok &= _check_nccl_bucket_cap(mesh, uneven_rows, n=args.n_weights, cols=args.cols, cap=args.cap_bytes, say=say) + ok &= _check_binding(mesh, even_rows=even_rows, uneven_rows=uneven_rows, cols=args.cols, say=say) + ok &= _check_mixed(mesh, even_rows=even_rows, uneven_rows=uneven_rows, cols=args.cols, say=say) + + if ok: + say("UNEVEN_PASS") + dist.destroy_process_group() + raise SystemExit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/tests/feature_tests/fsdp/symm_helper/__init__.py b/tests/feature_tests/fsdp/symm_helper/__init__.py new file mode 100644 index 0000000..3eaa44a --- /dev/null +++ b/tests/feature_tests/fsdp/symm_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/symm_helper/verify_symm_e2e.py b/tests/feature_tests/fsdp/symm_helper/verify_symm_e2e.py new file mode 100644 index 0000000..90f448d --- /dev/null +++ b/tests/feature_tests/fsdp/symm_helper/verify_symm_e2e.py @@ -0,0 +1,262 @@ +# 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. + +"""Step-3 acceptance: the whole copy-engine path, from a meta-built model. + +The other verification scripts each cut the chain somewhere. This one runs it +end to end the way inference actually does, which is the only way to exercise +the piece that cannot be faked: the weights are never allocated as ordinary +tensors at all. The model is built under ``torch.device("meta")``, sharded by +SimpleFSDP, and materialized with a single ``root.to_empty(cuda)`` -- and the +decorated block's patched ``_apply`` claims its own subtree on the way down. +The builder here is deliberately ignorant of symmetric memory, because the real +one is too. + +Then a checkpoint load writes into the buffer views (``copy_``, not ``assign``), +``@magi_compile`` compiles the block, the rewrite pass retargets the gathers, +and the reorder pass hoists them. Checked at the end: + + 1. **Placement** -- every block shard lives in a symmetric window, and the + head, outside the decorated block, does not. + 2. **Rewrite** -- the gathers really became ``magi::ce_all_gather``. A pass + that silently no-ops leaves a correct, NCCL-transported model behind, so + correctness alone cannot detect it. + 3. **Numerics** -- output matches an unsharded eager model holding the same + checkpoint, and it must still match on the SECOND step, when the resident + slots hold the previous step's weights. + 4. **Bytes moved** -- the resident-slot footprint, which is the memory this + transport trades for SM occupancy. + +Driven by ``tests/feature_tests/test_symm_e2e.py`` at two ranks; the checks are +asserted through the ``CHECK ...`` / ``E2E_PASS`` markers printed below. Run it +directly for a bigger, more realistic shape (2+ NVLink-connected GPUs):: + + torchrun --nproc_per_node=8 tests/feature_tests/symm_helper/verify_symm_e2e.py +""" + +from __future__ import annotations + +import argparse +import os + +os.environ.setdefault("TORCH_SYMM_MEM_DISABLE_MULTICAST", "1") + +import torch # noqa: E402 +import torch.distributed as dist # noqa: E402 +import torch.nn as nn # noqa: E402 +from torch.distributed.device_mesh import init_device_mesh # noqa: E402 + +from magi_compiler import magi_compile # noqa: E402 +from magi_compiler.config import CompileMode, CudaGraphMode # noqa: E402 + + +def _block_cls(transport: str) -> type: + """A decorated block, exactly as a model author would write one. + + Decoration happens at class-definition time, before any instance exists -- + which is what lets the transport choose where the weights are allocated. + """ + + class Block(nn.Module): + def __init__(self, hidden: int, n_layers: int, dtype: torch.dtype): + super().__init__() + self.layers = nn.ModuleList(nn.Linear(hidden, hidden, bias=False, dtype=dtype) for _ in range(n_layers)) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + def patch(cfg): + cfg.compile_mode = CompileMode.MAGI_COMPILE + cfg.cudagraph_mode = CudaGraphMode.NONE + cfg.disable_graph_split = True + cfg.fsdp_config.enable_fsdp = True + cfg.fsdp_config.transport = transport + return cfg + + return magi_compile(Block, config_patch=patch, dynamic_arg_dims={"x": 0}) + + +class Root(nn.Module): + """``to_empty`` is called here, never on the block: the interception has to + survive the recursion, or it never fires in a real model.""" + + def __init__(self, block_cls: type, hidden: int, n_layers: int, dtype: torch.dtype): + super().__init__() + self.block = block_cls(hidden, n_layers, dtype) + self.head = nn.Linear(hidden, hidden, bias=False, dtype=dtype) + + def forward(self, x): + return self.head(self.block(x)) + + +class _PlainBlock(nn.Module): + """Undecorated twin of ``Block``: same parameter names, so one state_dict + fits both, and no FSDP, so it gives an independent answer.""" + + def __init__(self, hidden: int, n_layers: int, dtype: torch.dtype): + super().__init__() + self.layers = nn.ModuleList(nn.Linear(hidden, hidden, bias=False, dtype=dtype) for _ in range(n_layers)) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +def _load_into_shards(model: nn.Module, state: dict[str, torch.Tensor]) -> None: + """Write each rank's slice of the checkpoint into the shard it already owns. + + ``copy_`` into the existing storage, never ``load_state_dict(assign=True)``: + assigning would swap the buffer views out for ordinary tensors and quietly + demote every gather back to NCCL. + """ + from torch.distributed.tensor import distribute_tensor + + with torch.no_grad(): + for name, p in _named_shards(model): + full = state[name] + want = distribute_tensor(full.to(p.device), p.device_mesh, p.placements) + p._local_tensor.copy_(want._local_tensor) + torch.cuda.synchronize() + + +def _named_shards(model: nn.Module): + """``(state_dict name, DTensor)`` for every parameter, reaching through + SimpleFSDP's parametrization the same way ``state_dict`` does.""" + from torch.distributed.tensor import DTensor + + for mod_name, mod in model.named_modules(): + for p_name, p in mod._parameters.items(): + if isinstance(p, DTensor): + yield (f"{mod_name}.{p_name}" if mod_name else p_name), p + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--hidden", type=int, default=4096) + ap.add_argument("--n-layers", type=int, default=6) + ap.add_argument("--n-tokens", type=int, default=4096) + ap.add_argument("--transport", default="copy_engine", choices=["nccl", "copy_engine"]) + args = ap.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("cpu:gloo,cuda:nccl", device_id=torch.device("cuda", local_rank)) + rank, world = dist.get_rank(), dist.get_world_size() + device = torch.device("cuda", local_rank) + dtype = torch.bfloat16 + mesh = init_device_mesh("cuda", (world,), mesh_dim_names=("dp",)) + os.environ.setdefault("MAGI_LOGGING_LEVEL", "INFO") + + def log(*a): + if rank == 0: + print(*a, flush=True) + + ce = args.transport == "copy_engine" + + # -- count the rewrites, without asking the compiled model to confess ---- + from magi_compiler.passes.fsdp_overlap import lower_and_bucket as _lb + + n_rewritten = 0 + orig_rewrite = _lb.rewrite_weight_ag_to_copy_engine + + def spy_rewrite(graph): + nonlocal n_rewritten + got = orig_rewrite(graph) + n_rewritten += got + return got + + _lb.rewrite_weight_ag_to_copy_engine = spy_rewrite + + # -- the reference: unsharded, eager, same checkpoint ------------------- + torch.manual_seed(1234) + ref = Root(_PlainBlock, args.hidden, args.n_layers, dtype).to(device) + with torch.no_grad(): + for p in ref.parameters(): + p.normal_(0.0, args.hidden**-0.5) + state = {k: v.detach().clone() for k, v in ref.state_dict().items()} + + x = torch.randn(args.n_tokens, args.hidden, device=device, dtype=dtype).mul_(0.02) + with torch.no_grad(): + truth = ref(x) + del ref + torch.cuda.empty_cache() + + # -- the model under test: meta build -> shard -> to_empty -> load ------ + from torchtitan.experiments.simple_fsdp.simple_fsdp import data_parallel + + with torch.device("meta"): + model = Root(_block_cls(args.transport), args.hidden, args.n_layers, dtype) + model = data_parallel(model, mesh, mode="fully_shard", ac_mode="full") + model.to_empty(device=device) + _load_into_shards(model, state) + + # (1) compile, then run twice. The second step is the one that would read a + # stale destination if the wait were missing. + with torch.no_grad(): + out1 = model(x) + torch.cuda.synchronize() + out2 = model(x) + torch.cuda.synchronize() + + _lb.rewrite_weight_ag_to_copy_engine = orig_rewrite + + # (2) placement: the block's shards are in a window, the head's are not. + # Checked after the first forward, not before it: the weights move during + # compilation, off the graph Dynamo captured. + from magi_compiler.symm_mem import lookup_shard, registered_buffers + + block_names = {n for n, _ in _named_shards(model) if n.startswith("block.")} + in_buffer = [n for n, p in _named_shards(model) if lookup_shard(p._local_tensor.data_ptr()) is not None] + head_in_buffer = [n for n in in_buffer if not n.startswith("block.")] + buffer_mib = sum(a.nbytes for a in registered_buffers()) / 2**20 + placement_ok = (set(in_buffer) == block_names) if ce else (in_buffer == []) + log( + f"CHECK placement: {len(in_buffer)}/{len(block_names)} block shards in {len(registered_buffers())} " + f"window(s) ({buffer_mib:.0f} MiB), {len(head_in_buffer)} stray -> {'ok' if placement_ok else 'WRONG'}" + ) + + expect_rewrites = args.n_layers if ce else 0 + rewrite_ok = n_rewritten == expect_rewrites + log(f"CHECK rewrite: {n_rewritten}/{expect_rewrites} gathers on magi::ce_all_gather -> {'ok' if rewrite_ok else 'WRONG'}") + + def rel(a, b): + return ((a.float() - b.float()).norm() / (b.float().norm() + 1e-6)).item() + + r1, r2 = rel(out1, truth), rel(out2, truth) + numeric_ok = bool(torch.isfinite(out1).all()) and max(r1, r2) < 5e-2 + log(f"CHECK numerics vs unsharded eager: step1 rel={r1:.6f} step2 rel={r2:.6f} -> {'ok' if numeric_ok else 'WRONG'}") + + log( + f"\nCONFIG world={world} hidden={args.hidden} layers={args.n_layers} tokens={args.n_tokens} " + f"transport={args.transport}\n" + f"MEMORY buffer={buffer_mib:.0f}MiB " + f"(one gathered layer = {args.hidden ** 2 * dtype.itemsize / 2**20:.0f}MiB)" + ) + + ok = placement_ok and rewrite_ok and numeric_ok + t = torch.tensor([1 if ok else 0], device=device) + dist.all_reduce(t, op=dist.ReduceOp.MIN) + all_ok = bool(t.item()) + log(f"\nE2E_{'PASS' if all_ok else 'FAIL'}") + + 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/test_ce_all_gather.py b/tests/feature_tests/fsdp/test_ce_all_gather.py new file mode 100644 index 0000000..240e2a3 --- /dev/null +++ b/tests/feature_tests/fsdp/test_ce_all_gather.py @@ -0,0 +1,268 @@ +# 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. + +"""Mechanics of ``magi::ce_all_gather`` (step 2 of the landing order). + +Single rank, so the copies are device-to-self and the values are trivially +checkable -- what is under test here is the plumbing that is easy to get subtly +wrong and hard to see: that the untouched ``wait_tensor`` really does pick up our +event through the work registry, that two in-flight gathers do not alias, and +that a shard the buffer never saw is rejected loudly rather than gathering +garbage. + +The transport itself (peer reads, overlap, ordering under load) needs several +NVLink-connected ranks and lives in +``example/inference/fsdp_overlap/verify_ce_ag_op.py``. +""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.nn as nn + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + +_WAIT = torch.ops._c10d_functional.wait_tensor + + +@pytest.fixture(scope="module") +def mesh_1rank(): + import torch.distributed as dist + + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29672") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + created = False + torch.cuda.set_device(0) + if not dist.is_initialized(): + dist.init_process_group("nccl", device_id=torch.device("cuda", 0)) + created = True + from torch.distributed.device_mesh import init_device_mesh + + yield init_device_mesh("cuda", (1,), mesh_dim_names=("dp",)) + if created: + dist.destroy_process_group() + + +@pytest.fixture(autouse=True) +def _clean_state(): + # Importing the module is what defines ``magi::ce_all_gather``. Nothing + # else here pulls it in, so without this the ops only exist when some other + # test in the same session happened to import them first. + import magi_compiler.symm_mem.all_gather # noqa: F401 + from magi_compiler.symm_mem import reset_registry + + reset_registry() + yield + reset_registry() + + +def _symm_model(mesh, hidden: int = 64, n_layers: int = 3, dtype=torch.bfloat16): + """A Shard(0)-sharded model whose weights have been bound into symmetric memory. + + Bound through ``bind_parameters`` rather than a graph: what this file tests is + the operator, and the operator only cares that the shard it is handed is a + registered one. + """ + from torch.distributed.tensor import Shard, distribute_tensor + + from magi_compiler.symm_mem import bind_parameters + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([nn.Linear(hidden, hidden, bias=False, dtype=dtype) for _ in range(n_layers)]) + + with torch.device("cuda", 0): + model = Block() + for mod in model.modules(): + for name, p in list(mod.named_parameters(recurse=False)): + mod.register_parameter(name, nn.Parameter(distribute_tensor(p, mesh, [Shard(0)]))) + + assert bind_parameters(list(model.parameters())) == n_layers + + shards = [p._local_tensor for p in model.parameters()] + for i, s in enumerate(shards): + s.copy_(torch.full_like(s, float(i + 1))) + torch.cuda.synchronize() + return model, shards + + +@requires_cuda +def test_gather_matches_nccl_bitwise(mesh_1rank): + _, shards = _symm_model(mesh_1rank) + + for shard in shards: + got = _WAIT(torch.ops.magi.ce_all_gather(shard, 1, "")) + ref = torch.empty_like(got) + torch.distributed.all_gather_into_tensor(ref, shard) + torch.cuda.synchronize() + assert torch.equal(got, ref), "disagrees with all_gather_into_tensor" + + +@requires_cuda +def test_coalesced_wrap_matches_per_member_gather(mesh_1rank): + """The thin wrap is per-member dests; each wait must match a single gather.""" + _, shards = _symm_model(mesh_1rank) + outs = torch.ops.magi.ce_all_gather_coalesced(list(shards), 1, "") + assert len(outs) == len(shards) + for out, shard in zip(outs, shards): + got = _WAIT(out) + ref = _WAIT(torch.ops.magi.ce_all_gather(shard, 1, "")) + assert torch.equal(got, ref) + + +@requires_cuda +def test_wait_tensor_picks_up_the_registered_event(mesh_1rank): + """The launch registers a Work; the *stock* wait_tensor must consume it.""" + _, shards = _symm_model(mesh_1rank) + shard = shards[0] + + out = _WAIT(torch.ops.magi.ce_all_gather(shard, 1, "")) + # No synchronize: the value must be correct because of the wait alone. + assert torch.equal(out, shard.expand_as(out)), "wait_tensor did not order the copies" + + +@requires_cuda +def test_each_gather_returns_a_fresh_buffer(mesh_1rank): + """Two live gathers must land in different buffers, or a prefetched weight + would be overwritten before its consumer ran.""" + _, shards = _symm_model(mesh_1rank) + a = torch.ops.magi.ce_all_gather(shards[0], 1, "") + b = torch.ops.magi.ce_all_gather(shards[1], 1, "") + _WAIT(a) + _WAIT(b) + torch.cuda.synchronize() + assert a.data_ptr() != b.data_ptr() + assert torch.equal(a, shards[0].expand_as(a)) + assert torch.equal(b, shards[1].expand_as(b)) + + +@requires_cuda +def test_unregistered_shard_is_rejected(mesh_1rank): + """A weight the buffer never claimed has no peer views, so gathering it would + read whatever happens to be at that address. Fail instead.""" + ordinary = torch.ones(8, 4, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="not a registered symmetric-memory shard"): + torch.ops.magi.ce_all_gather(ordinary, 1, "") + + +@requires_cuda +def test_group_size_mismatch_is_rejected(mesh_1rank): + _, shards = _symm_model(mesh_1rank) + with pytest.raises(RuntimeError, match="peers but the gather asks for"): + torch.ops.magi.ce_all_gather(shards[0], 4, "") + + +@requires_cuda +def test_coalesced_validates_every_member_before_allocating(mesh_1rank): + """One bad member must fail the whole call. Allocating the dests first and + discovering it halfway through would leave the earlier members' copies in + flight against buffers nobody waits on.""" + _, shards = _symm_model(mesh_1rank) + ordinary = torch.ones(8, 4, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="not a registered symmetric-memory shard"): + torch.ops.magi.ce_all_gather_coalesced([shards[0], ordinary], 1, "") + + +@requires_cuda +def test_meta_kernel_shape(mesh_1rank): + with torch.device("meta"): + local = torch.empty(4, 6, dtype=torch.bfloat16) + out = torch.ops.magi.ce_all_gather(local, 8, "") + assert out.shape == (32, 6) and out.is_meta + + +@requires_cuda +def test_coalesced_meta_kernel_shapes(mesh_1rank): + """Bucket members have different row counts, so the meta kernel cannot + broadcast one shape across the list -- Dynamo would trace the wrong dest.""" + with torch.device("meta"): + shards = [torch.empty(4, 6, dtype=torch.bfloat16), torch.empty(2, 6, dtype=torch.bfloat16)] + outs = torch.ops.magi.ce_all_gather_coalesced(shards, 8, "") + assert [tuple(o.shape) for o in outs] == [(32, 6), (16, 6)] + assert all(o.is_meta for o in outs) + + +# --------------------------------------------------------------------------- +# The copy plan -- raw pointer arithmetic, so nothing downstream can catch a +# wrong offset: it reads whatever is at that address. +# --------------------------------------------------------------------------- +@requires_cuda +def test_copy_plan_lands_each_rank_at_its_own_offset(): + """Rank r's shard belongs at row ``r * rows``. An off-by-one here is a + silently wrong weight, not a crash.""" + from magi_compiler.symm_mem.all_gather import _copy_triples, _gather_dest + + world = 4 + local = torch.ones(8, 4, device="cuda", dtype=torch.bfloat16) + peers = tuple(torch.full_like(local, r) for r in range(world)) + out = _gather_dest(local, world) + assert out.shape == (8 * world, 4) + + triples = _copy_triples([(out, local, peers)]) + nbytes = local.numel() * local.element_size() + assert triples == [(out.data_ptr() + r * nbytes, peers[r].data_ptr(), nbytes) for r in range(world)] + + +@requires_cuda +def test_copy_plan_concatenates_members_of_one_batch(): + """A coalesced gather is one submission covering every member; the members + must not share a destination.""" + from magi_compiler.symm_mem.all_gather import _copy_triples, _gather_dest + + world = 2 + a = torch.ones(8, 4, device="cuda", dtype=torch.bfloat16) + b = torch.ones(2, 4, device="cuda", dtype=torch.bfloat16) + gathers = [(_gather_dest(t, world), t, tuple(torch.empty_like(t) for _ in range(world))) for t in (a, b)] + + triples = _copy_triples(gathers) + assert len(triples) == 2 * world + assert {b for _d, _s, b in triples} == {a.numel() * a.element_size(), b.numel() * b.element_size()} + assert len({d for d, _s, _b in triples}) == 2 * world, "two copies target the same address" + + +@requires_cuda +def test_batch_plan_preserves_order_and_count(): + """``cudaMemcpyBatchAsync`` reads three parallel arrays; a reordering between + them pairs a destination with the wrong source.""" + from magi_compiler.symm_mem.all_gather import BatchMemcpy + + triples = [(0x1000, 0x2000, 64), (0x3000, 0x4000, 128)] + dsts, srcs, sizes, n = BatchMemcpy.plan(triples) + assert n == 2 + assert [dsts[i] for i in range(n)] == [0x1000, 0x3000] + assert [srcs[i] for i in range(n)] == [0x2000, 0x4000] + assert [sizes[i] for i in range(n)] == [64, 128] + + +@requires_cuda +def test_per_peer_fallback_matches_the_batched_path(mesh_1rank, monkeypatch): + """Older CUDA runtimes have no ``cudaMemcpyBatchAsync``; the fallback is the + only path there and is never exercised on the boxes we develop on.""" + from magi_compiler.symm_mem import all_gather as ag_mod + + _, shards = _symm_model(mesh_1rank) + batched = _WAIT(torch.ops.magi.ce_all_gather(shards[0], 1, "")) + torch.cuda.synchronize() + + monkeypatch.setattr(ag_mod, "_batcher", lambda: None) + fallback = _WAIT(torch.ops.magi.ce_all_gather(shards[0], 1, "")) + torch.cuda.synchronize() + + assert torch.equal(fallback, batched) diff --git a/tests/feature_tests/fsdp/test_copy_engine.py b/tests/feature_tests/fsdp/test_copy_engine.py new file mode 100644 index 0000000..eb886da --- /dev/null +++ b/tests/feature_tests/fsdp/test_copy_engine.py @@ -0,0 +1,380 @@ +# 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. + +"""Which gathers the copy engine may serve, and the retarget that acts on it. + +Two passes, one question each. ``copy_engine_weight_candidates`` decides what is +*possible* from the graph alone -- a gather whose input is a cast or a pad reads a +tensor the caching allocator produced, not a symmetric window. The retarget +decides nothing: it acts on ``node_meta.CE_BOUND``, which only binding sets, so a +weight whose allocation never happened cannot be retargeted into an operator that +would reject it at run time. + +Deliberately allocation-free, on a 1-rank gloo mesh: the allocation half is +``test_symm_bind.py``'s job, and keeping it out of here is what lets these run +anywhere. +""" + +from __future__ import annotations + +import operator +import os + +import pytest +import torch +import torch.fx as fx + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + +_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 +_TO_COPY = torch.ops.aten._to_copy.default +_PAD = torch.ops.aten.constant_pad_nd.default + + +@pytest.fixture(scope="module") +def mesh_1rank(): + import torch.distributed as dist + + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29673") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + created = False + torch.cuda.set_device(0) + if not dist.is_initialized(): + dist.init_process_group("gloo") + created = True + from torch.distributed.device_mesh import init_device_mesh + + yield init_device_mesh("cuda", (1,)) + if created: + dist.destroy_process_group() + + +def _symm_op(): + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER + + return CE_ALL_GATHER + + +def _graph_with_gathers( + mesh, n: int, *, derive: str | None = None, marked: bool = True, shapes: list[tuple] | None = None, uneven: bool = False +): + """``n`` weight gathers reading their shards, in program order. + + ``derive`` inserts a cast or a pad between the shard and the gather, the two + shapes the lowering pass emits for mixed precision and uneven sharding. + ``shapes`` gives the per-gather local shard shape. ``uneven`` sets the + rank-identical flag the lowering pass puts on a gather whose ``Shard(0)`` does + not divide evenly -- independently of ``derive``, because the ranks that own a + full chunk of an uneven weight get no pad at all. + """ + from torch.distributed.tensor import Shard, distribute_tensor + + from magi_compiler.passes.fsdp_overlap.node_meta import mark_weight_ag + + shapes = shapes or [(8, 4)] * n + + g = fx.Graph() + weights = [] + for i in range(n): + local = distribute_tensor(torch.randn(*shapes[i], device="cuda", dtype=torch.bfloat16), mesh, [Shard(0)]) + w = g.placeholder(f"layer_{i}_weight") + w.meta["example_value"] = local + weights.append((w, local)) + outs = [] + for w, local in weights: + cur = g.call_method("to_local", (w,)) + cur.meta["example_value"] = local._local_tensor + if derive == "cast": + cur = g.call_function(_TO_COPY, (cur,), {"dtype": torch.float32}) + cur.meta["example_value"] = local._local_tensor.to(torch.float32) + elif derive == "pad": + cur = g.call_function(_PAD, (cur, [0, 0, 0, 2], 0.0)) + cur.meta["example_value"] = local._local_tensor.new_empty((10, 4)) + ag = g.call_function(_AG, (cur, 1, "dummy_group")) + ag.meta["example_value"] = local._local_tensor.new_empty(local._local_tensor.shape) + if marked: + mark_weight_ag(ag, uneven=uneven) + outs.append(g.call_function(_WAIT, (ag,))) + g.output(tuple(outs)) + return fx.GraphModule(torch.nn.Module(), g) + + +def _gathers(gm, target): + return [n for n in gm.graph.nodes if n.op == "call_function" and n.target is target] + + +def _candidates(gm): + from magi_compiler.passes.fsdp_overlap import copy_engine_weight_candidates + + return copy_engine_weight_candidates(gm) + + +def _pretend_bound(gm) -> int: + """Mark every candidate, standing in for a ``bind_graph_weights`` that + succeeded. Keeps this file allocation-free while still deriving the mark from + the real selection pass rather than hand-placing it.""" + from magi_compiler.passes.fsdp_overlap.node_meta import mark_ce_bound + + marked = 0 + for gather, _holder in _candidates(gm): + mark_ce_bound(gather) + marked += 1 + return marked + + +# --------------------------------------------------------------------------- +# Selection +# --------------------------------------------------------------------------- +@requires_cuda +def test_plain_weight_gathers_are_candidates(mesh_1rank): + gm = _graph_with_gathers(mesh_1rank, 3) + candidates = _candidates(gm) + + assert [g for g, _h in candidates] == _gathers(gm, _AG) + assert all(h.op == "placeholder" for _g, h in candidates) + + +@requires_cuda +@pytest.mark.parametrize("derive", ["cast", "pad"]) +def test_derived_shards_are_not_candidates(mesh_1rank, derive): + """A cast or pad output is not in any symmetric window, so there would be no + peer views to read it from. + + Both graphs here are the same on every rank: ``forward_dtype`` is a model-wide + setting, and the pad is spliced into every gather. The case where only *some* + ranks see the pad is a different property, covered by + ``test_uneven_shards_are_not_candidates``. + """ + gm = _graph_with_gathers(mesh_1rank, 3, derive=derive) + assert _candidates(gm) == [] + + +@requires_cuda +def test_uneven_shards_are_not_candidates(mesh_1rank): + """The half of an uneven ``Shard(0)`` that the pad does not mark. + + ``ceil(F / world)`` rows go to the leading ranks and the remainder to the trailing + ones, so only the trailing ranks get a pad. A rank that owns a full chunk sees a + clean ``to_local`` and nothing in its own graph says the weight is unevenly split. + If it decides from the input shape alone it moves to the copy engine while its + peers stay on NCCL, and their all-gather waits on a rank that will never join it. + So the decision comes off the flag, which is derived from F and world and is + therefore the same everywhere. + """ + gm = _graph_with_gathers(mesh_1rank, 3, uneven=True) + assert _candidates(gm) == [] + + +@requires_cuda +def test_unmarked_gathers_are_not_candidates(mesh_1rank): + """CP / TP / MoE gathers are never marked as SimpleFSDP weight gathers, and + must survive copy-engine mode untouched -- the transports coexist in one graph.""" + gm = _graph_with_gathers(mesh_1rank, 3, marked=False) + assert _candidates(gm) == [] + + +# --------------------------------------------------------------------------- +# Retarget +# --------------------------------------------------------------------------- +@requires_cuda +def test_bound_gathers_are_retargeted_and_waits_untouched(mesh_1rank): + from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine + + gm = _graph_with_gathers(mesh_1rank, 4) + assert _pretend_bound(gm) == 4 + assert rewrite_weight_ag_to_copy_engine(gm) == 4 + + assert not _gathers(gm, _AG) + symm = _gathers(gm, _symm_op()) + assert len(symm) == 4 + # The wait is the load-bearing part of the design: it must be the same stock + # node, still reading the gather. + waits = _gathers(gm, _WAIT) + assert len(waits) == 4 + assert [w.args[0] for w in waits] == symm + + +@requires_cuda +def test_group_args_are_preserved(mesh_1rank): + """group_size / group_name stay in place, so the cost model reads them the + same way for either transport.""" + from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine + + gm = _graph_with_gathers(mesh_1rank, 3) + _pretend_bound(gm) + rewrite_weight_ag_to_copy_engine(gm) + assert all(n.args[1:3] == (1, "dummy_group") for n in _gathers(gm, _symm_op())) + + +@requires_cuda +def test_an_unbound_gather_is_never_retargeted(mesh_1rank): + """The retarget holds no opinion of its own: only binding knows whether the + allocation actually happened, and a second opinion here is a second chance to + disagree with the other ranks.""" + from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine + + gm = _graph_with_gathers(mesh_1rank, 3) # candidates, but nothing bound them + assert rewrite_weight_ag_to_copy_engine(gm) == 0 + assert len(_gathers(gm, _AG)) == 3 + assert not _gathers(gm, _symm_op()) + + +@requires_cuda +def test_mixed_graph_splits_by_transport(mesh_1rank): + from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine + + gm = _graph_with_gathers(mesh_1rank, 2) + # Splice in one unmarked gather reading a graph input directly. + node = _gathers(gm, _AG)[0] + with gm.graph.inserting_before(node): + other = gm.graph.placeholder("activation") + other.meta["example_value"] = torch.randn(4, 4, device="cuda") + extra = gm.graph.call_function(_AG, (other, 1, "dummy_group")) + extra.meta["example_value"] = torch.randn(4, 4, device="cuda") + gm.graph.lint() + + assert _pretend_bound(gm) == 2 + assert rewrite_weight_ag_to_copy_engine(gm) == 2 + assert len(_gathers(gm, _AG)) == 1 + assert len(_gathers(gm, _symm_op())) == 2 + + +@requires_cuda +def test_rewriting_nothing_leaves_the_graph_untouched(mesh_1rank): + """An NCCL-only graph passed through copy-engine mode must not be recompiled + into a different graph; the transports share this pass.""" + from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine + + gm = _graph_with_gathers(mesh_1rank, 2, marked=False) + before = gm.print_readable(print_output=False) + assert rewrite_weight_ag_to_copy_engine(gm) == 0 + assert gm.print_readable(print_output=False) == before + + +@requires_cuda +def test_selection_matches_what_the_lowering_pass_emits(mesh_1rank): + """The candidate shape has to match what lowering really produces, not what + this file thinks it produces.""" + from magi_compiler.passes.fsdp_overlap import lower_prim_redistribute_to_collectives, rewrite_weight_ag_to_copy_engine + + from .test_fsdp_overlap_lowering import _build_redistribute_graph + + gm = _build_redistribute_graph(mesh_1rank, "model_fc1_weight_parameter") + lower_prim_redistribute_to_collectives(gm) + + assert _pretend_bound(gm) == 1 + assert rewrite_weight_ag_to_copy_engine(gm) == 1 + assert not _gathers(gm, _AG) + symm = _gathers(gm, _symm_op()) + assert _gathers(gm, _WAIT)[0].args[0] is symm[0] + + +# --------------------------------------------------------------------------- +# Bucketing carries the mark +# --------------------------------------------------------------------------- +@requires_cuda +def test_a_bucket_of_bound_gathers_is_retargeted(mesh_1rank): + """Bucketing runs after binding, so the coalesced node it builds has to inherit + the mark from its members or the whole bucket falls back.""" + from magi_compiler.passes.fsdp_overlap import bucket_weight_all_gather_coalesced, rewrite_weight_ag_to_copy_engine + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED + + gm = _graph_with_gathers(mesh_1rank, 4) + _pretend_bound(gm) + assert bucket_weight_all_gather_coalesced(gm, bucket_size_bytes=0) == 1 + assert rewrite_weight_ag_to_copy_engine(gm) == 1 + + assert not _gathers(gm, _AG_COALESCED) + assert len(_gathers(gm, CE_ALL_GATHER_COALESCED)) == 1 + assert len(_gathers(gm, _WAIT)) == 4 + + +@requires_cuda +def test_a_bucket_is_all_or_nothing(mesh_1rank): + """A bucket is one submission, so it can only go to the copy engine if *every* + member is bound -- one unbound member has to keep the whole bucket on NCCL + rather than being gathered from an address with no peers.""" + from magi_compiler.passes.fsdp_overlap import bucket_weight_all_gather_coalesced, rewrite_weight_ag_to_copy_engine + from magi_compiler.passes.fsdp_overlap.node_meta import CE_BOUND + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED + + gm = _graph_with_gathers(mesh_1rank, 3) + _pretend_bound(gm) + _gathers(gm, _AG)[1].meta[CE_BOUND] = False + + # No eligibility filter, so the unbound member lands in the bucket anyway -- + # the case the propagation exists to catch. + assert bucket_weight_all_gather_coalesced(gm, bucket_size_bytes=0) == 1 + assert rewrite_weight_ag_to_copy_engine(gm) == 0 + assert len(_gathers(gm, _AG_COALESCED)) == 1 + assert not _gathers(gm, CE_ALL_GATHER_COALESCED) + + +def _coalesced_graph(mesh, n: int, *, marked: bool = True, bound: bool = True): + """One ``all_gather_into_tensor_coalesced`` over ``n`` shards. + + Bucketing normally builds this node, but it is also what a pre-bucketed graph + hands the rewrite, so the membership check gets its own graph rather than being + reached only through the bucket pass. + """ + from torch.distributed.tensor import Shard, distribute_tensor + + from magi_compiler.passes.fsdp_overlap.node_meta import mark_ce_bound, mark_weight_ag + + g = fx.Graph() + locals_ = [] + for i in range(n): + local = distribute_tensor(torch.randn(8, 4, device="cuda", dtype=torch.bfloat16), mesh, [Shard(0)]) + w = g.placeholder(f"layer_{i}_weight") + w.meta["example_value"] = local + cur = g.call_method("to_local", (w,)) + cur.meta["example_value"] = local._local_tensor + locals_.append(cur) + + ag = g.call_function(_AG_COALESCED, (locals_, 1, "dummy_group")) + ag.meta["example_value"] = [torch.empty(8, 4, device="cuda", dtype=torch.bfloat16) for _ in range(n)] + if marked: + mark_weight_ag(ag, uneven=False) + if bound: + mark_ce_bound(ag) + outs = [] + for i in range(n): + item = g.call_function(operator.getitem, (ag, i)) + outs.append(g.call_function(_WAIT, (item,))) + g.output(tuple(outs)) + return fx.GraphModule(torch.nn.Module(), g) + + +@requires_cuda +def test_a_pre_bucketed_bound_gather_is_retargeted(mesh_1rank): + from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED + + gm = _coalesced_graph(mesh_1rank, 3) + assert rewrite_weight_ag_to_copy_engine(gm) == 1 + assert not _gathers(gm, _AG_COALESCED) + assert len(_gathers(gm, CE_ALL_GATHER_COALESCED)) == 1 + + +@requires_cuda +def test_unmarked_coalesced_stays_on_nccl(mesh_1rank): + from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine + + gm = _coalesced_graph(mesh_1rank, 2, marked=False) + assert rewrite_weight_ag_to_copy_engine(gm) == 0 + assert len(_gathers(gm, _AG_COALESCED)) == 1 diff --git a/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py b/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py index 4d3d00a..72e1312 100644 --- a/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py +++ b/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py @@ -43,19 +43,26 @@ 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 + ``specs``: list of dicts, each ``{shape, dtype, compute_before?, local_shape?}``. + 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. + + ``shape`` is the chunk every rank agrees on, so the gathered output is + ``chunk * world`` rows. ``local_shape`` overrides only the placeholder, which is + how a trailing rank of an uneven ``Shard(0)`` looks: fewer rows locally, same + gathered size. That is the one thing a bucketing decision must not depend on. """ + from magi_compiler.passes.fsdp_overlap.node_meta import mark_weight_ag + 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") + loc.meta["example_value"] = torch.empty(*s.get("local_shape", s["shape"]), dtype=s["dtype"], device="meta") locs.append(loc) outs = [] @@ -70,7 +77,7 @@ def _build_ag_graph(specs, world=2, group="grp0"): 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 + mark_weight_ag(ag, uneven=False) w = g.call_function(_WAIT, (ag,)) w.meta["example_value"] = gathered outs.append(w) @@ -145,6 +152,38 @@ def test_bucket_size_bytes_caps_run(): assert _n(gm2, _AG_COALESCED) == 2 +def _bucket_sizes(gm) -> list[int]: + return [len(n.args[0]) for n in gm.graph.nodes if n.op == "call_function" and n.target is _AG_COALESCED] + + +def test_bucket_cap_does_not_depend_on_the_local_shard_length(): + """Where the byte cap cuts a bucket must be the same on every rank, and it must not + take a pad to make that true. + + A rank that fits an extra member submits a coalesced launch with a membership its + peers never submit, and the collective never completes. Graphs from this repo's + lowering are safe either way, because an uneven ``Shard(0)`` is padded back up to + ``chunk`` before the gather -- but ``_gathers_a_weight`` also matches gathers that + were never lowered here and have no pad, so the accounting is taken from the + gathered size instead of the input's. + + 4 gathers, 64 B per rank each (8x8 bf16 gathered over world=2); a 144 B cap fits + two of them and not three. In the ``tail`` graph this rank owns 1 row instead of + 4 for the second weight, with no pad to hide it, and must still cut in the same + place. + """ + even = [{"shape": (4, 8), "dtype": torch.bfloat16} for _ in range(4)] + tail = [dict(s) for s in even] + tail[1]["local_shape"] = (1, 8) + + gm_even, gm_tail = _build_ag_graph(even), _build_ag_graph(tail) + n_even = bucket_weight_all_gather_coalesced(gm_even, bucket_size_bytes=144) + n_tail = bucket_weight_all_gather_coalesced(gm_tail, bucket_size_bytes=144) + + assert (n_even, n_tail) == (2, 2) + assert _bucket_sizes(gm_even) == _bucket_sizes(gm_tail) == [2, 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).""" diff --git a/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py b/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py index 9374f18..e21977f 100644 --- a/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py +++ b/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py @@ -15,14 +15,14 @@ """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). +The pass matches a redistribute + to_local pair whose input is a weight +placeholder carrying a ``Shard(0)`` DTensor ``example_value``, and rewrites it +into explicit ``all_gather_into_tensor`` + ``wait_tensor``. Dynamo emits that +pair in two shapes -- ``prim_redistribute`` / ``prim_to_local`` call_functions +through torch 2.9, plain ``redistribute`` / ``to_local`` call_methods from 2.12 -- +and both are covered here, because a shape the pass fails to match produces no +error, just an un-lowered graph. Neither shape can be constructed by calling +into torch, so the graphs are synthetic, with REAL 1-rank DTensor metas. Uses a 1-rank process group + device mesh (GPU required). """ @@ -91,8 +91,36 @@ def _build_redistribute_graph(mesh, weight_name, dtype=torch.bfloat16, rows=8, c return fx.GraphModule(torch.nn.Module(), g) +def _build_method_redistribute_graph(mesh, weight_name, *, forward_dtype=None, dtype=torch.bfloat16, rows=8, cols=4): + """The same chain in the shape Dynamo emits from torch 2.12 on. + + There is no on-the-fly prim anymore: the method call itself is traced, and the + placements and dtypes that used to live in the prim's closure ride along in + ``kwargs``. A pass that only knows the older shape leaves this graph + un-lowered, and the only symptom is a copy-engine rewrite that finds nothing. + """ + from torch.distributed.tensor import Partial, 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_method( + "redistribute", (w,), {"placements": [Replicate()], "forward_dtype": forward_dtype, "backward_dtype": None} + ) + rd.meta["example_value"] = replicated + tl = g.call_method("to_local", (rd,), {"grad_placements": [Partial()]}) + 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 +_TO_COPY = torch.ops.aten._to_copy.default def _targets(gm): @@ -113,7 +141,9 @@ def test_lowering_rewrites_shard0_weight(dist_1rank): # 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")] + from magi_compiler.passes.fsdp_overlap import is_weight_ag + + tagged = [x for x in gm.graph.nodes if is_weight_ag(x)] assert len(tagged) == 1 @@ -130,6 +160,57 @@ def test_lowering_skips_non_weight_input(dist_1rank): assert prim_redistribute in targets # untouched +@requires_cuda +def test_lowering_rewrites_method_shaped_redistribute(dist_1rank): + """The 2.12+ shape must lower exactly like the prim one, marked gather included.""" + from magi_compiler.passes.fsdp_overlap import is_weight_ag, lower_prim_redistribute_to_collectives + + gm = _build_method_redistribute_graph(dist_1rank, "model_fc1_weight_parameter") + assert lower_prim_redistribute_to_collectives(gm) == 1 + + targets = _targets(gm) + assert _AG in targets and _WAIT in targets + methods = [n.target for n in gm.graph.nodes if n.op == "call_method"] + assert "redistribute" not in methods # consumed; only the to_local of the shard is left + assert len([x for x in gm.graph.nodes if is_weight_ag(x)]) == 1 + + +@requires_cuda +def test_lowering_reads_forward_dtype_from_method_kwargs(dist_1rank): + """Mixed precision is a kwarg on the method node, not a closure variable: miss + it and the gather moves the shard in its stored dtype, silently.""" + from magi_compiler.passes.fsdp_overlap import lower_prim_redistribute_to_collectives + + gm = _build_method_redistribute_graph(dist_1rank, "layer_weight", forward_dtype=torch.float32) + assert lower_prim_redistribute_to_collectives(gm) == 1 + + cast = [n for n in gm.graph.nodes if n.op == "call_function" and n.target is _TO_COPY] + assert len(cast) == 1 + assert cast[0].kwargs["dtype"] is torch.float32 + ag = [n for n in gm.graph.nodes if n.op == "call_function" and n.target is _AG] + assert ag[0].args[0] is cast[0] # the gather moves the cast shard, not the stored one + + +@requires_cuda +def test_lowering_records_shard_evenness_on_the_gather(dist_1rank): + """The gather carries whether its ``Shard(0)`` divides evenly. + + This is the flag the copy-engine rewrite and the bucketing read to decide the + transport, and it has to be present: a missing key reads as falsy, i.e. "even", + which is exactly the wrong default. It is computed from F and world, so it is the + same on every rank -- unlike the local shard length, which is what makes the pad + appear on the trailing ranks only. + """ + from magi_compiler.passes.fsdp_overlap import UNEVEN_SHARD, lower_prim_redistribute_to_collectives + + gm = _build_redistribute_graph(dist_1rank, "layer_weight", rows=8) + assert lower_prim_redistribute_to_collectives(gm) == 1 + + ag = [n for n in gm.graph.nodes if n.op == "call_function" and n.target is _AG] + assert UNEVEN_SHARD in ag[0].meta + assert ag[0].meta[UNEVEN_SHARD] is False # 8 rows over a 1-rank mesh + + @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 diff --git a/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py b/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py index 39388d4..17994a5 100644 --- a/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py +++ b/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py @@ -175,3 +175,22 @@ def test_reorder_graph_mismatch_slot_mode(): assert "REORDER_SLOT" in p.stdout, out[-3000:] # ... and it chose SLOT consensus assert "REORDER_SKELETON ok=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_copy_engine_recognition(): + """world=2, gathers served by the copy engine out of a symmetric buffer. + + The gather is an opaque fallback kernel, so if the pass fails to see through + it nothing is planned and ``gathers`` comes out 0. Destinations are fresh + ``empty``s -- Inductor owns liveness -- so the only extra invariant is that + the three launches are recognized and the compiled answer matches NCCL. + """ + p = _run(2, "--copy-engine", port="29634") + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-3000:]}" + assert "REORDER_CALLED gathers=3" in p.stdout, out[-3000:] # all three were recognized + assert "REORDER_FINITE ok=True" in p.stdout, out[-3000:] # ... and match the NCCL answer + assert "REORDER_PASS" in p.stdout, out[-3000:] diff --git a/tests/feature_tests/fsdp/test_profiling_estimator.py b/tests/feature_tests/fsdp/test_profiling_estimator.py index 6164539..2420684 100644 --- a/tests/feature_tests/fsdp/test_profiling_estimator.py +++ b/tests/feature_tests/fsdp/test_profiling_estimator.py @@ -14,7 +14,8 @@ """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. +shape helper, the estimator's memoization/deepcopy, the extern replay measure, +and the copy-engine gather branch. """ import copy @@ -26,7 +27,19 @@ 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 +from magi_compiler.profiling import runtime_estimator as re_mod +from magi_compiler.profiling.runtime_estimator import ( + ProfileEntry, + _ce_ag_label, + _ce_ag_launch_wait, + _ce_ag_spec, + _is_ce_ag_coalesced_ir, + _leaf_ce_ag, + _measure_ce_ag, + _measure_extern, + _realize_arg, + _static, +) from magi_compiler.utils.envs import TORCH_VERSION requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -265,7 +278,6 @@ def test_internal_collective_extern_not_measured_in_sync_warmup(monkeypatch, syn -> hang). It is seeded analytical + stashed for warm_and_sync. In non-sync mode the normal measurement path still runs.""" from magi_compiler.profiling import register_materialize_inputs - from magi_compiler.profiling import runtime_estimator as re_mod from magi_compiler.profiling.materialize_inputs import _INTERNAL_COLLECTIVE_OPS, _MATERIALIZE_INPUT_HOOKS measured = {"called": False} @@ -293,6 +305,25 @@ def _boom(*a, **k): _INTERNAL_COLLECTIVE_OPS.discard("aten::mm") +# --------------------------------------------------------------------------- +# window scale: need = comm * scale + margin +# --------------------------------------------------------------------------- +def test_reorder_window_params_survive_deepcopy(): + """Both terms of the window budget must survive the deepcopy Inductor does + when it folds this pass into the fx-graph cache key. ``__deepcopy__`` + copies the fields by hand, so a dropped line silently reverts one of them + to its default instead of raising.""" + import copy + + from magi_compiler.passes.fsdp_overlap import FsdpOverlapReorder + + pass_obj = FsdpOverlapReorder(comm_overlap_window_scale=2.0, comm_overlap_window_margin_ns=1234.0) + clone = copy.deepcopy(pass_obj) + for obj in (pass_obj, clone): + assert obj.comm_overlap_window_scale == pytest.approx(2.0) + assert obj.comm_overlap_window_margin_ns == pytest.approx(1234.0) + + import torch._inductor.config as inductor_config # noqa: E402 # =========================================================================== @@ -509,3 +540,303 @@ def test_collective_profile_accuracy_multi_rank(): # key-set intersection: shared keys still measured; no hang on mismatch assert "COLL_MISMATCH ok=True" in p.stdout, out[-3000:] assert "COLL_PASS" in p.stdout, out[-3000:] + + +# =========================================================================== +# COPY-ENGINE GATHER (``fsdp_config.transport="copy_engine"``). +# +# The reorder pass sizes each gather's overlap window from what the estimator +# returns here, and this is the one branch with no self-check: a gather priced +# at 0ns is not a wrong number that surfaces as a wrong answer, it is a gather +# with no window -- so nothing is hoisted, the transport still works, and the +# only symptom is that the speedup is missing. Both ways that has happened are +# pinned below: reading the spec off ``get_origin_node()`` (Inductor leaves it +# unset on these nodes) and timing the launch without its wait (the copies run +# on a side stream, so the launch alone times a ~3us CPU issue). +# +# The IR these helpers read is three attributes deep, so the structural tests +# use stand-ins rather than a full compile. The replay tests need a real +# symmetric window and run on one rank, where a peer read reads our own shard. +# =========================================================================== +_CE_ROWS = 128 + + +def _ce_ops(): + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED + + return CE_ALL_GATHER, CE_ALL_GATHER_COALESCED + + +class _FakeLayout: + def __init__(self, size, dtype): + self.size = size + self.dtype = dtype + + +class _FakeBuf: + def __init__(self, shape, dtype): + self.layout = _FakeLayout(tuple(shape), dtype) + + +class _FakeCeAgIR: + """A ``magi::ce_all_gather`` FallbackKernel. + + ``get_origin_node`` returns None on purpose: that is what the real node + does, and it is why the spec has to come from ``constant_args``. + """ + + def __init__(self, op, shapes, dtype=torch.bfloat16, group_size=1, group_name="grp"): + self.op_overload = op + self.inputs = [_FakeBuf(s, dtype) for s in shapes] + self.constant_args = (group_size, group_name) + + def get_origin_node(self): + return None + + +class _FakeMmIR: + def __init__(self): + self.op_overload = torch.ops.aten.mm.default + self.inputs = [] + self.constant_args = () + + +class _FakeGatherSnode: + def __init__(self, node=None, snodes=()): + self.node = node + self.snodes = list(snodes) + + +def _ce_snode(shape=(_CE_ROWS, 64), group_size=1): + ag, _ = _ce_ops() + return _FakeGatherSnode(node=_FakeCeAgIR(ag, [shape], group_size=group_size)) + + +# --- spec / recognition / label: structural, no window needed --------------- +def test_ce_ag_spec_is_read_off_constant_args_not_the_origin_node(): + """``get_origin_node()`` is unset on these nodes; reading the group from it + yields None, which degraded the gather to a zero cost and a zero window.""" + ag, _ = _ce_ops() + node = _FakeCeAgIR(ag, [(_CE_ROWS, 64)], group_size=8, group_name="fsdp") + assert node.get_origin_node() is None # the trap this test exists for + + shapes, dtype, group_size, group_name = _ce_ag_spec(node) + assert shapes == ((_CE_ROWS, 64),) + assert dtype is torch.bfloat16 + assert group_size == 8 + assert group_name == "fsdp" + + +def test_ce_ag_spec_of_a_coalesced_gather_has_one_shape_per_member(): + """The bucket's cost is the whole batch, so every member's shape must be in + the spec -- and in the cache key built from it.""" + _, coalesced = _ce_ops() + shapes, _dtype, group_size, _gn = _ce_ag_spec(_FakeCeAgIR(coalesced, [(8, 4), (16, 4), (32, 4)], group_size=4)) + assert shapes == ((8, 4), (16, 4), (32, 4)) + assert group_size == 4 + + +@pytest.mark.parametrize("constant_args", [(), (1,)], ids=["empty", "group_name_only"]) +def test_ce_ag_spec_is_none_when_the_group_args_are_missing(constant_args): + ag, _ = _ce_ops() + node = _FakeCeAgIR(ag, [(8, 4)]) + node.constant_args = constant_args + assert _ce_ag_spec(node) is None + + +def test_ce_ag_spec_is_none_when_the_node_has_no_inputs(): + ag, _ = _ce_ops() + assert _ce_ag_spec(_FakeCeAgIR(ag, [])) is None + + +def test_leaf_ce_ag_finds_the_gather_in_a_plain_and_in_a_fused_snode(): + """Inductor may hand the pass either the gather's own snode or a fused snode + that contains it; the cost model has to price both.""" + ag, _ = _ce_ops() + node = _FakeCeAgIR(ag, [(8, 4)]) + + assert _leaf_ce_ag(_FakeGatherSnode(node=node)) is node + fused = _FakeGatherSnode(node=None, snodes=[_FakeGatherSnode(node=_FakeMmIR()), _FakeGatherSnode(node=node)]) + assert _leaf_ce_ag(fused) is node + + +def test_leaf_ce_ag_is_none_for_an_ordinary_kernel(): + """A copy-engine gather is a plain FallbackKernel, so the probe cannot key on + 'is a fallback' -- an mm must not be mistaken for one.""" + assert _leaf_ce_ag(_FakeGatherSnode(node=_FakeMmIR())) is None + assert _leaf_ce_ag(_FakeGatherSnode()) is None + + +def test_coalesced_gather_is_distinguished_from_a_single_one(): + ag, coalesced = _ce_ops() + assert _is_ce_ag_coalesced_ir(_FakeCeAgIR(coalesced, [(8, 4)])) + assert not _is_ce_ag_coalesced_ir(_FakeCeAgIR(ag, [(8, 4)])) + + +def test_ce_ag_label_reports_transport_world_size_and_member_count(): + """``summary()`` is diffed against nsys traces by hand, so a copy-engine + gather must not be labelled like the NCCL one it replaced.""" + ag, coalesced = _ce_ops() + single = _ce_ag_label(_FakeGatherSnode(node=_FakeCeAgIR(ag, [(8, 4)], group_size=2))) + batch = _ce_ag_label(_FakeGatherSnode(node=_FakeCeAgIR(coalesced, [(8, 4), (16, 4)], group_size=2))) + + assert single == "ce_all_gather(ws=2,8x4)" + assert batch == "ce_all_gather_coalesced(ws=2,n=2,8x4)" + + +# --- replay: needs a real symmetric window ---------------------------------- +@pytest.fixture +def symm_registry(): + """CE tests register real shards; leaking one across tests would let a stale + layout answer ``find_shard_by_layout`` after its window is gone.""" + import magi_compiler.symm_mem.all_gather # noqa: F401 - importing defines the ops + from magi_compiler.symm_mem import reset_registry + + reset_registry() + yield + reset_registry() + + +def _ce_group_name() -> str: + import torch.distributed as dist + + return dist.group.WORLD.group_name + + +def _register_shards(shapes, dtype=torch.bfloat16): + """Register ``shapes`` as real symmetric-memory shards, filled distinctly.""" + import torch.distributed as dist + + from magi_compiler.symm_mem import alloc_shard + + shards = [] + for i, shape in enumerate(shapes): + s = alloc_shard(shape, dtype, torch.device("cuda", 0), _ce_group_name()) + s.fill_(i + 1) + shards.append(s) + # A peer read is only legal once every rank has filled its shard. + torch.cuda.synchronize() + dist.barrier() + return shards + + +def _ce_replay_snode(shapes, coalesced=False): + ag, ag_coalesced = _ce_ops() + op = ag_coalesced if coalesced else ag + return _FakeGatherSnode(node=_FakeCeAgIR(op, shapes, group_size=1, group_name=_ce_group_name())) + + +@requires_cuda +def test_ce_ag_replay_gathers_the_registered_shard(pg_1rank, symm_registry): + """The replay must run the real op on a real SymmBuffer shard: a gather of an + ordinary ``empty`` has no peers and would be rejected, leaving the cost model + on the analytical estimate it was installed to replace.""" + (shard,) = _register_shards([(_CE_ROWS, 64)]) + + launch, wait = _ce_ag_launch_wait(_ce_replay_snode([(_CE_ROWS, 64)])) + out = wait(launch()) + # No synchronize: if the timed closure did not include the wait, this is a + # race -- which is exactly the cost-model bug being pinned. + assert torch.equal(out, shard.expand_as(out)) + + +@requires_cuda +def test_ce_ag_coalesced_replay_covers_every_member(pg_1rank, symm_registry): + shapes = [(_CE_ROWS, 64), (_CE_ROWS, 32)] + shards = _register_shards(shapes) + + launch, wait = _ce_ag_launch_wait(_ce_replay_snode(shapes, coalesced=True)) + outs = wait(launch()) + assert len(outs) == len(shards) + for out, shard in zip(outs, shards): + assert torch.equal(out, shard.expand_as(out)) + + +@requires_cuda +def test_measure_ce_ag_prices_the_gather_above_zero(pg_1rank, symm_registry): + """The whole point of the CE branch: a real number, not Inductor's 0us for a + fallback kernel.""" + _register_shards([(1024, 1024)]) + assert _measure_ce_ag(_ce_replay_snode([(1024, 1024)])) > 0.0 + + +@requires_cuda +def test_measure_ce_ag_degrades_quietly_when_no_shard_has_that_layout(pg_1rank, symm_registry): + """Cast/pad gathers keep NCCL, so a CE-shaped node with no matching shard is + reachable. It must fall back, not raise inside the scheduler callback.""" + _register_shards([(_CE_ROWS, 64)]) + snode = _ce_replay_snode([(7, 5)]) + + assert _ce_ag_launch_wait(snode) is None + assert _measure_ce_ag(snode) == 0.0 + + +# --- __call__: the ("ce_ag", ...) cache ----------------------------------- +@pytest.fixture +def ce_measure_calls(monkeypatch): + """Drive ``__call__`` straight into the copy-engine branch and count measures.""" + calls = [] + + def fake_measure(snode): + calls.append(snode) + return 222.0 + + monkeypatch.setattr(re_mod, "contains_wait", lambda s: False) + monkeypatch.setattr(re_mod, "_is_multi_output_unpack", lambda s: False) + monkeypatch.setattr(re_mod, "_safe_analytical", lambda s: 111.0) + monkeypatch.setattr(re_mod, "_measure_ce_ag", fake_measure) + return calls + + +def test_isomorphic_ce_ag_gathers_are_measured_once(ce_measure_calls): + """Every layer gathers the same shape. Measuring each one would make compile + time linear in depth for no new information.""" + est = ProfilingRuntimeEstimator() + + first = est(_ce_snode()) + second = est(_ce_snode()) + + assert first == second == 222.0 + assert len(ce_measure_calls) == 1, "the second gather re-measured instead of hitting the cache" + assert est.n_cache_hits == 1 + assert len(est.table) == 1 + + +def test_ce_ag_cache_key_separates_shape_and_world_size(ce_measure_calls): + """Sharing an entry across shapes would price a small gather like a large one + and size its window from the wrong transfer.""" + est = ProfilingRuntimeEstimator() + est(_ce_snode(shape=(_CE_ROWS, 64))) + est(_ce_snode(shape=(_CE_ROWS, 32))) + est(_ce_snode(shape=(_CE_ROWS, 64), group_size=8)) + + assert len(ce_measure_calls) == 3 + assert len(est.table) == 3 + assert est.n_cache_hits == 0 + + +def test_ce_ag_entry_is_tagged_as_a_copy_engine_gather(ce_measure_calls): + est = ProfilingRuntimeEstimator() + est(_ce_snode()) + + (entry,) = est.table.values() + assert entry.kind == "ce_ag" + assert entry.measured + assert entry.ns == 222.0 + assert "ce_all_gather" in entry.label + + +def test_ce_ag_sync_mode_defers_the_measurement_to_warm_and_sync(ce_measure_calls): + """In profile_sync mode every rank must measure the same keys in the same + order. Measuring here instead would let a rank whose graph reaches the + gather first run copies the others have not issued.""" + est = ProfilingRuntimeEstimator() + est._sync_across_ranks = True + + ns = est(_ce_snode()) + + assert ns == 111.0, "sync mode must seed with the analytical estimate" + assert ce_measure_calls == [], "sync mode measured inside __call__" + (entry,) = est.table.values() + assert not entry.measured + assert list(est._key_snode) == list(est.table), "snode not stashed for warm_and_sync" diff --git a/tests/feature_tests/fsdp/test_symm_bind.py b/tests/feature_tests/fsdp/test_symm_bind.py new file mode 100644 index 0000000..27d0454 --- /dev/null +++ b/tests/feature_tests/fsdp/test_symm_bind.py @@ -0,0 +1,561 @@ +# 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. + +"""Graph-driven binding of weights into symmetric memory. + +Two properties matter here and neither is about bytes moving. The first is +*scope*: exactly the weights the graph gathers get an allocation, and a weight +sitting next to them in the same model does not. The second is *identity*: the +parameter object Dynamo already guarded on has to survive the move, because the +call that triggered this compilation is going to run against those same objects. + +Single rank on purpose -- selection, identity, layout and every fallback are +rank-independent, and a 1-rank window exercises the same allocation path. The +cross-rank half (peer reads, plan agreement) is covered by ``test_symm_e2e.py`` +and ``test_uneven_shard_transport.py``. +""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.fx as fx + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + +_AG = torch.ops._c10d_functional.all_gather_into_tensor.default +_WAIT = torch.ops._c10d_functional.wait_tensor.default +_TO_COPY = torch.ops.aten._to_copy.default + + +@pytest.fixture(scope="module") +def mesh_1rank(): + """A single-rank NCCL group + cuda device mesh (symmetric memory needs both).""" + 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 + torch.cuda.set_device(0) + if not dist.is_initialized(): + dist.init_process_group("nccl", device_id=torch.device("cuda", 0)) + created = True + from torch.distributed.device_mesh import init_device_mesh + + yield init_device_mesh("cuda", (1,), mesh_dim_names=("dp",)) + if created: + dist.destroy_process_group() + + +@pytest.fixture +def group_name(mesh_1rank) -> str: + return mesh_1rank.get_group("dp").group_name + + +@pytest.fixture(autouse=True) +def _clean_registry(): + from magi_compiler.symm_mem import reset_registry + + reset_registry() + yield + reset_registry() + + +def _param(mesh, shape=(8, 4), dtype=torch.bfloat16, placement=None): + from torch.distributed.tensor import Shard, distribute_tensor + + full = torch.randn(*shape, device="cuda", dtype=dtype) + return distribute_tensor(full, mesh, [placement or Shard(0)]) + + +def _graph(params: list, *, derive: str | None = None, uneven: bool = False, extra_placeholder=None): + """One marked weight all-gather per parameter, in program order. + + ``derive`` splices the mixed-precision cast between the shard and the gather; + ``uneven`` sets the rank-identical flag the lowering pass puts on a weight + whose ``Shard(0)`` does not divide evenly. Returns the module and the + ``example_inputs`` Dynamo would hand the backend -- the live parameters, in + placeholder order. + """ + from magi_compiler.passes.fsdp_overlap.node_meta import mark_weight_ag + + g = fx.Graph() + examples: list[object] = [] + + # Every placeholder first, as Dynamo emits them: bucketing hoists a member's + # shard prep above the first gather, which is only legal if the weights it + # reads are already defined up there. + weights = [] + for i, p in enumerate(params): + w = g.placeholder(f"layer_{i}_weight") + w.meta["example_value"] = p + examples.append(p) + weights.append((w, p._local_tensor if hasattr(p, "_local_tensor") else p)) + if extra_placeholder is not None: + # A weight the graph never gathers: it must come out of binding untouched. + g.placeholder("spare_weight") + examples.append(extra_placeholder) + + outs = [] + for w, local in weights: + cur = g.call_method("to_local", (w,)) + cur.meta["example_value"] = local + if derive == "cast": + cur = g.call_function(_TO_COPY, (cur,), {"dtype": torch.float32}) + cur.meta["example_value"] = local.to(torch.float32) + ag = g.call_function(_AG, (cur, 1, "dummy_group")) + ag.meta["example_value"] = local.new_empty(local.shape) # world size 1 + mark_weight_ag(ag, uneven=uneven) + outs.append(g.call_function(_WAIT, (ag,))) + g.output(tuple(outs)) + return fx.GraphModule(torch.nn.Module(), g), examples + + +def _bind(gm, examples, **kw): + """Binding on its own, standing in for what the backend does around it.""" + from magi_compiler.passes.fsdp_overlap import copy_engine_weight_candidates + from magi_compiler.symm_mem import bind_graph_weights + + names = (n.name for n in gm.graph.find_nodes(op="placeholder")) + return bind_graph_weights(gm, copy_engine_weight_candidates(gm), dict(zip(names, examples)), **kw) + + +# --------------------------------------------------------------------------- +# The allocation itself +# --------------------------------------------------------------------------- +@requires_cuda +def test_a_shard_starts_at_its_own_storage(mesh_1rank, group_name): + """Every shard owns its whole window. + + ``storage_offset`` 0 is not cosmetic: it is what Dynamo memoized for these + parameters before the backend ran, and a shard part-way into a shared window + makes the next fakification of the same source disagree with itself. The + 16B alignment is what Inductor assumes of every graph input. + """ + from magi_compiler.symm_mem import alloc_shard + + shard = alloc_shard((8, 4), torch.bfloat16, torch.device("cuda", 0), group_name) + assert shard.storage_offset() == 0 + assert shard.shape == (8, 4) + assert shard.is_contiguous() + assert shard.data_ptr() % 16 == 0 + + +@requires_cuda +def test_peer_view_round_trips_to_the_shard_itself(mesh_1rank, group_name): + """At one rank the only peer view is this rank's, so it must alias the shard -- + the cheapest check that the window really is what got registered.""" + from magi_compiler.symm_mem import alloc_shard, lookup_shard + + shard = alloc_shard((8, 4), torch.bfloat16, torch.device("cuda", 0), group_name) + shard.fill_(7.0) + + entry = lookup_shard(shard.data_ptr()) + assert entry is not None + assert len(entry.peer_views) == 1 + assert torch.equal(entry.peer_views[0], shard) + + +@requires_cuda +def test_contains_rejects_memory_outside_the_allocation(mesh_1rank, group_name): + """A caching-allocator tensor must never look like a symmetric shard.""" + from magi_compiler.symm_mem import alloc_shard, registered_buffers + + shard = alloc_shard((8, 8), torch.bfloat16, torch.device("cuda", 0), group_name) + (buffer,) = registered_buffers() + + assert buffer.contains(shard) + assert not buffer.contains(torch.empty(8, 8, device="cuda", dtype=torch.bfloat16)) + + +@requires_cuda +def test_find_shard_by_layout_matches_on_shape_and_dtype(mesh_1rank, group_name): + """The cost model replays a gather on *some* registered shard with the right + layout -- it cannot use a generic ``empty``, which has no peers. A miss has to + be a None it can degrade on, not a wrong-dtype shard it would gather.""" + from magi_compiler.symm_mem import alloc_shard, find_shard_by_layout + + dev = torch.device("cuda", 0) + small = alloc_shard((8, 4), torch.bfloat16, dev, group_name) + large = alloc_shard((16, 4), torch.bfloat16, dev, group_name) + + assert find_shard_by_layout((8, 4), torch.bfloat16) is small + assert find_shard_by_layout((16, 4), torch.bfloat16) is large + assert find_shard_by_layout((8, 5), torch.bfloat16) is None + assert find_shard_by_layout((8, 4), torch.float32) is None + + +@requires_cuda +def test_reset_registry_drops_buffers_and_shards(mesh_1rank, group_name): + """Tests and the multi-model path rebuild in-process; a stale entry would let a + freed shard's address answer a lookup.""" + from magi_compiler.symm_mem import alloc_shard, find_shard_by_layout, lookup_shard, registered_buffers, reset_registry + + shard = alloc_shard((8, 4), torch.bfloat16, torch.device("cuda", 0), group_name) + assert lookup_shard(shard.data_ptr()) is not None + + reset_registry() + assert registered_buffers() == [] + assert lookup_shard(shard.data_ptr()) is None + assert find_shard_by_layout((8, 4), torch.bfloat16) is None + + +@requires_cuda +def test_group_name_of_refuses_a_mesh_it_cannot_resolve(mesh_1rank): + """Defaulting to WORLD would open the window on the wrong group and read peers + holding somebody else's rows, which is silent corruption rather than a crash.""" + from magi_compiler.symm_mem import group_name_of + + assert group_name_of(_param(mesh_1rank)) == mesh_1rank.get_group("dp").group_name + with pytest.raises(RuntimeError, match="cannot resolve the process group"): + group_name_of(torch.empty(4, 4, device="cuda")) + + +# --------------------------------------------------------------------------- +# Selection: what the graph says gets bound, and nothing else +# --------------------------------------------------------------------------- +@requires_cuda +def test_only_gathered_weights_are_bound(mesh_1rank): + """The whole point of reading the graph: a weight the model owns but never + all-gathers gets no window. The patch this replaced took the entire subtree.""" + from magi_compiler.symm_mem import lookup_shard, registered_buffers + + gathered = [_param(mesh_1rank), _param(mesh_1rank)] + spare = _param(mesh_1rank) + gm, examples = _graph(gathered, extra_placeholder=spare) + + assert len(_bind(gm, examples)) == 2 + # Same dtype and group, so the two land in one pooled window. + (buffer,) = registered_buffers() + assert all(buffer.contains(p._local_tensor) for p in gathered) + assert lookup_shard(spare._local_tensor.data_ptr()) is None + + +@requires_cuda +def test_binding_keeps_the_parameter_object_and_its_values(mesh_1rank): + """Dynamo has already read this call's inputs and installed guards against + these exact objects, so the move has to be a pointer swap under them: a + replacement parameter would leave *this* invocation gathering the old, + unregistered storage.""" + from magi_compiler.symm_mem import lookup_shard + + param = _param(mesh_1rank) + local = param._local_tensor + before = local.clone() + + gm, examples = _graph([param]) + assert len(_bind(gm, examples)) == 1 + + assert param._local_tensor is local + assert torch.equal(local, before) + assert lookup_shard(local.data_ptr()) is not None + + +@requires_cuda +def test_binding_preserves_layout(mesh_1rank): + from magi_compiler.symm_mem import registered_buffers + + param = _param(mesh_1rank, shape=(16, 8)) + local = param._local_tensor + stride = local.stride() + + gm, examples = _graph([param]) + _bind(gm, examples) + + assert local.storage_offset() == 0 + assert local.stride() == stride + assert local.data_ptr() % 16 == 0 + assert registered_buffers()[0].contains(local) + + +@requires_cuda +def test_a_cast_before_the_gather_binds_nothing(mesh_1rank): + """Under mixed precision the gather reads the cast's output, which no window + backs. Binding the shard anyway would spend symmetric memory on a weight that + still goes over NCCL -- the failure mode of deciding this at build time.""" + from magi_compiler.symm_mem import registered_buffers + + gm, examples = _graph([_param(mesh_1rank), _param(mesh_1rank)], derive="cast") + + assert _bind(gm, examples) == set() + assert registered_buffers() == [] + + +@requires_cuda +def test_an_uneven_shard_binds_nothing(mesh_1rank): + """Read off the flag, never the graph shape: the pad appears only on the ranks + that own fewer rows, so a shape-derived answer splits one collective across two + transports and never completes.""" + from magi_compiler.symm_mem import registered_buffers + + gm, examples = _graph([_param(mesh_1rank)], uneven=True) + + assert _bind(gm, examples) == set() + assert registered_buffers() == [] + + +@requires_cuda +def test_a_replicated_weight_is_not_bound(mesh_1rank): + from torch.distributed.tensor import Replicate + + from magi_compiler.symm_mem import registered_buffers + + gm, examples = _graph([_param(mesh_1rank, placement=Replicate())]) + + assert _bind(gm, examples) == set() + assert registered_buffers() == [] + + +@requires_cuda +def test_a_meta_weight_is_not_bound(mesh_1rank): + """Ahead-of-time compilation hands the backend fake inputs. There is no + allocated weight to move, so binding declines and the graph stays on NCCL -- + those callers reach for ``bind_parameters`` instead.""" + from torch.distributed.tensor import DTensor, Shard + + from magi_compiler.symm_mem import registered_buffers + + fake = DTensor.from_local(torch.empty(8, 4, device="meta"), mesh_1rank, [Shard(0)], run_check=False) + gm, examples = _graph([fake]) + + assert _bind(gm, examples) == set() + assert registered_buffers() == [] + + +@requires_cuda +def test_an_unresolvable_graph_input_is_not_bound(mesh_1rank): + """No ``example_inputs`` means no live parameter behind the placeholder.""" + from magi_compiler.symm_mem import registered_buffers + + gm, _examples = _graph([_param(mesh_1rank)]) + + assert _bind(gm, {}) == set() + assert registered_buffers() == [] + + +@requires_cuda +def test_shards_below_the_size_floor_are_not_bound(mesh_1rank): + """Each shard costs an allocation and a rendezvous, and a small gather is + launch-bound anyway.""" + from magi_compiler.symm_mem import registered_buffers + + gm, examples = _graph([_param(mesh_1rank, shape=(8, 4))]) + + assert _bind(gm, examples, min_shard_bytes=1 << 20) == set() + assert registered_buffers() == [] + + +# --------------------------------------------------------------------------- +# Pooling +# --------------------------------------------------------------------------- +@requires_cuda +def test_many_shards_share_few_windows(mesh_1rank): + """Windows are a capped driver resource -- roughly 128 per process, however + small. A window per shard looks fine at this scale and dies at shard 92 of + gaga4's 808, in ``rendezvous``, with tens of GiB of the card still free. So + the invariant worth holding is a count that does not track the shard count.""" + from magi_compiler.symm_mem import registered_buffers + + params = [_param(mesh_1rank) for _ in range(16)] + gm, examples = _graph(params) + + assert len(_bind(gm, examples)) == 16 + # One dtype, one group -- and 16 shards nowhere near the size cap. + (buffer,) = registered_buffers() + assert all(buffer.contains(p._local_tensor) for p in params) + + +@requires_cuda +def test_a_window_holds_one_dtype_per_group(mesh_1rank): + """A window is a single ``symm_mem.empty``, so it has exactly one dtype; the + shards of a second dtype need a second one.""" + from magi_compiler.symm_mem import registered_buffers + + params = [_param(mesh_1rank), _param(mesh_1rank, dtype=torch.float32)] + gm, examples = _graph(params) + + assert len(_bind(gm, examples)) == 2 + buffers = registered_buffers() + assert len(buffers) == 2 + assert {b.dtype for b in buffers} == {torch.bfloat16, torch.float32} + + +@requires_cuda +def test_a_pooled_shard_still_starts_its_own_storage(mesh_1rank): + """Dynamo memoized ``storage_offset == 0`` for these parameters before the + backend ran, so a slot handed out as ``window[off:off+n]`` makes the next + fakification of the same source contradict the shape env -- an assert inside + ``create_symintnode``, nowhere near here.""" + params = [_param(mesh_1rank) for _ in range(3)] + gm, examples = _graph(params) + + assert len(_bind(gm, examples)) == 3 + assert [p._local_tensor.storage_offset() for p in params] == [0, 0, 0] + assert len({p._local_tensor.data_ptr() for p in params}) == 3 + + +@requires_cuda +def test_pooled_shards_do_not_overlap(mesh_1rank): + """Suballocation is only safe if the slots are disjoint: a bad offset would let + one weight silently overwrite the next, and the gather would still 'work'.""" + params = [_param(mesh_1rank) for _ in range(4)] + for i, p in enumerate(params): + p._local_tensor.fill_(i + 1) + gm, examples = _graph(params) + + assert len(_bind(gm, examples)) == 4 + for i, p in enumerate(params): + assert torch.equal(p._local_tensor, torch.full_like(p._local_tensor, i + 1)) + + +# --------------------------------------------------------------------------- +# Aliasing and repetition +# --------------------------------------------------------------------------- +@requires_cuda +def test_a_tied_weight_is_bound_once_and_serves_both_gathers(mesh_1rank): + """Two placeholders, one tensor. Allocating twice would leave the second + gather reading a window nobody wrote into.""" + from magi_compiler.symm_mem import registered_buffers + + param = _param(mesh_1rank) + gm, examples = _graph([param, param]) + + assert len(_bind(gm, examples)) == 2 + assert len(registered_buffers()) == 1 + + +@requires_cuda +def test_binding_twice_reuses_the_first_allocation(mesh_1rank): + """A model can be compiled more than once (a second entry point, a recompile). + A second allocation per rank would also desynchronize the rendezvous count.""" + from magi_compiler.symm_mem import registered_buffers + + param = _param(mesh_1rank) + gm, examples = _graph([param]) + + assert len(_bind(gm, examples)) == 1 + ptr = param._local_tensor.data_ptr() + + gm2, examples2 = _graph([param]) + assert len(_bind(gm2, examples2)) == 1 + assert len(registered_buffers()) == 1 + assert param._local_tensor.data_ptr() == ptr + + +# --------------------------------------------------------------------------- +# The explicit entry point, and the pipeline as a whole +# --------------------------------------------------------------------------- +@requires_cuda +def test_bind_parameters_moves_an_explicit_list(mesh_1rank): + """The escape hatch for callers with no graph: op-level benchmarks and AOT.""" + from magi_compiler.symm_mem import bind_parameters, lookup_shard + + params = [_param(mesh_1rank), _param(mesh_1rank)] + assert bind_parameters(params) == 2 + assert all(lookup_shard(p._local_tensor.data_ptr()) is not None for p in params) + + +@requires_cuda +def test_bind_parameters_skips_what_it_cannot_gather(mesh_1rank): + from torch.distributed.tensor import Replicate + + from magi_compiler.symm_mem import bind_parameters, registered_buffers + + assert bind_parameters([_param(mesh_1rank, placement=Replicate())]) == 0 + assert bind_parameters([torch.empty(8, 4, device="cuda")]) == 0 + assert registered_buffers() == [] + + +@requires_cuda +def test_bound_weights_are_bucketed_and_retargeted(mesh_1rank): + """The pipeline end to end: bind, then bucket only what bound, then retarget. + Binding before bucketing is what keeps a bucket homogeneous.""" + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED + + gm, examples = _graph([_param(mesh_1rank) for _ in range(4)]) + n = lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine", example_inputs=examples) + + assert n == 1 + coalesced = [x for x in gm.graph.nodes if x.target is CE_ALL_GATHER_COALESCED] + assert len(coalesced) == 1 + assert len(coalesced[0].args[0]) == 4 + + +@requires_cuda +def test_an_unbound_weight_keeps_its_neighbours_on_the_copy_engine(mesh_1rank): + """Only the weight that could not be bound loses the copy engine. Excluding it + must not split the bucket around it, or one odd weight would cost throughput + across the whole model.""" + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.passes.fsdp_overlap.node_meta import UNEVEN_SHARD + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED + + gm, examples = _graph([_param(mesh_1rank) for _ in range(4)]) + gathers = [n for n in gm.graph.nodes if n.target is _AG] + gathers[1].meta[UNEVEN_SHARD] = True + + lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine", example_inputs=examples) + + coalesced = [n for n in gm.graph.nodes if n.target is CE_ALL_GATHER_COALESCED] + assert len(coalesced) == 1 + assert len(coalesced[0].args[0]) == 3 # the three bound weights, in one bucket + assert len([n for n in gm.graph.nodes if n.target is _AG]) == 1 # the uneven one, on NCCL + + +@requires_cuda +def test_unbound_weights_are_still_bucketed_as_nccl(mesh_1rank): + """Losing the copy engine must not also lose bucketing. + + Bucketing partitions by transport rather than filtering on it: if unbound + gathers were simply dropped out of the pass, a model that declines binding + wholesale would go from a handful of coalesced launches to one launch per + weight, which costs far more memory than the copy engine ever saved. + """ + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.passes.fsdp_overlap.node_meta import UNEVEN_SHARD + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED + + _AG_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default + + gm, examples = _graph([_param(mesh_1rank) for _ in range(4)]) + for node in [n for n in gm.graph.nodes if n.target is _AG][2:]: + node.meta[UNEVEN_SHARD] = True # two weights the copy engine cannot serve + + assert ( + lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine", example_inputs=examples) + == 2 + ) + + (ce,) = [n for n in gm.graph.nodes if n.target is CE_ALL_GATHER_COALESCED] + (nccl,) = [n for n in gm.graph.nodes if n.target is _AG_COALESCED] + assert len(ce.args[0]) == 2 + assert len(nccl.args[0]) == 2 + assert not [n for n in gm.graph.nodes if n.target is _AG] + + +@requires_cuda +def test_nccl_transport_binds_nothing(mesh_1rank): + """Symmetric memory is a copy-engine cost; the default transport must not pay it.""" + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.symm_mem import registered_buffers + + gm, examples = _graph([_param(mesh_1rank) for _ in range(2)]) + lower_and_bucket_full_graph(gm, "coalesced", transport="nccl", example_inputs=examples) + + assert registered_buffers() == [] diff --git a/tests/feature_tests/fsdp/test_symm_e2e.py b/tests/feature_tests/fsdp/test_symm_e2e.py new file mode 100644 index 0000000..944aa91 --- /dev/null +++ b/tests/feature_tests/fsdp/test_symm_e2e.py @@ -0,0 +1,92 @@ +# 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. + +"""End-to-end guard for ``fsdp_config.transport="copy_engine"``. + +The unit tests each stub something out: the buffer tests materialize by calling +``_apply`` with a hand-forged lambda, and the gather tests run on one rank where +the peer view is the local shard. The property that only shows up when the +whole chain runs -- meta build, SimpleFSDP, ``to_empty``, checkpoint load, +compile, rewrite, reorder -- is that the weights are never ordinary tensors at +any point, and that is what the helper script asserts. It needs a real process +group and two ranks, so it runs as a ``torchrun`` subprocess and this file +asserts on its stdout markers. +""" + +import os +import shutil +import subprocess +from pathlib import Path + +import pytest +import torch + +_SCRIPT = Path(__file__).parent / "symm_helper" / "verify_symm_e2e.py" + +requires_2gpu = pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") +requires_torchrun = pytest.mark.skipif(shutil.which("torchrun") is None, reason="requires torchrun") + + +def _run(transport: str, port: str) -> subprocess.CompletedProcess: + env = os.environ.copy() + env["MAGI_LOGGING_LEVEL"] = env.get("MAGI_LOGGING_LEVEL", "info") + return subprocess.run( + [ + "torchrun", + "--nproc_per_node=2", + f"--master_port={port}", + str(_SCRIPT), + "--transport", + transport, + # Small enough to compile quickly, big enough that the gather is a + # real cross-rank transfer rather than a rounding error. + "--hidden", + "1024", + "--n-layers", + "4", + "--n-tokens", + "512", + ], + env=env, + capture_output=True, + text=True, + timeout=900, + ) + + +@requires_2gpu +@requires_torchrun +def test_copy_engine_end_to_end(): + """Weights land in symmetric memory, the gathers get retargeted, and two + consecutive steps both match an unsharded eager model -- the second step + is where a missing wait would show up.""" + p = _run("copy_engine", "29641") + out = p.stdout + p.stderr + assert p.returncode == 0, f"script failed:\n{out[-4000:]}" + assert "CHECK placement: 4/4 block shards" in p.stdout, out[-4000:] + assert "CHECK rewrite: 4/4 gathers" in p.stdout, out[-4000:] + assert "E2E_PASS" in p.stdout, out[-4000:] + + +@requires_2gpu +@requires_torchrun +def test_nccl_transport_is_untouched(): + """The default transport must allocate no window and rewrite no gather: the + control that says the copy-engine result above came from the flag.""" + p = _run("nccl", "29642") + out = p.stdout + p.stderr + assert p.returncode == 0, f"script failed:\n{out[-4000:]}" + assert "CHECK placement: 0/4 block shards in 0 window(s)" in p.stdout, out[-4000:] + assert "CHECK rewrite: 0/0 gathers" in p.stdout, out[-4000:] + assert "E2E_PASS" in p.stdout, out[-4000:] diff --git a/tests/feature_tests/fsdp/test_uneven_shard_transport.py b/tests/feature_tests/fsdp/test_uneven_shard_transport.py new file mode 100644 index 0000000..e1d6275 --- /dev/null +++ b/tests/feature_tests/fsdp/test_uneven_shard_transport.py @@ -0,0 +1,82 @@ +# 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. + +"""Guard: an uneven ``Shard(0)`` weight must get the same transport on every rank. + +A weight gather is a collective, so which transport carries it is a joint decision. +An uneven ``Shard(0)`` is where that decision used to split: the lowering pads the +shards of the ranks that own fewer rows, and copy-engine eligibility was decided from +the gather's input, which the pad changes. Only the trailing ranks refused the copy +engine; the rest moved on without them and their NCCL all-gather never completed. + +The divergence is invisible at one rank -- each rank's graph is individually +reasonable -- so the check has to compare graphs ACROSS ranks, which needs a real +process group and two of them. The helper runs under ``torchrun`` and all-gathers +its own decisions; this file asserts on its markers. +""" + +import os +import shutil +import subprocess +from pathlib import Path + +import pytest +import torch + +_SCRIPT = Path(__file__).parent / "fsdp_overlap_helper" / "uneven_shard_helper.py" + +requires_2gpu = pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") +requires_torchrun = pytest.mark.skipif(shutil.which("torchrun") is None, reason="requires torchrun") + + +def _run(port: str) -> subprocess.CompletedProcess: + env = os.environ.copy() + env["MAGI_LOGGING_LEVEL"] = env.get("MAGI_LOGGING_LEVEL", "info") + return subprocess.run( + ["torchrun", "--nproc_per_node=2", f"--master_port={port}", str(_SCRIPT)], + env=env, + capture_output=True, + text=True, + timeout=600, + ) + + +@requires_2gpu +@requires_torchrun +def test_uneven_shard_transport_is_rank_identical(): + """Both the uneven weight and the even control must be decided unanimously. + + The even case is the control that keeps the fix honest: refusing the copy engine + for everything would satisfy the uneven assertion and quietly cost the whole + feature, so ``rows=4`` has to still come out as one copy-engine bucket. + + The mixed graph interleaves the two: even weights stay on the copy engine, uneven + ones stay on NCCL, each as their own coalesced bucket, and every rank agrees. + """ + p = _run("29645") + out = p.stdout + p.stderr + assert p.returncode == 0, f"helper failed:\n{out[-4000:]}" + # even control: still bucketed onto the copy engine + assert "UNEVEN_TRANSPORT rows=4 agree=True targets={'ce_coalesced': 1}" in p.stdout, out[-4000:] + # uneven: every rank keeps it on NCCL -- but still buckets it there. Losing the + # copy engine must not also cost bucketing, or one odd weight turns N gathers + # into N launches. + assert "UNEVEN_TRANSPORT rows=3 agree=True targets={'nccl_coalesced': 1}" in p.stdout, out[-4000:] + assert "UNEVEN_NCCL_BUCKETS rows=3 agree=True" in p.stdout, out[-4000:] + assert "UNEVEN_SYMM agree=True in_buffer=['even']" in p.stdout, out[-4000:] + assert ( + "UNEVEN_MIXED agree=True targets={'ce_coalesced': 1, 'nccl_coalesced': 1} " + "sizes=[('ce_coalesced', 2), ('nccl_coalesced', 2)]" in p.stdout + ), out[-4000:] + assert "UNEVEN_PASS" in p.stdout, out[-4000:] diff --git a/tests/feature_tests/pass/test_python_scalar_constant_folding.py b/tests/feature_tests/pass/test_python_scalar_constant_folding.py index b0ec83a..775b39f 100644 --- a/tests/feature_tests/pass/test_python_scalar_constant_folding.py +++ b/tests/feature_tests/pass/test_python_scalar_constant_folding.py @@ -107,9 +107,9 @@ def test_recaptured_negated_float_constant_is_folded_before_split(monkeypatch): split_calls: list[torch.fx.GraphModule] = [] original_split = magi_backend_module.MagiBackend._split_graph - def capture_split(self, graph): + def capture_split(self, graph, example_inputs): split_calls.append(graph) - return original_split(self, graph) + return original_split(self, graph, example_inputs) monkeypatch.setattr(magi_backend_module.MagiBackend, "_split_graph", capture_split)