From 80e9af74578809fc88a2ab1037c8cef4402dd8d5 Mon Sep 17 00:00:00 2001 From: wtr Date: Thu, 27 Aug 2026 17:19:34 +0800 Subject: [PATCH 01/16] [Feat] Support Copy Engine Allgather in FSDP --- magi_compiler/_api.py | 15 + magi_compiler/config.py | 12 +- magi_compiler/magi_backend/magi_backend.py | 7 +- magi_compiler/passes/fsdp_overlap/__init__.py | 2 + .../passes/fsdp_overlap/bucket_all_gather.py | 4 +- .../passes/fsdp_overlap/lower_and_bucket.py | 26 +- magi_compiler/passes/fsdp_overlap/reorder.py | 111 ++++-- .../passes/fsdp_overlap/symm_ag_rewrite.py | 82 +++++ magi_compiler/profiling/runtime_estimator.py | 230 +++++++++--- magi_compiler/runtime/__init__.py | 13 + magi_compiler/runtime/symm_all_gather.py | 233 +++++++++++++ magi_compiler/runtime/symm_arena.py | 329 ++++++++++++++++++ .../fsdp_overlap_helper/reorder_helper.py | 73 +++- .../fsdp/test_fsdp_overlap_reorder.py | 19 + .../fsdp/test_profiling_estimator.py | 19 + 15 files changed, 1094 insertions(+), 81 deletions(-) create mode 100644 magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py create mode 100644 magi_compiler/runtime/__init__.py create mode 100644 magi_compiler/runtime/symm_all_gather.py create mode 100644 magi_compiler/runtime/symm_arena.py diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 033a020..ec3c93e 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -184,6 +184,11 @@ def _lazy_init_magi_state( if getattr(state_holder, state_attr, None) is not None: return + if conf.fsdp_config.transport == "copy_engine": + from magi_compiler.runtime.symm_arena import barrier_after_load + + barrier_after_load() + compilation_counter.num_models_seen += 1 setattr( @@ -218,6 +223,11 @@ def _magi_compile_class( if issubclass(cls, nn.Module) and conf.offload_config.model_cpu_offload: _patch_cpu_offload_apply(cls, conf) + if issubclass(cls, nn.Module) and conf.fsdp_config.transport == "copy_engine": + from magi_compiler.runtime.symm_arena import patch_symm_arena_apply + + patch_symm_arena_apply(cls) + old_init = cls.__init__ @functools.wraps(old_init) @@ -241,6 +251,11 @@ def _magi_compile_bound_method( if getattr(instance, installed_attr, False): return instance + if conf.fsdp_config.transport == "copy_engine" and isinstance(instance, nn.Module): + from magi_compiler.runtime.symm_arena import migrate_to_arenas + + migrate_to_arenas(instance) + old_method = getattr(instance, method_name) @torch.compiler.disable() diff --git a/magi_compiler/config.py b/magi_compiler/config.py index f833cdf..c27bff6 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,16 @@ 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=( + "How weight all-gathers move bytes. 'nccl': ring kernels on the SMs. " + "'copy_engine': weight shards are allocated in symmetric memory at model build time and " + "gathered by peer copy-engine reads -- zero SM occupancy and no per-step cross-rank barrier, " + "at a lower raw bandwidth. Requires all ranks of the FSDP mesh dim to be NVLink-connected " + "within one node, and static weights (inference)." + ), + ) 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..1beb364 100644 --- a/magi_compiler/magi_backend/magi_backend.py +++ b/magi_compiler/magi_backend/magi_backend.py @@ -615,9 +615,12 @@ 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 + ) 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, diff --git a/magi_compiler/passes/fsdp_overlap/__init__.py b/magi_compiler/passes/fsdp_overlap/__init__.py index d507a8a..49b1ff1 100644 --- a/magi_compiler/passes/fsdp_overlap/__init__.py +++ b/magi_compiler/passes/fsdp_overlap/__init__.py @@ -16,10 +16,12 @@ from .lower_and_bucket import lower_and_bucket_full_graph from .redistribute_lowering import lower_prim_redistribute_to_collectives from .reorder import FsdpOverlapReorder +from .symm_ag_rewrite import rewrite_weight_ag_to_copy_engine __all__ = [ "bucket_weight_all_gather_coalesced", "lower_prim_redistribute_to_collectives", "lower_and_bucket_full_graph", + "rewrite_weight_ag_to_copy_engine", "FsdpOverlapReorder", ] diff --git a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py index c4ef6f6..24a6dab 100644 --- a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py +++ b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py @@ -177,7 +177,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, eligible=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`` @@ -204,6 +204,8 @@ def bucket_weight_all_gather_coalesced(graph: fx.GraphModule, bucket_size_bytes: for node in graph.graph.nodes: if not _is_weight_all_gather(node): continue + if eligible is not None and not eligible(node): + continue _, _world, group_name = node.args groups[group_name].append(node) diff --git a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py index 34123dc..b9938f6 100644 --- a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py +++ b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py @@ -20,9 +20,12 @@ from .bucket_all_gather import bucket_weight_all_gather_coalesced from .redistribute_lowering import lower_prim_redistribute_to_collectives +from .symm_ag_rewrite import rewrite_weight_ag_to_copy_engine -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" +) -> int: """Lower SimpleFSDP weight redistribute -> explicit collectives, then optionally bucket them across the WHOLE graph (no subgraph partitioning). @@ -36,19 +39,28 @@ 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"`` buckets *first* (only arena-shard gathers, so + cast/pad stays out of the bucket), then retargets both the leftover singles + and the coalesced launches at the copy-engine ops. The wrapper still runs + one gather per member; reorder just sees one comm node per bucket. + 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 - + n = 0 if bucket_mode == "coalesced": - n = bucket_weight_all_gather_coalesced(graph, bucket_size_bytes=bucket_size_bytes) - else: + from .symm_ag_rewrite import _input_is_arena_shard + + eligible = _input_is_arena_shard if transport == "copy_engine" else None + n = bucket_weight_all_gather_coalesced(graph, bucket_size_bytes=bucket_size_bytes, eligible=eligible) + 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/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index 3c80ab1..2f4e607 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -51,7 +51,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 _symm_ag_ops(): + """Copy-engine gather ops, imported lazily so this pass stays importable + without a CUDA build.""" + try: + from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + + return tuple(op for op in (SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED) if op is not None) + except Exception: # noqa: BLE001 + return () + + +_SYMM_AG_OPS = _symm_ag_ops() +_WEIGHT_AG_OPS = tuple(op for op in (_AG, _AG_COALESCED, *_SYMM_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 +73,43 @@ _DEFAULT_WINDOW_MARGIN_NS = 5_000.0 +def _is_symm_ag_ir(node) -> bool: + """ + ``magi::symm_all_gather`` lowers to an ordinary FallbackKernel, so + Inductor's ``is_collective`` does not recognize it. + """ + return getattr(node, "op_overload", None) in _SYMM_AG_OPS + + +def _is_symm_ag_coalesced(node) -> bool: + try: + from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER_COALESCED + except Exception: # noqa: BLE001 + return False + return SYMM_ALL_GATHER_COALESCED is not None and getattr(node, "op_overload", None) is SYMM_ALL_GATHER_COALESCED + + +def _is_gather_ir(node) -> bool: + return node is not None and (is_collective(node) or _is_symm_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 +145,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 +251,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 +286,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 @@ -311,8 +350,10 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: # 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 | " + "FSDP overlap placement: launch %s(%s) cur=%d -> target=%d fc=%d lower=%d | " "comm=%.1fus acc_upstream=%.1fus need=%.1fus %s", + launch.get_name(), + getattr(_leaf_collective_node(launch), "op_overload", "?"), cur, target, fc_idx, @@ -476,7 +517,11 @@ 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) + # The upper bound keeps the gather below the next collective, but it + # must never win against the floor: with no collective above (or, on a + # pure copy-engine graph, no skeleton at all) it degenerates to the + # launch's current index and would silently undo a legal hoist. + 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 +545,53 @@ 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_symm_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_symm_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_symm_ag_ir(node): + # A FallbackKernel's result is re-exposed through an alias snode + # (``buf1 = buf0`` in the generated code), which is what the wait and + # the consumer actually read. Inductor's own collectives have no such + # layer, so this is the one structural difference the copy-engine + # transport introduces -- and it has to move with the launch. + 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/passes/fsdp_overlap/symm_ag_rewrite.py b/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py new file mode 100644 index 0000000..6d93ecc --- /dev/null +++ b/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.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. + +from __future__ import annotations + +import torch +import torch.fx as fx + +from magi_compiler.utils import magi_logger + +_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default +_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default + + +def _input_is_arena_shard(node: fx.Node) -> bool: + """True when the gather input is ``to_local(placeholder/get_attr)`` with no intervening op.""" + src = node.args[0] if node.args else None + if not isinstance(src, fx.Node) or src.op != "call_method" or src.target != "to_local": + return False + owner = src.args[0] if src.args else None + return isinstance(owner, fx.Node) and owner.op in ("placeholder", "get_attr") + + +def _coalesced_inputs_are_arena_shards(node: fx.Node) -> bool: + locs = node.args[0] if node.args else None + if not isinstance(locs, (list, tuple)) or not locs: + return False + return all(isinstance(loc, fx.Node) and _input_is_arena_shard_from_local(loc) for loc in locs) + + +def _input_is_arena_shard_from_local(src: fx.Node) -> bool: + if src.op != "call_method" or src.target != "to_local": + return False + owner = src.args[0] if src.args else None + return isinstance(owner, fx.Node) and owner.op in ("placeholder", "get_attr") + + +def rewrite_weight_ag_to_copy_engine(graph: fx.GraphModule) -> int: + """Retarget marked weight gathers to ``magi::symm_all_gather``. Returns count rewritten.""" + from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + + rewritten = 0 + skipped = 0 + for node in graph.graph.nodes: + if node.op != "call_function": + continue + if not node.meta.get("magi_fsdp_weight_ag"): + continue + if node.target is _ALL_GATHER: + if not _input_is_arena_shard(node): + skipped += 1 + continue + node.target = SYMM_ALL_GATHER + rewritten += 1 + elif node.target is _ALL_GATHER_COALESCED: + if not _coalesced_inputs_are_arena_shards(node): + skipped += 1 + continue + node.target = SYMM_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 (input is a cast/pad of the shard, not the shard itself)", + rewritten, + skipped, + ) + return rewritten diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index db3246a..e2ed9a0 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -216,19 +216,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.""" + hook (``materialize_inputs``) that rebuilds value-consistent metadata. + """ fx_node = snode.node.get_origin_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 +244,43 @@ 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: + return 0.0 + 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 +343,117 @@ def _collective_spec(node): return op, group_name, group_size, specs +def _symm_ag_ops(): + """Copy-engine gather ops, or empty when the runtime is unavailable.""" + try: + from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + + return tuple(op for op in (SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED) if op is not None) + except Exception: # noqa: BLE001 + return () + + +def _leaf_symm_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 = _symm_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_symm_ag_coalesced_ir(node) -> bool: + try: + from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER_COALESCED + except Exception: # noqa: BLE001 + return False + return SYMM_ALL_GATHER_COALESCED is not None and getattr(node, "op_overload", None) is SYMM_ALL_GATHER_COALESCED + + +def _symm_ag_spec(node): + """(shapes, dtype, group_size, group_name) for a copy-engine gather. + + ``shapes`` is a tuple of local-shard shapes (one member, or one per + coalesced input). Read off ``constant_args``, the way ``_collective_spec`` + reads a group name. Not off ``get_origin_node()``: Inductor leaves that + unset on these nodes, and the resulting ``None`` silently degraded the + gather to a zero cost. + """ + 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 _symm_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.runtime.symm_arena import find_shard_by_layout + + node = _leaf_symm_ag(snode) + spec = _symm_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.runtime.symm_all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + + if _is_symm_ag_coalesced_ir(node): + op = SYMM_ALL_GATHER_COALESCED + return (lambda: op(shards, group_size, group_name)), lambda outs: [_WAIT(o) for o in outs] + return (lambda: SYMM_ALL_GATHER(shards[0], group_size, group_name)), _WAIT + + +def _measure_symm_ag(snode: BaseSchedulerNode) -> float: + """Time the copy-engine gather TOGETHER WITH its wait. + + Timing the launch alone would measure the CPU issue cost and nothing else: + the copies run on a side stream, so without the wait the timing events on the + current stream close before a single byte has moved. That reads as ~3us for + a gather that really takes tens of microseconds, and the reorder pass then + sizes a window an order of magnitude too small. + """ + pair = _symm_ag_launch_wait(snode) + if pair is None: + return 0.0 + launch, wait = pair + return _time_fixed(lambda: wait(launch())) + + +def _symm_ag_label(snode: BaseSchedulerNode) -> str: + node = _leaf_symm_ag(snode) + spec = _symm_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"symm_all_gather_coalesced(ws={group_size},n={len(shapes)},{shape0})" + return f"symm_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 +465,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: @@ -508,6 +619,8 @@ def _measure_one(self, snode: BaseSchedulerNode) -> float: """Lockstep-safe single measurement (fixed iters for anything containing a collective); never raises -- falls back to the analytical estimate.""" try: + if _leaf_symm_ag(snode) is not None: + return _measure_symm_ag(snode) if contains_collective(snode): return _measure_collective_op(snode) if isinstance(snode, ExternKernelSchedulerNode): @@ -548,6 +661,29 @@ def __call__(self, snode: BaseSchedulerNode) -> float: if _is_multi_output_unpack(snode): return 0.0 + if _leaf_symm_ag(snode) is not None: + node = _leaf_symm_ag(snode) + spec = _symm_ag_spec(node) + if spec is None: + return _safe_analytical(snode) + shapes, dtype, group_size, _gn = spec + ckey = ("symm_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="symm_ag", label=_symm_ag_label(snode), measured=False) + if self._sync_across_ranks: + self._key_snode[ckey] = snode + else: + ns = _measure_symm_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/runtime/__init__.py b/magi_compiler/runtime/__init__.py new file mode 100644 index 0000000..3eaa44a --- /dev/null +++ b/magi_compiler/runtime/__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/magi_compiler/runtime/symm_all_gather.py b/magi_compiler/runtime/symm_all_gather.py new file mode 100644 index 0000000..6d969c8 --- /dev/null +++ b/magi_compiler/runtime/symm_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_arena 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 arena. +_SCHEMA = "symm_all_gather(Tensor local, int group_size, str group_name) -> Tensor" +_SCHEMA_COALESCED = "symm_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::symm_all_gather got a tensor that is not a registered symmetric-memory shard. " + "Only weights materialized through the arena 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 _symm_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 _symm_all_gather_meta(local: torch.Tensor, group_size: int, group_name: str) -> torch.Tensor: + return _gather_dest(local, group_size) + + +def _symm_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 _symm_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("symm_all_gather", _symm_all_gather, "CUDA") + _LIB.impl("symm_all_gather", _symm_all_gather_meta, "Meta") + + _LIB.define(_SCHEMA_COALESCED) + _LIB.impl("symm_all_gather_coalesced", _symm_all_gather_coalesced, "CUDA") + _LIB.impl("symm_all_gather_coalesced", _symm_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. +SYMM_ALL_GATHER = torch.ops.magi.symm_all_gather.default +SYMM_ALL_GATHER_COALESCED = torch.ops.magi.symm_all_gather_coalesced.default diff --git a/magi_compiler/runtime/symm_arena.py b/magi_compiler/runtime/symm_arena.py new file mode 100644 index 0000000..a8266cb --- /dev/null +++ b/magi_compiler/runtime/symm_arena.py @@ -0,0 +1,329 @@ +# 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 inspect +from dataclasses import dataclass + +import torch +import torch.distributed as dist +import torch.nn as nn + +from magi_compiler.utils import magi_logger + +# Only hijack ``to_empty``'s lambda. ``.cuda()`` / ``.to()`` / ``.float()`` also +# go through ``_apply``; intercepting those would change builder semantics. +_TO_EMPTY_LAMBDA = "Module.to_empty.." + + +class SymmArena: + """One symmetric-memory window, suballocated to many weight shards. + + One window per (decorated block, dtype, process group): a single rendezvous, + and because every rank walks the module tree in the same order, offset ``k`` + is the same shard on every peer. Two meshes sharing a dtype get two windows. + """ + + # 256 bf16 elems = 512B, which the copy engine wants for peak throughput. + ALIGN = 256 + + def __init__(self, dtype: torch.dtype, device: torch.device, group_name: str) -> None: + self.dtype = dtype + self.device = device + self.group_name = group_name + self.buf: torch.Tensor | None = None + self.handle = None + self.peers: list[torch.Tensor] = [] + self._reserved = 0 + self._cursor = 0 + + # -- build phase ------------------------------------------------------ + def reserve(self, numel: int) -> None: + self._reserved += self._round(numel) + + def commit(self) -> None: + import torch.distributed._symmetric_memory as symm_mem + + symm_mem.enable_symm_mem_for_group(self.group_name) + self.buf = symm_mem.empty(self._reserved, dtype=self.dtype, device=self.device) + self.handle = symm_mem.rendezvous(self.buf, self.group_name) + # Slice each peer's whole window once. VMM maps every window into one + # VA space so these are directly copyable; NCCL-backend peer pointers are not. + self.peers = [self.handle.get_buffer(r, (self._reserved,), self.dtype) for r in range(self.handle.world_size)] + + def take(self, shape: torch.Size | tuple[int, ...]) -> torch.Tensor: + numel = 1 + for s in shape: + numel *= int(s) + off = self._cursor + self._cursor += self._round(numel) + if self._cursor > self._reserved: + raise RuntimeError( + f"symmetric arena overflow: wanted {self._cursor} elems, reserved {self._reserved}. " + "The sizing walk and the dispensing walk must visit the same shards in the same order." + ) + return self.buf[off : off + numel].view(shape) + + # -- query ------------------------------------------------------------ + @property + def nbytes(self) -> int: + return self._reserved * self.dtype.itemsize + + def offset_of(self, t: torch.Tensor) -> int: + return (t.data_ptr() - self.buf.data_ptr()) // self.buf.element_size() + + def peer_views(self, t: torch.Tensor) -> list[torch.Tensor]: + """``world_size`` views of the same shard, one per rank, in rank order.""" + off, numel = self.offset_of(t), t.numel() + return [p[off : off + numel].view(t.shape) for p in self.peers] + + def contains(self, t: torch.Tensor) -> bool: + if self.buf is None: + return False + base = self.buf.data_ptr() + return base <= t.data_ptr() < base + self.nbytes + + @classmethod + def _round(cls, numel: int) -> int: + return (numel + cls.ALIGN - 1) // cls.ALIGN * cls.ALIGN + + +@dataclass(frozen=True) +class ShardEntry: + """What the run-time gather needs to know about one local shard.""" + + arena: SymmArena + offset: int + local: torch.Tensor + peer_views: tuple[torch.Tensor, ...] + + @property + def shape(self) -> tuple[int, ...]: + return tuple(self.local.shape) + + +# Keyed by ``data_ptr()``: the gather op only sees a plain tensor. +_SHARD_REGISTRY: dict[int, ShardEntry] = {} +_ARENAS: list[SymmArena] = [] +_BARRIER_DONE = False + + +def register_shard(local: torch.Tensor, arena: SymmArena) -> ShardEntry: + entry = ShardEntry(arena=arena, offset=arena.offset_of(local), local=local, peer_views=tuple(arena.peer_views(local))) + _SHARD_REGISTRY[local.data_ptr()] = entry + return entry + + +def lookup_shard(data_ptr: int) -> ShardEntry | None: + return _SHARD_REGISTRY.get(data_ptr) + + +def registered_arenas() -> list[SymmArena]: + return list(_ARENAS) + + +def find_shard_by_layout(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor | None: + """Any registered shard with this layout -- the cost model cannot replay a gather on a generic ``empty`` (no peers).""" + want = tuple(int(s) for s in shape) + for entry in _SHARD_REGISTRY.values(): + if entry.shape == want and entry.local.dtype == dtype: + return entry.local + return None + + +def reset_registry() -> None: + """Test-only: drop every arena so a new model can be built in-process.""" + global _BARRIER_DONE + _SHARD_REGISTRY.clear() + _ARENAS.clear() + _BARRIER_DONE = False + + +def barrier_after_load() -> None: + """Publish every rank's freshly written shards, once per process. + + Must run after weights are loaded and before the first peer read. + """ + global _BARRIER_DONE + if _BARRIER_DONE or not _ARENAS: + return + if dist.is_available() and dist.is_initialized(): + torch.cuda.synchronize() + dist.barrier() + _BARRIER_DONE = True + magi_logger.info( + "Symmetric arena: published %d arena(s), %.1f MiB, %d shards; steady state is barrier-free", + len(_ARENAS), + sum(a.nbytes for a in _ARENAS) / 2**20, + len(_SHARD_REGISTRY), + ) + + +def _is_gatherable_shard(t: object) -> bool: + """A Shard(0) DTensor on a 1-D mesh -- the only placement the copy-engine gather handles.""" + from torch.distributed.tensor import DTensor, Shard + + if not isinstance(t, DTensor): + return False + placements = t.placements + return len(placements) == 1 and isinstance(placements[0], Shard) and placements[0].dim == 0 + + +def _group_name_of(t) -> str | None: + try: + return t.device_mesh._dim_group_names[0] + except Exception: # noqa: BLE001 + return None + + +def _arena_key(t) -> tuple[torch.dtype, str]: + """One window per (dtype, process group). Same dtype on two meshes (gaga4 FSDP + edp) must not share a window.""" + group_name = _group_name_of(t) + if group_name is None: + raise RuntimeError(f"cannot resolve the process group of a Shard(0) parameter on mesh {t.device_mesh}") + return (t.dtype, group_name) + + +def _apply_order_entries(mod: nn.Module): + """``(owner, name, param)`` in ``_apply`` order: post-order, every ``_parameters`` entry. + + Not ``named_parameters()`` (pre-order, dedups shared tensors). A different + walk would break cross-rank offset symmetry. Walking ``_parameters`` also + finds SimpleFSDP weights in ``parametrizations.weight.original``. + """ + for child in mod.children(): + yield from _apply_order_entries(child) + for name, p in mod._parameters.items(): + if p is not None: + yield mod, name, p + + +def _plan_arenas(shards: list, device: torch.device) -> dict[tuple[torch.dtype, str], SymmArena]: + """Size and commit one window per (dtype, group). Dedup by identity so a tied weight reserves a single slot.""" + arenas: dict[tuple[torch.dtype, str], SymmArena] = {} + seen: set[int] = set() + for p in shards: + if id(p) in seen: + continue + seen.add(id(p)) + key = _arena_key(p) + arena = arenas.get(key) + if arena is None: + arena = arenas[key] = SymmArena(p.dtype, device, key[1]) + arena.reserve(p._local_tensor.numel()) + + for arena in arenas.values(): + arena.commit() # the only collective, once per window + _ARENAS.extend(arenas.values()) + return arenas + + +def materialize_into_arenas(mod: nn.Module, device: torch.device) -> dict[tuple[torch.dtype, str], SymmArena]: + """Size windows for ``mod``'s Shard(0) shards while they are still on meta. Non-gatherable params are left to the caller.""" + shards = [p for _, _, p in _apply_order_entries(mod) if _is_gatherable_shard(p)] + if not shards: + return {} + return _plan_arenas(shards, device) + + +def migrate_to_arenas(root: nn.Module) -> dict[tuple[torch.dtype, str], SymmArena]: + """Copy already-allocated Shard(0) shards into symmetric memory. + + Used when ``magi_compile(model, ...)`` is given a live model rather than a + meta + ``to_empty`` path. ``load_state_dict(assign=True)`` after this would + replace arena views with ordinary tensors; the gather then rejects them. + """ + entries = [(m, n, p) for m, n, p in _apply_order_entries(root) if _is_gatherable_shard(p)] + if not entries: + return {} + + device = entries[0][2]._local_tensor.device + if device.type != "cuda": + raise RuntimeError(f"symmetric memory needs the shards on cuda, found {device}") + arenas = _plan_arenas([p for _, _, p in entries], device) + + from torch.distributed.tensor import DTensor + + views: dict[int, torch.Tensor] = {} + for owner, name, p in entries: + local = views.get(id(p)) + if local is None: + arena = arenas[_arena_key(p)] + local = views[id(p)] = arena.take(p._local_tensor.shape) + local.copy_(p._local_tensor) + register_shard(local, arena) + moved = DTensor.from_local(local, p.device_mesh, p.placements, run_check=False) + owner.register_parameter(name, nn.Parameter(moved, requires_grad=p.requires_grad)) + + magi_logger.info( + "Symmetric arena: migrated %d shard(s) into %.1f MiB across %d window(s)", + len(views), + sum(a.nbytes for a in arenas.values()) / 2**20, + len(arenas), + ) + return arenas + + +def patch_symm_arena_apply(cls: type[nn.Module]) -> None: + """Install the ``_apply`` interception on a decorated class. + + Mirrors ``_patch_cpu_offload_apply``: take over for ``to_empty``'s lambda, delegate everything else. + """ + if getattr(cls, "_magi_symm_apply_patched", False): + return + orig_apply = cls._apply + magi_logger.info("Symmetric arena: intercepting %s._apply for copy-engine FSDP", cls.__name__) + + def _symm_apply(self, fn, recurse: bool = True): + if getattr(fn, "__qualname__", "") != _TO_EMPTY_LAMBDA: + return orig_apply(self, fn, recurse) + if getattr(self, "_magi_symm_arenas", None) is not None: + return orig_apply(self, fn, recurse) + + device = torch.device(inspect.getclosurevars(fn).nonlocals["device"]) + from torch.distributed.tensor import DTensor + + arenas = materialize_into_arenas(self, device) + if not arenas: + return orig_apply(self, fn, recurse) + + views: dict[int, torch.Tensor] = {} + + def materialize(t: torch.Tensor) -> torch.Tensor: + if not _is_gatherable_shard(t): + return torch.empty_like(t, device=device) + # Tied weight: same view so tying survives materialization. + local = views.get(id(t)) + if local is None: + arena = arenas[_arena_key(t)] + local = views[id(t)] = arena.take(t._local_tensor.shape) + register_shard(local, arena) + return DTensor.from_local(local, t.device_mesh, t.placements, run_check=False) + + # Do not forge to_empty's qualname: a nested decorated block must fail + # the check above and delegate, so its params land in *this* arena. + out = orig_apply(self, materialize, recurse) + self._magi_symm_arenas = arenas + magi_logger.info( + "Symmetric arena: %s materialized %d shard(s) into %.1f MiB across %d window(s)", + cls.__name__, + len(views), + sum(a.nbytes for a in arenas.values()) / 2**20, + len(arenas), + ) + return out + + cls._apply = _symm_apply + cls._magi_symm_apply_patched = True 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..2198a6e 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::symm_all_gather`` reading a symmetric arena. 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 """ @@ -110,6 +121,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 arena instead of NCCL") args = ap.parse_args() if args.modes_only: @@ -138,6 +150,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 +165,47 @@ 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.runtime.symm_all_gather import SYMM_ALL_GATHER + from magi_compiler.runtime.symm_arena import SymmArena, register_shard + + arena = SymmArena(torch.bfloat16, torch.device("cuda", dev), grp) + for _ in range(N_CE_LAYERS): + arena.reserve(H * H) + arena.commit() + for i in range(N_CE_LAYERS): + s = arena.take((H, H)) + s.normal_(0.0, H**-0.5).add_(0.01 * i) + register_shard(s, arena) + ce_shards.append(s) + # A peer read is only legal once that peer has written its shard. + 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(SYMM_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 +238,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_symm_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 +257,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 +286,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/test_fsdp_overlap_reorder.py b/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py index 39388d4..61b7c57 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 arena. + + 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..b603ac6 100644 --- a/tests/feature_tests/fsdp/test_profiling_estimator.py +++ b/tests/feature_tests/fsdp/test_profiling_estimator.py @@ -293,6 +293,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 # =========================================================================== From 6290ee7369d996be359e55a1fe20f6a0b13fd9e5 Mon Sep 17 00:00:00 2001 From: wtr Date: Thu, 27 Aug 2026 20:03:30 +0800 Subject: [PATCH 02/16] [Refactor] refactor symm_mempory code --- magi_compiler/_api.py | 6 +- magi_compiler/passes/fsdp_overlap/reorder.py | 4 +- .../passes/fsdp_overlap/symm_ag_rewrite.py | 2 +- magi_compiler/profiling/runtime_estimator.py | 8 +- magi_compiler/runtime/__init__.py | 13 - magi_compiler/runtime/symm_all_gather.py | 233 ------------- magi_compiler/runtime/symm_arena.py | 329 ------------------ .../fsdp_overlap_helper/reorder_helper.py | 4 +- 8 files changed, 12 insertions(+), 587 deletions(-) delete mode 100644 magi_compiler/runtime/__init__.py delete mode 100644 magi_compiler/runtime/symm_all_gather.py delete mode 100644 magi_compiler/runtime/symm_arena.py diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index ec3c93e..2018c2b 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -185,7 +185,7 @@ def _lazy_init_magi_state( return if conf.fsdp_config.transport == "copy_engine": - from magi_compiler.runtime.symm_arena import barrier_after_load + from magi_compiler.symm_mem import barrier_after_load barrier_after_load() @@ -224,7 +224,7 @@ def _magi_compile_class( _patch_cpu_offload_apply(cls, conf) if issubclass(cls, nn.Module) and conf.fsdp_config.transport == "copy_engine": - from magi_compiler.runtime.symm_arena import patch_symm_arena_apply + from magi_compiler.symm_mem import patch_symm_arena_apply patch_symm_arena_apply(cls) @@ -252,7 +252,7 @@ def _magi_compile_bound_method( return instance if conf.fsdp_config.transport == "copy_engine" and isinstance(instance, nn.Module): - from magi_compiler.runtime.symm_arena import migrate_to_arenas + from magi_compiler.symm_mem import migrate_to_arenas migrate_to_arenas(instance) diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index 2f4e607..da4a7df 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -57,7 +57,7 @@ def _symm_ag_ops(): """Copy-engine gather ops, imported lazily so this pass stays importable without a CUDA build.""" try: - from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED return tuple(op for op in (SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED) if op is not None) except Exception: # noqa: BLE001 @@ -83,7 +83,7 @@ def _is_symm_ag_ir(node) -> bool: def _is_symm_ag_coalesced(node) -> bool: try: - from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED except Exception: # noqa: BLE001 return False return SYMM_ALL_GATHER_COALESCED is not None and getattr(node, "op_overload", None) is SYMM_ALL_GATHER_COALESCED diff --git a/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py b/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py index 6d93ecc..1871cfb 100644 --- a/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py +++ b/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py @@ -48,7 +48,7 @@ def _input_is_arena_shard_from_local(src: fx.Node) -> bool: def rewrite_weight_ag_to_copy_engine(graph: fx.GraphModule) -> int: """Retarget marked weight gathers to ``magi::symm_all_gather``. Returns count rewritten.""" - from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED rewritten = 0 skipped = 0 diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index e2ed9a0..c8248e3 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -346,7 +346,7 @@ def _collective_spec(node): def _symm_ag_ops(): """Copy-engine gather ops, or empty when the runtime is unavailable.""" try: - from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED return tuple(op for op in (SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED) if op is not None) except Exception: # noqa: BLE001 @@ -370,7 +370,7 @@ def _leaf_symm_ag(snode: BaseSchedulerNode): def _is_symm_ag_coalesced_ir(node) -> bool: try: - from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED except Exception: # noqa: BLE001 return False return SYMM_ALL_GATHER_COALESCED is not None and getattr(node, "op_overload", None) is SYMM_ALL_GATHER_COALESCED @@ -402,7 +402,7 @@ def _symm_ag_launch_wait(snode: BaseSchedulerNode): Split in two rather than one fused closure so the cost model can time ``wait(launch())`` as a unit. """ - from magi_compiler.runtime.symm_arena import find_shard_by_layout + from magi_compiler.symm_mem import find_shard_by_layout node = _leaf_symm_ag(snode) spec = _symm_ag_spec(node) if node is not None else None @@ -418,7 +418,7 @@ def _symm_ag_launch_wait(snode: BaseSchedulerNode): dtype, ) return None - from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED if _is_symm_ag_coalesced_ir(node): op = SYMM_ALL_GATHER_COALESCED diff --git a/magi_compiler/runtime/__init__.py b/magi_compiler/runtime/__init__.py deleted file mode 100644 index 3eaa44a..0000000 --- a/magi_compiler/runtime/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) 2026 SandAI. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. diff --git a/magi_compiler/runtime/symm_all_gather.py b/magi_compiler/runtime/symm_all_gather.py deleted file mode 100644 index 6d969c8..0000000 --- a/magi_compiler/runtime/symm_all_gather.py +++ /dev/null @@ -1,233 +0,0 @@ -# Copyright (c) 2026 SandAI. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -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_arena 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 arena. -_SCHEMA = "symm_all_gather(Tensor local, int group_size, str group_name) -> Tensor" -_SCHEMA_COALESCED = "symm_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::symm_all_gather got a tensor that is not a registered symmetric-memory shard. " - "Only weights materialized through the arena 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 _symm_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 _symm_all_gather_meta(local: torch.Tensor, group_size: int, group_name: str) -> torch.Tensor: - return _gather_dest(local, group_size) - - -def _symm_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 _symm_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("symm_all_gather", _symm_all_gather, "CUDA") - _LIB.impl("symm_all_gather", _symm_all_gather_meta, "Meta") - - _LIB.define(_SCHEMA_COALESCED) - _LIB.impl("symm_all_gather_coalesced", _symm_all_gather_coalesced, "CUDA") - _LIB.impl("symm_all_gather_coalesced", _symm_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. -SYMM_ALL_GATHER = torch.ops.magi.symm_all_gather.default -SYMM_ALL_GATHER_COALESCED = torch.ops.magi.symm_all_gather_coalesced.default diff --git a/magi_compiler/runtime/symm_arena.py b/magi_compiler/runtime/symm_arena.py deleted file mode 100644 index a8266cb..0000000 --- a/magi_compiler/runtime/symm_arena.py +++ /dev/null @@ -1,329 +0,0 @@ -# Copyright (c) 2026 SandAI. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from __future__ import annotations - -import inspect -from dataclasses import dataclass - -import torch -import torch.distributed as dist -import torch.nn as nn - -from magi_compiler.utils import magi_logger - -# Only hijack ``to_empty``'s lambda. ``.cuda()`` / ``.to()`` / ``.float()`` also -# go through ``_apply``; intercepting those would change builder semantics. -_TO_EMPTY_LAMBDA = "Module.to_empty.." - - -class SymmArena: - """One symmetric-memory window, suballocated to many weight shards. - - One window per (decorated block, dtype, process group): a single rendezvous, - and because every rank walks the module tree in the same order, offset ``k`` - is the same shard on every peer. Two meshes sharing a dtype get two windows. - """ - - # 256 bf16 elems = 512B, which the copy engine wants for peak throughput. - ALIGN = 256 - - def __init__(self, dtype: torch.dtype, device: torch.device, group_name: str) -> None: - self.dtype = dtype - self.device = device - self.group_name = group_name - self.buf: torch.Tensor | None = None - self.handle = None - self.peers: list[torch.Tensor] = [] - self._reserved = 0 - self._cursor = 0 - - # -- build phase ------------------------------------------------------ - def reserve(self, numel: int) -> None: - self._reserved += self._round(numel) - - def commit(self) -> None: - import torch.distributed._symmetric_memory as symm_mem - - symm_mem.enable_symm_mem_for_group(self.group_name) - self.buf = symm_mem.empty(self._reserved, dtype=self.dtype, device=self.device) - self.handle = symm_mem.rendezvous(self.buf, self.group_name) - # Slice each peer's whole window once. VMM maps every window into one - # VA space so these are directly copyable; NCCL-backend peer pointers are not. - self.peers = [self.handle.get_buffer(r, (self._reserved,), self.dtype) for r in range(self.handle.world_size)] - - def take(self, shape: torch.Size | tuple[int, ...]) -> torch.Tensor: - numel = 1 - for s in shape: - numel *= int(s) - off = self._cursor - self._cursor += self._round(numel) - if self._cursor > self._reserved: - raise RuntimeError( - f"symmetric arena overflow: wanted {self._cursor} elems, reserved {self._reserved}. " - "The sizing walk and the dispensing walk must visit the same shards in the same order." - ) - return self.buf[off : off + numel].view(shape) - - # -- query ------------------------------------------------------------ - @property - def nbytes(self) -> int: - return self._reserved * self.dtype.itemsize - - def offset_of(self, t: torch.Tensor) -> int: - return (t.data_ptr() - self.buf.data_ptr()) // self.buf.element_size() - - def peer_views(self, t: torch.Tensor) -> list[torch.Tensor]: - """``world_size`` views of the same shard, one per rank, in rank order.""" - off, numel = self.offset_of(t), t.numel() - return [p[off : off + numel].view(t.shape) for p in self.peers] - - def contains(self, t: torch.Tensor) -> bool: - if self.buf is None: - return False - base = self.buf.data_ptr() - return base <= t.data_ptr() < base + self.nbytes - - @classmethod - def _round(cls, numel: int) -> int: - return (numel + cls.ALIGN - 1) // cls.ALIGN * cls.ALIGN - - -@dataclass(frozen=True) -class ShardEntry: - """What the run-time gather needs to know about one local shard.""" - - arena: SymmArena - offset: int - local: torch.Tensor - peer_views: tuple[torch.Tensor, ...] - - @property - def shape(self) -> tuple[int, ...]: - return tuple(self.local.shape) - - -# Keyed by ``data_ptr()``: the gather op only sees a plain tensor. -_SHARD_REGISTRY: dict[int, ShardEntry] = {} -_ARENAS: list[SymmArena] = [] -_BARRIER_DONE = False - - -def register_shard(local: torch.Tensor, arena: SymmArena) -> ShardEntry: - entry = ShardEntry(arena=arena, offset=arena.offset_of(local), local=local, peer_views=tuple(arena.peer_views(local))) - _SHARD_REGISTRY[local.data_ptr()] = entry - return entry - - -def lookup_shard(data_ptr: int) -> ShardEntry | None: - return _SHARD_REGISTRY.get(data_ptr) - - -def registered_arenas() -> list[SymmArena]: - return list(_ARENAS) - - -def find_shard_by_layout(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor | None: - """Any registered shard with this layout -- the cost model cannot replay a gather on a generic ``empty`` (no peers).""" - want = tuple(int(s) for s in shape) - for entry in _SHARD_REGISTRY.values(): - if entry.shape == want and entry.local.dtype == dtype: - return entry.local - return None - - -def reset_registry() -> None: - """Test-only: drop every arena so a new model can be built in-process.""" - global _BARRIER_DONE - _SHARD_REGISTRY.clear() - _ARENAS.clear() - _BARRIER_DONE = False - - -def barrier_after_load() -> None: - """Publish every rank's freshly written shards, once per process. - - Must run after weights are loaded and before the first peer read. - """ - global _BARRIER_DONE - if _BARRIER_DONE or not _ARENAS: - return - if dist.is_available() and dist.is_initialized(): - torch.cuda.synchronize() - dist.barrier() - _BARRIER_DONE = True - magi_logger.info( - "Symmetric arena: published %d arena(s), %.1f MiB, %d shards; steady state is barrier-free", - len(_ARENAS), - sum(a.nbytes for a in _ARENAS) / 2**20, - len(_SHARD_REGISTRY), - ) - - -def _is_gatherable_shard(t: object) -> bool: - """A Shard(0) DTensor on a 1-D mesh -- the only placement the copy-engine gather handles.""" - from torch.distributed.tensor import DTensor, Shard - - if not isinstance(t, DTensor): - return False - placements = t.placements - return len(placements) == 1 and isinstance(placements[0], Shard) and placements[0].dim == 0 - - -def _group_name_of(t) -> str | None: - try: - return t.device_mesh._dim_group_names[0] - except Exception: # noqa: BLE001 - return None - - -def _arena_key(t) -> tuple[torch.dtype, str]: - """One window per (dtype, process group). Same dtype on two meshes (gaga4 FSDP + edp) must not share a window.""" - group_name = _group_name_of(t) - if group_name is None: - raise RuntimeError(f"cannot resolve the process group of a Shard(0) parameter on mesh {t.device_mesh}") - return (t.dtype, group_name) - - -def _apply_order_entries(mod: nn.Module): - """``(owner, name, param)`` in ``_apply`` order: post-order, every ``_parameters`` entry. - - Not ``named_parameters()`` (pre-order, dedups shared tensors). A different - walk would break cross-rank offset symmetry. Walking ``_parameters`` also - finds SimpleFSDP weights in ``parametrizations.weight.original``. - """ - for child in mod.children(): - yield from _apply_order_entries(child) - for name, p in mod._parameters.items(): - if p is not None: - yield mod, name, p - - -def _plan_arenas(shards: list, device: torch.device) -> dict[tuple[torch.dtype, str], SymmArena]: - """Size and commit one window per (dtype, group). Dedup by identity so a tied weight reserves a single slot.""" - arenas: dict[tuple[torch.dtype, str], SymmArena] = {} - seen: set[int] = set() - for p in shards: - if id(p) in seen: - continue - seen.add(id(p)) - key = _arena_key(p) - arena = arenas.get(key) - if arena is None: - arena = arenas[key] = SymmArena(p.dtype, device, key[1]) - arena.reserve(p._local_tensor.numel()) - - for arena in arenas.values(): - arena.commit() # the only collective, once per window - _ARENAS.extend(arenas.values()) - return arenas - - -def materialize_into_arenas(mod: nn.Module, device: torch.device) -> dict[tuple[torch.dtype, str], SymmArena]: - """Size windows for ``mod``'s Shard(0) shards while they are still on meta. Non-gatherable params are left to the caller.""" - shards = [p for _, _, p in _apply_order_entries(mod) if _is_gatherable_shard(p)] - if not shards: - return {} - return _plan_arenas(shards, device) - - -def migrate_to_arenas(root: nn.Module) -> dict[tuple[torch.dtype, str], SymmArena]: - """Copy already-allocated Shard(0) shards into symmetric memory. - - Used when ``magi_compile(model, ...)`` is given a live model rather than a - meta + ``to_empty`` path. ``load_state_dict(assign=True)`` after this would - replace arena views with ordinary tensors; the gather then rejects them. - """ - entries = [(m, n, p) for m, n, p in _apply_order_entries(root) if _is_gatherable_shard(p)] - if not entries: - return {} - - device = entries[0][2]._local_tensor.device - if device.type != "cuda": - raise RuntimeError(f"symmetric memory needs the shards on cuda, found {device}") - arenas = _plan_arenas([p for _, _, p in entries], device) - - from torch.distributed.tensor import DTensor - - views: dict[int, torch.Tensor] = {} - for owner, name, p in entries: - local = views.get(id(p)) - if local is None: - arena = arenas[_arena_key(p)] - local = views[id(p)] = arena.take(p._local_tensor.shape) - local.copy_(p._local_tensor) - register_shard(local, arena) - moved = DTensor.from_local(local, p.device_mesh, p.placements, run_check=False) - owner.register_parameter(name, nn.Parameter(moved, requires_grad=p.requires_grad)) - - magi_logger.info( - "Symmetric arena: migrated %d shard(s) into %.1f MiB across %d window(s)", - len(views), - sum(a.nbytes for a in arenas.values()) / 2**20, - len(arenas), - ) - return arenas - - -def patch_symm_arena_apply(cls: type[nn.Module]) -> None: - """Install the ``_apply`` interception on a decorated class. - - Mirrors ``_patch_cpu_offload_apply``: take over for ``to_empty``'s lambda, delegate everything else. - """ - if getattr(cls, "_magi_symm_apply_patched", False): - return - orig_apply = cls._apply - magi_logger.info("Symmetric arena: intercepting %s._apply for copy-engine FSDP", cls.__name__) - - def _symm_apply(self, fn, recurse: bool = True): - if getattr(fn, "__qualname__", "") != _TO_EMPTY_LAMBDA: - return orig_apply(self, fn, recurse) - if getattr(self, "_magi_symm_arenas", None) is not None: - return orig_apply(self, fn, recurse) - - device = torch.device(inspect.getclosurevars(fn).nonlocals["device"]) - from torch.distributed.tensor import DTensor - - arenas = materialize_into_arenas(self, device) - if not arenas: - return orig_apply(self, fn, recurse) - - views: dict[int, torch.Tensor] = {} - - def materialize(t: torch.Tensor) -> torch.Tensor: - if not _is_gatherable_shard(t): - return torch.empty_like(t, device=device) - # Tied weight: same view so tying survives materialization. - local = views.get(id(t)) - if local is None: - arena = arenas[_arena_key(t)] - local = views[id(t)] = arena.take(t._local_tensor.shape) - register_shard(local, arena) - return DTensor.from_local(local, t.device_mesh, t.placements, run_check=False) - - # Do not forge to_empty's qualname: a nested decorated block must fail - # the check above and delegate, so its params land in *this* arena. - out = orig_apply(self, materialize, recurse) - self._magi_symm_arenas = arenas - magi_logger.info( - "Symmetric arena: %s materialized %d shard(s) into %.1f MiB across %d window(s)", - cls.__name__, - len(views), - sum(a.nbytes for a in arenas.values()) / 2**20, - len(arenas), - ) - return out - - cls._apply = _symm_apply - cls._magi_symm_apply_patched = True 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 2198a6e..b3099b0 100644 --- a/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py +++ b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py @@ -167,8 +167,8 @@ def fn(x, w0, shard): ce_shards: list = [] if args.copy_engine: - from magi_compiler.runtime.symm_all_gather import SYMM_ALL_GATHER - from magi_compiler.runtime.symm_arena import SymmArena, register_shard + from magi_compiler.symm_mem import SymmArena, register_shard + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER arena = SymmArena(torch.bfloat16, torch.device("cuda", dev), grp) for _ in range(N_CE_LAYERS): From f9c76ba743ae1b816557a82decd1bb76e07db94e Mon Sep 17 00:00:00 2001 From: wtr Date: Thu, 27 Aug 2026 21:19:30 +0800 Subject: [PATCH 03/16] add ci tests --- magi_compiler/passes/fsdp_overlap/reorder.py | 9 - magi_compiler/profiling/runtime_estimator.py | 18 +- .../fsdp/test_profiling_estimator.py | 320 ++++++++++- tests/feature_tests/test_symm_ag_rewrite.py | 309 +++++++++++ tests/feature_tests/test_symm_all_gather.py | 282 ++++++++++ tests/feature_tests/test_symm_arena.py | 503 ++++++++++++++++++ tests/feature_tests/test_symm_e2e.py | 91 ++++ 7 files changed, 1504 insertions(+), 28 deletions(-) create mode 100644 tests/feature_tests/test_symm_ag_rewrite.py create mode 100644 tests/feature_tests/test_symm_all_gather.py create mode 100644 tests/feature_tests/test_symm_arena.py create mode 100644 tests/feature_tests/test_symm_e2e.py diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index da4a7df..c7521f1 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -517,10 +517,6 @@ 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] - # The upper bound keeps the gather below the next collective, but it - # must never win against the floor: with no collective above (or, on a - # pure copy-engine graph, no skeleton at all) it degenerates to the - # launch's current index and would silently undo a legal hoist. new_target = min(max(target, slot_lo), max(slot_hi, slot_lo)) targets[launch] = (new_target, group) magi_logger.debug( @@ -558,11 +554,6 @@ def _launch_group(self, launch, order, buf_to_snode, users) -> list[BaseSchedule if deps and all(d.name in produced for d in deps): group.append(s) elif _is_symm_ag_ir(node): - # A FallbackKernel's result is re-exposed through an alias snode - # (``buf1 = buf0`` in the generated code), which is what the wait and - # the consumer actually read. Inductor's own collectives have no such - # layer, so this is the one structural difference the copy-engine - # transport introduces -- and it has to move with the launch. for s in order: if s is launch or contains_wait(s) or not self._is_transparent(s): continue diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index c8248e3..c3fa09f 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -377,14 +377,7 @@ def _is_symm_ag_coalesced_ir(node) -> bool: def _symm_ag_spec(node): - """(shapes, dtype, group_size, group_name) for a copy-engine gather. - - ``shapes`` is a tuple of local-shard shapes (one member, or one per - coalesced input). Read off ``constant_args``, the way ``_collective_spec`` - reads a group name. Not off ``get_origin_node()``: Inductor leaves that - unset on these nodes, and the resulting ``None`` silently degraded the - gather to a zero cost. - """ + """(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 @@ -427,14 +420,7 @@ def _symm_ag_launch_wait(snode: BaseSchedulerNode): def _measure_symm_ag(snode: BaseSchedulerNode) -> float: - """Time the copy-engine gather TOGETHER WITH its wait. - - Timing the launch alone would measure the CPU issue cost and nothing else: - the copies run on a side stream, so without the wait the timing events on the - current stream close before a single byte has moved. That reads as ~3us for - a gather that really takes tens of microseconds, and the reorder pass then - sizes a window an order of magnitude too small. - """ + """Time ``wait(launch())``. Launch-only is ~3us CPU issue; copies run on a side stream.""" pair = _symm_ag_launch_wait(snode) if pair is None: return 0.0 diff --git a/tests/feature_tests/fsdp/test_profiling_estimator.py b/tests/feature_tests/fsdp/test_profiling_estimator.py index b603ac6..275a9f4 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, + _is_symm_ag_coalesced_ir, + _leaf_symm_ag, + _measure_extern, + _measure_symm_ag, + _realize_arg, + _static, + _symm_ag_label, + _symm_ag_launch_wait, + _symm_ag_spec, +) 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} @@ -528,3 +540,305 @@ 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 SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + + return SYMM_ALL_GATHER, SYMM_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 _FakeSymmAgIR: + """A ``magi::symm_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=_FakeSymmAgIR(ag, [shape], group_size=group_size)) + + +# --- spec / recognition / label: structural, no window needed --------------- +def test_symm_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 = _FakeSymmAgIR(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 = _symm_ag_spec(node) + assert shapes == ((_CE_ROWS, 64),) + assert dtype is torch.bfloat16 + assert group_size == 8 + assert group_name == "fsdp" + + +def test_symm_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 = _symm_ag_spec(_FakeSymmAgIR(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_symm_ag_spec_is_none_when_the_group_args_are_missing(constant_args): + ag, _ = _ce_ops() + node = _FakeSymmAgIR(ag, [(8, 4)]) + node.constant_args = constant_args + assert _symm_ag_spec(node) is None + + +def test_symm_ag_spec_is_none_when_the_node_has_no_inputs(): + ag, _ = _ce_ops() + assert _symm_ag_spec(_FakeSymmAgIR(ag, [])) is None + + +def test_leaf_symm_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 = _FakeSymmAgIR(ag, [(8, 4)]) + + assert _leaf_symm_ag(_FakeGatherSnode(node=node)) is node + fused = _FakeGatherSnode(node=None, snodes=[_FakeGatherSnode(node=_FakeMmIR()), _FakeGatherSnode(node=node)]) + assert _leaf_symm_ag(fused) is node + + +def test_leaf_symm_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_symm_ag(_FakeGatherSnode(node=_FakeMmIR())) is None + assert _leaf_symm_ag(_FakeGatherSnode()) is None + + +def test_coalesced_gather_is_distinguished_from_a_single_one(): + ag, coalesced = _ce_ops() + assert _is_symm_ag_coalesced_ir(_FakeSymmAgIR(coalesced, [(8, 4)])) + assert not _is_symm_ag_coalesced_ir(_FakeSymmAgIR(ag, [(8, 4)])) + + +def test_symm_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 = _symm_ag_label(_FakeGatherSnode(node=_FakeSymmAgIR(ag, [(8, 4)], group_size=2))) + batch = _symm_ag_label(_FakeGatherSnode(node=_FakeSymmAgIR(coalesced, [(8, 4), (16, 4)], group_size=2))) + + assert single == "symm_all_gather(ws=2,8x4)" + assert batch == "symm_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.""" + from magi_compiler.symm_mem import SymmArena, register_shard + + arena = SymmArena(dtype, torch.device("cuda", 0), _ce_group_name()) + for shape in shapes: + arena.reserve(shape[0] * shape[1]) + arena.commit() + + shards = [] + for i, shape in enumerate(shapes): + s = arena.take(shape) + s.fill_(i + 1) + register_shard(s, arena) + shards.append(s) + torch.cuda.synchronize() + return shards + + +def _ce_replay_snode(shapes, coalesced=False): + ag, ag_coalesced = _ce_ops() + op = ag_coalesced if coalesced else ag + return _FakeGatherSnode(node=_FakeSymmAgIR(op, shapes, group_size=1, group_name=_ce_group_name())) + + +@requires_cuda +def test_symm_ag_replay_gathers_the_registered_shard(pg_1rank, symm_registry): + """The replay must run the real op on a real arena 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 = _symm_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_symm_ag_coalesced_replay_covers_every_member(pg_1rank, symm_registry): + shapes = [(_CE_ROWS, 64), (_CE_ROWS, 32)] + shards = _register_shards(shapes) + + launch, wait = _symm_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_symm_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_symm_ag(_ce_replay_snode([(1024, 1024)])) > 0.0 + + +@requires_cuda +def test_measure_symm_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 _symm_ag_launch_wait(snode) is None + assert _measure_symm_ag(snode) == 0.0 + + +# --- __call__: the ("symm_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_symm_ag", fake_measure) + return calls + + +def test_isomorphic_symm_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_symm_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_symm_ag_entry_is_tagged_as_a_copy_engine_gather(ce_measure_calls): + est = ProfilingRuntimeEstimator() + est(_ce_snode()) + + (entry,) = est.table.values() + assert entry.kind == "symm_ag" + assert entry.measured + assert entry.ns == 222.0 + assert "symm_all_gather" in entry.label + + +def test_symm_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/test_symm_ag_rewrite.py b/tests/feature_tests/test_symm_ag_rewrite.py new file mode 100644 index 0000000..8d0253e --- /dev/null +++ b/tests/feature_tests/test_symm_ag_rewrite.py @@ -0,0 +1,309 @@ +# Copyright (c) 2026 SandAI. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the copy-engine retarget pass (step 3 of the landing order). + +The pass is deliberately tiny -- swap one node's target -- so what +is worth testing is what it *refuses* to touch. A gather whose input is a cast +or a pad reads a tensor the caching allocator produced, not the symmetric window, +and retargeting it would make the operator reject it at run time. A gather that +was never marked as a SimpleFSDP weight gather (CP, TP, MoE) must stay on NCCL +even in copy-engine mode. + +Built on the real lowering pass so the node shapes are the ones production +produces, with a 1-rank mesh (nothing here touches the transport). +""" + +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 SYMM_ALL_GATHER + + return SYMM_ALL_GATHER + + +def _graph_with_gathers(mesh, n: int, *, derive: str | None = None, marked: bool = True, shapes: list[tuple] | None = None): + """``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. + """ + from torch.distributed.tensor import Shard, distribute_tensor + + 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: + ag.meta["magi_fsdp_weight_ag"] = True + 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] + + +@requires_cuda +def test_marked_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 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) + 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 +@pytest.mark.parametrize("derive", ["cast", "pad"]) +def test_derived_shards_stay_on_nccl(mesh_1rank, derive): + """A cast or pad output is not in the symmetric window; retargeting it would + be rejected at run time, so it must stay on NCCL.""" + from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine + + gm = _graph_with_gathers(mesh_1rank, 3, derive=derive) + 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_unmarked_gathers_are_left_alone(mesh_1rank): + """CP / TP / MoE gathers are never marked, and must survive copy-engine mode + untouched -- the transports coexist in one graph.""" + from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine + + gm = _graph_with_gathers(mesh_1rank, 3, marked=False) + assert rewrite_weight_ag_to_copy_engine(gm) == 0 + assert len(_gathers(gm, _AG)) == 3 + + +@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 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_end_to_end_lowering_then_rewrite(mesh_1rank): + """The pass has to match what the lowering pass really emits, not what this + file thinks it emits.""" + from test_fsdp_overlap_lowering import _build_redistribute_graph + + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + + gm = _build_redistribute_graph(mesh_1rank, "model_fc1_weight_parameter") + lower_and_bucket_full_graph(gm, "none", transport="copy_engine") + + assert not _gathers(gm, _AG) + symm = _gathers(gm, _symm_op()) + assert len(symm) == 1 + assert _gathers(gm, _WAIT)[0].args[0] is symm[0] + + +@requires_cuda +def test_copy_engine_buckets_then_rewrites_coalesced(mesh_1rank): + """Phase-1 wrap: arena gathers are bucketed first, then the coalesced + launch is retargeted. Members stay separate dests underneath.""" + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + + gm = _graph_with_gathers(mesh_1rank, 4) + n = lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine") + assert n == 1 + assert not _gathers(gm, _AG) + assert not _gathers(gm, _AG_COALESCED) + assert not _gathers(gm, SYMM_ALL_GATHER) + assert len(_gathers(gm, SYMM_ALL_GATHER_COALESCED)) == 1 + waits = _gathers(gm, _WAIT) + assert len(waits) == 4 + + +@requires_cuda +def test_copy_engine_does_not_bucket_cast_gathers(mesh_1rank): + """Cast outputs are not arena shards; they must stay on NCCL and not join a CE bucket.""" + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED + + gm = _graph_with_gathers(mesh_1rank, 3, derive="cast") + n = lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine") + assert n == 0 + assert len(_gathers(gm, _AG)) == 3 + assert not _gathers(gm, SYMM_ALL_GATHER_COALESCED) + + +def _coalesced_graph(mesh, n: int, *, derive: str | None = None, marked: 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 + + 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 + if derive == "cast": + cur = g.call_function(_TO_COPY, (cur,), {"dtype": torch.float32}) + cur.meta["example_value"] = local._local_tensor.to(torch.float32) + 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: + ag.meta["magi_fsdp_weight_ag"] = True + 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_coalesced_of_arena_shards_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 SYMM_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, SYMM_ALL_GATHER_COALESCED)) == 1 + + +@requires_cuda +def test_coalesced_is_all_or_nothing(mesh_1rank): + """A bucket is one submission, so it can only go to the copy engine if + *every* member is an arena shard -- one cast 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 rewrite_weight_ag_to_copy_engine + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED + + gm = _coalesced_graph(mesh_1rank, 3, derive="cast") + assert rewrite_weight_ag_to_copy_engine(gm) == 0 + assert len(_gathers(gm, _AG_COALESCED)) == 1 + assert not _gathers(gm, SYMM_ALL_GATHER_COALESCED) + + +@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 + + +@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 diff --git a/tests/feature_tests/test_symm_all_gather.py b/tests/feature_tests/test_symm_all_gather.py new file mode 100644 index 0000000..0fac81c --- /dev/null +++ b/tests/feature_tests/test_symm_all_gather.py @@ -0,0 +1,282 @@ +# 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::symm_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 arena 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_symm_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::symm_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 _arena_model(mesh, hidden: int = 64, n_layers: int = 3, dtype=torch.bfloat16): + """A meta-built, Shard(0)-sharded model materialized into a symmetric arena, + exactly the shape step 1 produces.""" + from torch.distributed.tensor import Shard, distribute_tensor + + from magi_compiler.symm_mem import materialize_into_arenas + + 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("meta"): + 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)]))) + + device = torch.device("cuda", 0) + from torch.distributed.tensor import DTensor + + from magi_compiler.symm_mem.arena import _arena_key, register_shard + + arenas = materialize_into_arenas(model, device) + views: dict[int, torch.Tensor] = {} + + def materialize(t): + if isinstance(t, DTensor): + arena = arenas[_arena_key(t)] + local = arena.take(t._local_tensor.shape) + register_shard(local, arena) + views[id(t)] = local + return DTensor.from_local(local, t.device_mesh, t.placements, run_check=False) + return torch.empty_like(t, device=device) + + materialize.__qualname__ = "Module.to_empty.." + nn.Module._apply(model, materialize) + + 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 = _arena_model(mesh_1rank) + + for shard in shards: + got = _WAIT(torch.ops.magi.symm_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 = _arena_model(mesh_1rank) + outs = torch.ops.magi.symm_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.symm_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 = _arena_model(mesh_1rank) + shard = shards[0] + + out = _WAIT(torch.ops.magi.symm_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 = _arena_model(mesh_1rank) + a = torch.ops.magi.symm_all_gather(shards[0], 1, "") + b = torch.ops.magi.symm_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 arena 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.symm_all_gather(ordinary, 1, "") + + +@requires_cuda +def test_group_size_mismatch_is_rejected(mesh_1rank): + _, shards = _arena_model(mesh_1rank) + with pytest.raises(RuntimeError, match="peers but the gather asks for"): + torch.ops.magi.symm_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 = _arena_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.symm_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.symm_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.symm_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 = _arena_model(mesh_1rank) + batched = _WAIT(torch.ops.magi.symm_all_gather(shards[0], 1, "")) + torch.cuda.synchronize() + + monkeypatch.setattr(ag_mod, "_batcher", lambda: None) + fallback = _WAIT(torch.ops.magi.symm_all_gather(shards[0], 1, "")) + torch.cuda.synchronize() + + assert torch.equal(fallback, batched) diff --git a/tests/feature_tests/test_symm_arena.py b/tests/feature_tests/test_symm_arena.py new file mode 100644 index 0000000..269c70d --- /dev/null +++ b/tests/feature_tests/test_symm_arena.py @@ -0,0 +1,503 @@ +# 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. + +"""Materialization tests for ``fsdp_config.transport="copy_engine"``. + +These cover step 1 of the landing order: the weights move into symmetric memory +and nothing else changes. The interesting property is *scope* -- the decorated +class must claim exactly its own subtree, from a ``to_empty`` issued on the root +model, without the builder participating. + +Single rank on purpose: the placement, scoping, aliasing and fallback logic is +rank-independent, and a 1-rank symmetric window exercises the same allocation +path. The multi-rank peer reads are covered by +``example/inference/fsdp_overlap/verify_symm_param_hook.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") + + +@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(autouse=True) +def _clean_registry(): + from magi_compiler.symm_mem import reset_registry + + reset_registry() + yield + reset_registry() + + +def _decorate(cls: type, transport: str) -> type: + """Run the real decorator, so the config wiring is under test too.""" + from magi_compiler import magi_compile + from magi_compiler.config import CompileMode + + def patch(cfg): + cfg.compile_mode = CompileMode.MAGI_COMPILE + cfg.fsdp_config.enable_fsdp = True + cfg.fsdp_config.transport = transport + return cfg + + return magi_compile(cls, config_patch=patch, dynamic_arg_dims={"x": 0}) + + +def _shard(model: nn.Module, mesh, placement=None) -> None: + """Wrap every parameter as a DTensor, like torchtitan ``data_parallel``.""" + from torch.distributed.tensor import Shard, distribute_tensor + + placements = [placement or Shard(0)] + 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, placements))) + + +class Layer(nn.Module): + def __init__(self, hidden: int, dtype: torch.dtype): + super().__init__() + self.wq = nn.Linear(hidden, hidden, bias=False, dtype=dtype) + self.wo = nn.Linear(hidden, hidden, bias=False, dtype=dtype) + + def forward(self, x): + return self.wo(self.wq(x)) + + +def _make_root(block_cls: type, hidden: int, n_layers: int, dtype: torch.dtype) -> nn.Module: + class Root(nn.Module): + """``to_empty`` is called on this, never on the decorated block.""" + + def __init__(self): + 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)) + + with torch.device("meta"): + return Root() + + +def _make_block_cls(transport: str) -> type: + class Block(nn.Module): + def __init__(self, hidden: int, n_layers: int, dtype: torch.dtype): + super().__init__() + self.layers = nn.ModuleList([Layer(hidden, dtype) for _ in range(n_layers)]) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + return _decorate(Block, transport) + + +@requires_cuda +def test_root_to_empty_puts_block_shards_in_arena(mesh_1rank): + """The decorated block claims its own subtree; the rest stays ordinary.""" + from magi_compiler.symm_mem import lookup_shard, registered_arenas + + hidden, n_layers = 64, 3 + model = _make_root(_make_block_cls("copy_engine"), hidden, n_layers, torch.bfloat16) + _shard(model, mesh_1rank) + model.to_empty(device=torch.device("cuda", 0)) + + arenas = registered_arenas() + assert len(arenas) == 1, "one window per (dtype, group), not one per weight" + arena = arenas[0] + + block_shards = list(model.block.parameters()) + assert len(block_shards) == 2 * n_layers + assert all(arena.contains(p._local_tensor) for p in block_shards) + assert all(lookup_shard(p._local_tensor.data_ptr()) is not None for p in block_shards) + + # Outside the decorated class: untouched by the interception. + assert not arena.contains(model.head.weight._local_tensor) + assert lookup_shard(model.head.weight._local_tensor.data_ptr()) is None + + # Still ordinary, working DTensors on real storage. + assert all(p._local_tensor.device.type == "cuda" for p in block_shards) + assert not any(p.is_meta for p in block_shards) + + +@requires_cuda +def test_nccl_transport_leaves_allocation_alone(mesh_1rank): + """The default transport must not install the interception at all.""" + from magi_compiler.symm_mem import registered_arenas + + model = _make_root(_make_block_cls("nccl"), 64, 2, torch.bfloat16) + _shard(model, mesh_1rank) + model.to_empty(device=torch.device("cuda", 0)) + + assert registered_arenas() == [] + assert all(p._local_tensor.device.type == "cuda" for p in model.parameters()) + + +@requires_cuda +def test_peer_view_round_trips_through_arena(mesh_1rank): + """A shard written locally must be visible through its own peer view: this is + the addressing the copy-engine gather depends on.""" + from magi_compiler.symm_mem import lookup_shard + + model = _make_root(_make_block_cls("copy_engine"), 64, 2, torch.bfloat16) + _shard(model, mesh_1rank) + model.to_empty(device=torch.device("cuda", 0)) + + for i, p in enumerate(model.block.parameters()): + p._local_tensor.fill_(i + 1) + torch.cuda.synchronize() + + for i, p in enumerate(model.block.parameters()): + entry = lookup_shard(p._local_tensor.data_ptr()) + assert len(entry.peer_views) == 1 + assert torch.equal(entry.peer_views[0], p._local_tensor), f"shard {i} aliases the wrong bytes" + + +@requires_cuda +def test_tied_weights_share_one_slot(mesh_1rank): + """A shared parameter is visited once per referencing module, so it must get + one slot and keep its identity -- two allocations would silently untie it and + overflow a window sized by a deduped walk. + + Note the tie has to be (re)established *after* sharding: torchtitan's + ``data_parallel`` calls ``distribute_tensor`` per module, which replaces each + entry with its own DTensor and unties them on its own. + """ + from magi_compiler.symm_mem import registered_arenas + + hidden = 64 + + class TiedBlock(nn.Module): + def __init__(self, hidden: int, n_layers: int, dtype: torch.dtype): + super().__init__() + self.a = nn.Linear(hidden, hidden, bias=False, dtype=dtype) + self.b = nn.Linear(hidden, hidden, bias=False, dtype=dtype) + + def forward(self, x): + return self.b(self.a(x)) + + model = _make_root(_decorate(TiedBlock, "copy_engine"), hidden, 1, torch.bfloat16) + _shard(model, mesh_1rank) + model.block.b.weight = model.block.a.weight + model.to_empty(device=torch.device("cuda", 0)) + + arena = registered_arenas()[0] + assert model.block.a.weight is model.block.b.weight, "tying must survive materialization" + assert arena.contains(model.block.a.weight._local_tensor) + # One slot in the window, not two. + slot = arena.ALIGN * ((hidden * hidden + arena.ALIGN - 1) // arena.ALIGN) + assert arena.nbytes == slot * torch.bfloat16.itemsize + + +@requires_cuda +def test_nested_decorated_block_shares_the_outer_arena(mesh_1rank): + """A decorated block inside a decorated block must not open a second window: + the inner one has to fail the lambda check and delegate.""" + from magi_compiler.symm_mem import registered_arenas + + inner_cls = _decorate( + type( + "Inner", + (nn.Module,), + { + "__init__": lambda self, hidden, dtype: ( + nn.Module.__init__(self), + setattr(self, "lin", nn.Linear(hidden, hidden, bias=False, dtype=dtype)), + )[0], + "forward": lambda self, x: self.lin(x), + }, + ), + "copy_engine", + ) + + class Outer(nn.Module): + def __init__(self, hidden: int, n_layers: int, dtype: torch.dtype): + super().__init__() + self.own = nn.Linear(hidden, hidden, bias=False, dtype=dtype) + self.inner = inner_cls(hidden, dtype) + + def forward(self, x): + return self.inner(self.own(x)) + + model = _make_root(_decorate(Outer, "copy_engine"), 64, 1, torch.bfloat16) + _shard(model, mesh_1rank) + model.to_empty(device=torch.device("cuda", 0)) + + arenas = registered_arenas() + assert len(arenas) == 1, f"nested decoration opened {len(arenas)} windows" + assert arenas[0].contains(model.block.own.weight._local_tensor) + assert arenas[0].contains(model.block.inner.lin.weight._local_tensor) + + +@requires_cuda +def test_non_shard0_placement_falls_back(mesh_1rank): + """Replicate weights are not gatherable, so they must be allocated normally + rather than silently placed in the arena.""" + from torch.distributed.tensor import Replicate + + from magi_compiler.symm_mem import registered_arenas + + model = _make_root(_make_block_cls("copy_engine"), 64, 2, torch.bfloat16) + _shard(model, mesh_1rank, placement=Replicate()) + model.to_empty(device=torch.device("cuda", 0)) + + assert registered_arenas() == [] + assert all(p._local_tensor.device.type == "cuda" for p in model.block.parameters()) + + +@requires_cuda +def test_two_process_groups_same_dtype_get_two_windows(mesh_1rank, monkeypatch): + """gaga4 shards bf16 dense weights on the FSDP mesh and bf16 experts on the + orthogonal edp mesh. One window per dtype would either rendezvous on the + wrong group or mix offsets that are only meaningful inside one group. + + Fake group names cannot rendezvous, so ``commit`` is stubbed; the assertion + is that planning opens two windows keyed by (dtype, group). + """ + from torch.distributed.tensor import Shard, distribute_tensor + + from magi_compiler.symm_mem import arena as sa + + hidden = 64 + a = nn.Parameter(distribute_tensor(torch.empty(hidden, hidden, dtype=torch.bfloat16), mesh_1rank, [Shard(0)])) + b = nn.Parameter(distribute_tensor(torch.empty(hidden, hidden, dtype=torch.bfloat16), mesh_1rank, [Shard(0)])) + + monkeypatch.setattr(sa, "_group_name_of", lambda p, _a=a: "dense_fsdp" if p is _a else "edp") + monkeypatch.setattr(sa.SymmArena, "commit", lambda self: None) + arenas = sa._plan_arenas([a, b], torch.device("cuda", 0)) + assert set(arenas) == {(torch.bfloat16, "dense_fsdp"), (torch.bfloat16, "edp")} + assert arenas[(torch.bfloat16, "dense_fsdp")].group_name == "dense_fsdp" + assert arenas[(torch.bfloat16, "edp")].group_name == "edp" + + +@requires_cuda +def test_barrier_after_load_is_idempotent(mesh_1rank): + from magi_compiler.symm_mem import barrier_after_load + + model = _make_root(_make_block_cls("copy_engine"), 64, 2, torch.bfloat16) + _shard(model, mesh_1rank) + model.to_empty(device=torch.device("cuda", 0)) + + barrier_after_load() + barrier_after_load() # a second call must not issue a second collective + + +# --------------------------------------------------------------------------- +# Window suballocation. Every rank must agree on which bytes are which shard, +# and nothing at run time re-derives that -- the gather trusts the offsets. +# --------------------------------------------------------------------------- +@pytest.fixture +def group_name(): + import torch.distributed as dist + + return dist.group.WORLD.group_name + + +def _committed_arena(group_name, numels, dtype=torch.bfloat16): + from magi_compiler.symm_mem import SymmArena + + arena = SymmArena(dtype, torch.device("cuda", 0), group_name) + for n in numels: + arena.reserve(n) + arena.commit() + return arena + + +@requires_cuda +def test_shards_are_dispensed_at_aligned_offsets(mesh_1rank, group_name): + """Slots are padded to ``ALIGN`` for copy-engine throughput, so the second + shard does not start where the first one ends. ``offset_of`` is what the + peer views are built from, so it has to agree with what ``take`` handed out.""" + from magi_compiler.symm_mem import SymmArena + + rows, cols = 3, 5 # 15 elems: deliberately not a multiple of ALIGN + arena = _committed_arena(group_name, [rows * cols, rows * cols]) + first = arena.take((rows, cols)) + second = arena.take((rows, cols)) + + assert arena.offset_of(first) == 0 + assert arena.offset_of(second) == SymmArena.ALIGN + assert arena.contains(first) and arena.contains(second) + assert first.shape == (rows, cols) + + +@requires_cuda +def test_dispensing_more_than_was_reserved_is_an_error(mesh_1rank, group_name): + """The sizing walk and the dispensing walk are two separate traversals; if + they ever disagree the shards silently overlap, so the window must run out + rather than hand back memory reserved for someone else.""" + arena = _committed_arena(group_name, [64]) + arena.take((8, 8)) + with pytest.raises(RuntimeError, match="symmetric arena overflow"): + arena.take((8, 8)) + + +@requires_cuda +def test_contains_rejects_memory_outside_the_window(mesh_1rank, group_name): + """``contains`` is how the rewrite decides a weight is gatherable; a caching + allocator tensor must never pass.""" + arena = _committed_arena(group_name, [64]) + arena.take((8, 8)) + assert not arena.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 must + be a None it can degrade on, not a wrong-dtype shard it would gather.""" + from magi_compiler.symm_mem import find_shard_by_layout, register_shard + + arena = _committed_arena(group_name, [8 * 4, 16 * 4]) + small = arena.take((8, 4)) + large = arena.take((16, 4)) + register_shard(small, arena) + register_shard(large, arena) + + 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_arenas_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 find_shard_by_layout, lookup_shard, register_shard, registered_arenas, reset_registry + + arena = _committed_arena(group_name, [8 * 4]) + shard = arena.take((8, 4)) + register_shard(shard, arena) + assert lookup_shard(shard.data_ptr()) is not None + + reset_registry() + assert registered_arenas() == [] + assert lookup_shard(shard.data_ptr()) is None + assert find_shard_by_layout((8, 4), torch.bfloat16) is None + + +# --------------------------------------------------------------------------- +# migrate_to_arenas -- the live-model path (magi_compile on an already +# materialized model, where there is no to_empty to intercept). +# --------------------------------------------------------------------------- +@requires_cuda +def test_migrate_moves_live_shards_and_keeps_their_values(mesh_1rank): + """Unlike ``to_empty``, this runs on weights that already hold data, so the + copy is load-bearing: dropping it would gather uninitialized memory.""" + from magi_compiler.symm_mem import lookup_shard, migrate_to_arenas + + hidden = 64 + + class Live(nn.Module): + def __init__(self): + super().__init__() + self.a = nn.Linear(hidden, hidden, bias=False, dtype=torch.bfloat16) + self.b = nn.Linear(hidden, hidden, bias=False, dtype=torch.bfloat16) + + model = Live().to("cuda") + _shard(model, mesh_1rank) + before = {n: p._local_tensor.clone() for n, p in model.named_parameters()} + + arenas = migrate_to_arenas(model) + assert len(arenas) == 1 + arena = next(iter(arenas.values())) + + for name, p in model.named_parameters(): + local = p._local_tensor + assert arena.contains(local), f"{name} was not migrated" + assert lookup_shard(local.data_ptr()) is not None + assert torch.equal(local, before[name]), f"{name} lost its values" + + +@requires_cuda +def test_migrate_gives_a_tied_weight_one_slot(mesh_1rank): + """Migration rebuilds each parameter, so python identity does not survive -- + what must survive is the storage, or the tie is gone and the window is + overflowed by a walk that sized it once.""" + from magi_compiler.symm_mem import migrate_to_arenas + + hidden = 64 + + class Tied(nn.Module): + def __init__(self): + super().__init__() + self.a = nn.Linear(hidden, hidden, bias=False, dtype=torch.bfloat16) + self.b = nn.Linear(hidden, hidden, bias=False, dtype=torch.bfloat16) + + model = Tied().to("cuda") + _shard(model, mesh_1rank) + model.b.weight = model.a.weight + + arenas = migrate_to_arenas(model) + arena = next(iter(arenas.values())) + assert model.a.weight._local_tensor.data_ptr() == model.b.weight._local_tensor.data_ptr() + slot = arena.ALIGN * ((hidden * hidden + arena.ALIGN - 1) // arena.ALIGN) + assert arena.nbytes == slot * torch.bfloat16.itemsize + + +@requires_cuda +def test_migrate_refuses_a_model_that_was_never_materialized(mesh_1rank): + """``migrate_to_arenas`` is the live-model entry point, so being handed a + still-on-meta model is the way it gets misused. Copying from meta silently + produces a window of uninitialized weights, so it has to fail instead.""" + from torch.distributed.tensor import DTensor, Shard + + from magi_compiler.symm_mem import migrate_to_arenas + + with torch.device("meta"): + model = nn.Linear(8, 8, bias=False, dtype=torch.bfloat16) + local = DTensor.from_local(model.weight.data, mesh_1rank, [Shard(0)], run_check=False) + model.register_parameter("weight", nn.Parameter(local)) + assert model.weight._local_tensor.is_meta # the state under test + + with pytest.raises(RuntimeError, match="needs the shards on cuda"): + migrate_to_arenas(model) + + +@requires_cuda +def test_migrate_leaves_a_model_with_no_gatherable_shards_alone(mesh_1rank): + """A plain (unsharded) model must not open an empty window.""" + from magi_compiler.symm_mem import migrate_to_arenas, registered_arenas + + model = nn.Linear(8, 8, bias=False, dtype=torch.bfloat16).to("cuda") + assert migrate_to_arenas(model) == {} + assert registered_arenas() == [] diff --git a/tests/feature_tests/test_symm_e2e.py b/tests/feature_tests/test_symm_e2e.py new file mode 100644 index 0000000..a233229 --- /dev/null +++ b/tests/feature_tests/test_symm_e2e.py @@ -0,0 +1,91 @@ +# 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 arena 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 example script asserts. Driving it here as a +subprocess keeps that assertion in CI instead of in someone's shell history. +""" + +import os +import shutil +import subprocess +from pathlib import Path + +import pytest +import torch + +_SCRIPT = Path(__file__).resolve().parents[2] / "example" / "inference" / "fsdp_overlap" / "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:] From 018fb9bddcd1d6e9f68e0bed6e5e0f493b8f9222 Mon Sep 17 00:00:00 2001 From: wtr Date: Thu, 27 Aug 2026 22:00:43 +0800 Subject: [PATCH 04/16] chore --- magi_compiler/symm_mem/__init__.py | 48 ++++ magi_compiler/symm_mem/all_gather.py | 233 +++++++++++++++++++ magi_compiler/symm_mem/arena.py | 329 +++++++++++++++++++++++++++ 3 files changed, 610 insertions(+) create mode 100644 magi_compiler/symm_mem/__init__.py create mode 100644 magi_compiler/symm_mem/all_gather.py create mode 100644 magi_compiler/symm_mem/arena.py diff --git a/magi_compiler/symm_mem/__init__.py b/magi_compiler/symm_mem/__init__.py new file mode 100644 index 0000000..8222d17 --- /dev/null +++ b/magi_compiler/symm_mem/__init__.py @@ -0,0 +1,48 @@ +# 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::symm_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 .arena import ( + ShardEntry, + SymmArena, + barrier_after_load, + find_shard_by_layout, + lookup_shard, + materialize_into_arenas, + migrate_to_arenas, + patch_symm_arena_apply, + register_shard, + registered_arenas, + reset_registry, +) + +__all__ = [ + "ShardEntry", + "SymmArena", + "barrier_after_load", + "find_shard_by_layout", + "lookup_shard", + "materialize_into_arenas", + "migrate_to_arenas", + "patch_symm_arena_apply", + "register_shard", + "registered_arenas", + "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..9ecc712 --- /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 .arena 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 arena. +_SCHEMA = "symm_all_gather(Tensor local, int group_size, str group_name) -> Tensor" +_SCHEMA_COALESCED = "symm_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::symm_all_gather got a tensor that is not a registered symmetric-memory shard. " + "Only weights materialized through the arena 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 _symm_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 _symm_all_gather_meta(local: torch.Tensor, group_size: int, group_name: str) -> torch.Tensor: + return _gather_dest(local, group_size) + + +def _symm_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 _symm_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("symm_all_gather", _symm_all_gather, "CUDA") + _LIB.impl("symm_all_gather", _symm_all_gather_meta, "Meta") + + _LIB.define(_SCHEMA_COALESCED) + _LIB.impl("symm_all_gather_coalesced", _symm_all_gather_coalesced, "CUDA") + _LIB.impl("symm_all_gather_coalesced", _symm_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. +SYMM_ALL_GATHER = torch.ops.magi.symm_all_gather.default +SYMM_ALL_GATHER_COALESCED = torch.ops.magi.symm_all_gather_coalesced.default diff --git a/magi_compiler/symm_mem/arena.py b/magi_compiler/symm_mem/arena.py new file mode 100644 index 0000000..a8266cb --- /dev/null +++ b/magi_compiler/symm_mem/arena.py @@ -0,0 +1,329 @@ +# 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 inspect +from dataclasses import dataclass + +import torch +import torch.distributed as dist +import torch.nn as nn + +from magi_compiler.utils import magi_logger + +# Only hijack ``to_empty``'s lambda. ``.cuda()`` / ``.to()`` / ``.float()`` also +# go through ``_apply``; intercepting those would change builder semantics. +_TO_EMPTY_LAMBDA = "Module.to_empty.." + + +class SymmArena: + """One symmetric-memory window, suballocated to many weight shards. + + One window per (decorated block, dtype, process group): a single rendezvous, + and because every rank walks the module tree in the same order, offset ``k`` + is the same shard on every peer. Two meshes sharing a dtype get two windows. + """ + + # 256 bf16 elems = 512B, which the copy engine wants for peak throughput. + ALIGN = 256 + + def __init__(self, dtype: torch.dtype, device: torch.device, group_name: str) -> None: + self.dtype = dtype + self.device = device + self.group_name = group_name + self.buf: torch.Tensor | None = None + self.handle = None + self.peers: list[torch.Tensor] = [] + self._reserved = 0 + self._cursor = 0 + + # -- build phase ------------------------------------------------------ + def reserve(self, numel: int) -> None: + self._reserved += self._round(numel) + + def commit(self) -> None: + import torch.distributed._symmetric_memory as symm_mem + + symm_mem.enable_symm_mem_for_group(self.group_name) + self.buf = symm_mem.empty(self._reserved, dtype=self.dtype, device=self.device) + self.handle = symm_mem.rendezvous(self.buf, self.group_name) + # Slice each peer's whole window once. VMM maps every window into one + # VA space so these are directly copyable; NCCL-backend peer pointers are not. + self.peers = [self.handle.get_buffer(r, (self._reserved,), self.dtype) for r in range(self.handle.world_size)] + + def take(self, shape: torch.Size | tuple[int, ...]) -> torch.Tensor: + numel = 1 + for s in shape: + numel *= int(s) + off = self._cursor + self._cursor += self._round(numel) + if self._cursor > self._reserved: + raise RuntimeError( + f"symmetric arena overflow: wanted {self._cursor} elems, reserved {self._reserved}. " + "The sizing walk and the dispensing walk must visit the same shards in the same order." + ) + return self.buf[off : off + numel].view(shape) + + # -- query ------------------------------------------------------------ + @property + def nbytes(self) -> int: + return self._reserved * self.dtype.itemsize + + def offset_of(self, t: torch.Tensor) -> int: + return (t.data_ptr() - self.buf.data_ptr()) // self.buf.element_size() + + def peer_views(self, t: torch.Tensor) -> list[torch.Tensor]: + """``world_size`` views of the same shard, one per rank, in rank order.""" + off, numel = self.offset_of(t), t.numel() + return [p[off : off + numel].view(t.shape) for p in self.peers] + + def contains(self, t: torch.Tensor) -> bool: + if self.buf is None: + return False + base = self.buf.data_ptr() + return base <= t.data_ptr() < base + self.nbytes + + @classmethod + def _round(cls, numel: int) -> int: + return (numel + cls.ALIGN - 1) // cls.ALIGN * cls.ALIGN + + +@dataclass(frozen=True) +class ShardEntry: + """What the run-time gather needs to know about one local shard.""" + + arena: SymmArena + offset: int + local: torch.Tensor + peer_views: tuple[torch.Tensor, ...] + + @property + def shape(self) -> tuple[int, ...]: + return tuple(self.local.shape) + + +# Keyed by ``data_ptr()``: the gather op only sees a plain tensor. +_SHARD_REGISTRY: dict[int, ShardEntry] = {} +_ARENAS: list[SymmArena] = [] +_BARRIER_DONE = False + + +def register_shard(local: torch.Tensor, arena: SymmArena) -> ShardEntry: + entry = ShardEntry(arena=arena, offset=arena.offset_of(local), local=local, peer_views=tuple(arena.peer_views(local))) + _SHARD_REGISTRY[local.data_ptr()] = entry + return entry + + +def lookup_shard(data_ptr: int) -> ShardEntry | None: + return _SHARD_REGISTRY.get(data_ptr) + + +def registered_arenas() -> list[SymmArena]: + return list(_ARENAS) + + +def find_shard_by_layout(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor | None: + """Any registered shard with this layout -- the cost model cannot replay a gather on a generic ``empty`` (no peers).""" + want = tuple(int(s) for s in shape) + for entry in _SHARD_REGISTRY.values(): + if entry.shape == want and entry.local.dtype == dtype: + return entry.local + return None + + +def reset_registry() -> None: + """Test-only: drop every arena so a new model can be built in-process.""" + global _BARRIER_DONE + _SHARD_REGISTRY.clear() + _ARENAS.clear() + _BARRIER_DONE = False + + +def barrier_after_load() -> None: + """Publish every rank's freshly written shards, once per process. + + Must run after weights are loaded and before the first peer read. + """ + global _BARRIER_DONE + if _BARRIER_DONE or not _ARENAS: + return + if dist.is_available() and dist.is_initialized(): + torch.cuda.synchronize() + dist.barrier() + _BARRIER_DONE = True + magi_logger.info( + "Symmetric arena: published %d arena(s), %.1f MiB, %d shards; steady state is barrier-free", + len(_ARENAS), + sum(a.nbytes for a in _ARENAS) / 2**20, + len(_SHARD_REGISTRY), + ) + + +def _is_gatherable_shard(t: object) -> bool: + """A Shard(0) DTensor on a 1-D mesh -- the only placement the copy-engine gather handles.""" + from torch.distributed.tensor import DTensor, Shard + + if not isinstance(t, DTensor): + return False + placements = t.placements + return len(placements) == 1 and isinstance(placements[0], Shard) and placements[0].dim == 0 + + +def _group_name_of(t) -> str | None: + try: + return t.device_mesh._dim_group_names[0] + except Exception: # noqa: BLE001 + return None + + +def _arena_key(t) -> tuple[torch.dtype, str]: + """One window per (dtype, process group). Same dtype on two meshes (gaga4 FSDP + edp) must not share a window.""" + group_name = _group_name_of(t) + if group_name is None: + raise RuntimeError(f"cannot resolve the process group of a Shard(0) parameter on mesh {t.device_mesh}") + return (t.dtype, group_name) + + +def _apply_order_entries(mod: nn.Module): + """``(owner, name, param)`` in ``_apply`` order: post-order, every ``_parameters`` entry. + + Not ``named_parameters()`` (pre-order, dedups shared tensors). A different + walk would break cross-rank offset symmetry. Walking ``_parameters`` also + finds SimpleFSDP weights in ``parametrizations.weight.original``. + """ + for child in mod.children(): + yield from _apply_order_entries(child) + for name, p in mod._parameters.items(): + if p is not None: + yield mod, name, p + + +def _plan_arenas(shards: list, device: torch.device) -> dict[tuple[torch.dtype, str], SymmArena]: + """Size and commit one window per (dtype, group). Dedup by identity so a tied weight reserves a single slot.""" + arenas: dict[tuple[torch.dtype, str], SymmArena] = {} + seen: set[int] = set() + for p in shards: + if id(p) in seen: + continue + seen.add(id(p)) + key = _arena_key(p) + arena = arenas.get(key) + if arena is None: + arena = arenas[key] = SymmArena(p.dtype, device, key[1]) + arena.reserve(p._local_tensor.numel()) + + for arena in arenas.values(): + arena.commit() # the only collective, once per window + _ARENAS.extend(arenas.values()) + return arenas + + +def materialize_into_arenas(mod: nn.Module, device: torch.device) -> dict[tuple[torch.dtype, str], SymmArena]: + """Size windows for ``mod``'s Shard(0) shards while they are still on meta. Non-gatherable params are left to the caller.""" + shards = [p for _, _, p in _apply_order_entries(mod) if _is_gatherable_shard(p)] + if not shards: + return {} + return _plan_arenas(shards, device) + + +def migrate_to_arenas(root: nn.Module) -> dict[tuple[torch.dtype, str], SymmArena]: + """Copy already-allocated Shard(0) shards into symmetric memory. + + Used when ``magi_compile(model, ...)`` is given a live model rather than a + meta + ``to_empty`` path. ``load_state_dict(assign=True)`` after this would + replace arena views with ordinary tensors; the gather then rejects them. + """ + entries = [(m, n, p) for m, n, p in _apply_order_entries(root) if _is_gatherable_shard(p)] + if not entries: + return {} + + device = entries[0][2]._local_tensor.device + if device.type != "cuda": + raise RuntimeError(f"symmetric memory needs the shards on cuda, found {device}") + arenas = _plan_arenas([p for _, _, p in entries], device) + + from torch.distributed.tensor import DTensor + + views: dict[int, torch.Tensor] = {} + for owner, name, p in entries: + local = views.get(id(p)) + if local is None: + arena = arenas[_arena_key(p)] + local = views[id(p)] = arena.take(p._local_tensor.shape) + local.copy_(p._local_tensor) + register_shard(local, arena) + moved = DTensor.from_local(local, p.device_mesh, p.placements, run_check=False) + owner.register_parameter(name, nn.Parameter(moved, requires_grad=p.requires_grad)) + + magi_logger.info( + "Symmetric arena: migrated %d shard(s) into %.1f MiB across %d window(s)", + len(views), + sum(a.nbytes for a in arenas.values()) / 2**20, + len(arenas), + ) + return arenas + + +def patch_symm_arena_apply(cls: type[nn.Module]) -> None: + """Install the ``_apply`` interception on a decorated class. + + Mirrors ``_patch_cpu_offload_apply``: take over for ``to_empty``'s lambda, delegate everything else. + """ + if getattr(cls, "_magi_symm_apply_patched", False): + return + orig_apply = cls._apply + magi_logger.info("Symmetric arena: intercepting %s._apply for copy-engine FSDP", cls.__name__) + + def _symm_apply(self, fn, recurse: bool = True): + if getattr(fn, "__qualname__", "") != _TO_EMPTY_LAMBDA: + return orig_apply(self, fn, recurse) + if getattr(self, "_magi_symm_arenas", None) is not None: + return orig_apply(self, fn, recurse) + + device = torch.device(inspect.getclosurevars(fn).nonlocals["device"]) + from torch.distributed.tensor import DTensor + + arenas = materialize_into_arenas(self, device) + if not arenas: + return orig_apply(self, fn, recurse) + + views: dict[int, torch.Tensor] = {} + + def materialize(t: torch.Tensor) -> torch.Tensor: + if not _is_gatherable_shard(t): + return torch.empty_like(t, device=device) + # Tied weight: same view so tying survives materialization. + local = views.get(id(t)) + if local is None: + arena = arenas[_arena_key(t)] + local = views[id(t)] = arena.take(t._local_tensor.shape) + register_shard(local, arena) + return DTensor.from_local(local, t.device_mesh, t.placements, run_check=False) + + # Do not forge to_empty's qualname: a nested decorated block must fail + # the check above and delegate, so its params land in *this* arena. + out = orig_apply(self, materialize, recurse) + self._magi_symm_arenas = arenas + magi_logger.info( + "Symmetric arena: %s materialized %d shard(s) into %.1f MiB across %d window(s)", + cls.__name__, + len(views), + sum(a.nbytes for a in arenas.values()) / 2**20, + len(arenas), + ) + return out + + cls._apply = _symm_apply + cls._magi_symm_apply_patched = True From 789b83667740f9094d24cfe70852145208fa1124 Mon Sep 17 00:00:00 2001 From: wtr Date: Fri, 28 Aug 2026 10:51:20 +0800 Subject: [PATCH 05/16] chore --- tests/feature_tests/symm_helper/__init__.py | 13 + .../symm_helper/verify_symm_e2e.py | 260 ++++++++++++++++++ tests/feature_tests/test_symm_e2e.py | 7 +- 3 files changed, 277 insertions(+), 3 deletions(-) create mode 100644 tests/feature_tests/symm_helper/__init__.py create mode 100644 tests/feature_tests/symm_helper/verify_symm_e2e.py diff --git a/tests/feature_tests/symm_helper/__init__.py b/tests/feature_tests/symm_helper/__init__.py new file mode 100644 index 0000000..3eaa44a --- /dev/null +++ b/tests/feature_tests/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/symm_helper/verify_symm_e2e.py b/tests/feature_tests/symm_helper/verify_symm_e2e.py new file mode 100644 index 0000000..38de7ad --- /dev/null +++ b/tests/feature_tests/symm_helper/verify_symm_e2e.py @@ -0,0 +1,260 @@ +# 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 arena 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::symm_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 + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.device_mesh import init_device_mesh + +from magi_compiler import magi_compile +from magi_compiler.config import CompileMode, CudaGraphMode + + +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 arena 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) placement: the block's shards are in a window, the head's are not. + from magi_compiler.symm_mem import lookup_shard, registered_arenas + + block_names = {n for n, _ in _named_shards(model) if n.startswith("block.")} + in_arena = [n for n, p in _named_shards(model) if lookup_shard(p._local_tensor.data_ptr()) is not None] + head_in_arena = [n for n in in_arena if not n.startswith("block.")] + arena_mib = sum(a.nbytes for a in registered_arenas()) / 2**20 + placement_ok = (set(in_arena) == block_names) if ce else (in_arena == []) + log( + f"CHECK placement: {len(in_arena)}/{len(block_names)} block shards in {len(registered_arenas())} " + f"window(s) ({arena_mib:.0f} MiB), {len(head_in_arena)} stray -> {'ok' if placement_ok else 'WRONG'}" + ) + + # (2) + (3): 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 + + 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::symm_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 arena={arena_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/test_symm_e2e.py b/tests/feature_tests/test_symm_e2e.py index a233229..31ef817 100644 --- a/tests/feature_tests/test_symm_e2e.py +++ b/tests/feature_tests/test_symm_e2e.py @@ -19,8 +19,9 @@ 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 example script asserts. Driving it here as a -subprocess keeps that assertion in CI instead of in someone's shell history. +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 @@ -31,7 +32,7 @@ import pytest import torch -_SCRIPT = Path(__file__).resolve().parents[2] / "example" / "inference" / "fsdp_overlap" / "verify_symm_e2e.py" +_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") From 4c46ccb21286ddb4167f93e8ff1068ea605f2034 Mon Sep 17 00:00:00 2001 From: wtr Date: Fri, 28 Aug 2026 18:59:21 +0800 Subject: [PATCH 06/16] chore --- .../fsdp_overlap/redistribute_lowering.py | 53 ++++++++----- .../cache/test_cache_topology_isolation.py | 3 +- .../fsdp_overlap_helper/reorder_helper.py | 23 +++++- .../fsdp/test_fsdp_overlap_lowering.py | 75 +++++++++++++++++-- .../fsdp/test_fsdp_overlap_reorder.py | 11 +++ .../symm_helper/verify_symm_e2e.py | 24 +++++- 6 files changed, 161 insertions(+), 28 deletions(-) diff --git a/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py index 8a5e262..cd5e6a7 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 @@ -33,6 +34,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 +76,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 +101,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 +142,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 @@ -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/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 b3099b0..a3b0477 100644 --- a/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py +++ b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py @@ -95,6 +95,16 @@ def __init__(self) -> None: self.node = _FakeIR() +def _is_driver_refusal(e: Exception) -> bool: + """True when the CUDA driver, not our code, rejected the symmetric-memory setup. + + Anything else -- a shape error, an arena overflow, a missing registration -- is a + real failure and must not be mistaken for an unequipped host. + """ + msg = str(e) + return "CUDA driver error" in msg or "not supported" in msg + + def _mode_ladder_selfcheck(rank: int) -> bool: """Assert every rung of ``_negotiate_mode``'s ladder, with rank 1 feeding the divergent input. All ranks walk the cases in the same order, so the symmetric @@ -173,7 +183,18 @@ def fn(x, w0, shard): arena = SymmArena(torch.bfloat16, torch.device("cuda", dev), grp) for _ in range(N_CE_LAYERS): arena.reserve(H * H) - arena.commit() + try: + arena.commit() + except RuntimeError as e: + # A symmetric window needs an initialized NVLink fabric -- on NVSwitch hosts a + # running nvidia-fabricmanager. Where it is not, the driver refuses the + # rendezvous on every rank before any of the pass under test runs, so report + # the host as unable rather than the pass as broken. + if not _is_driver_refusal(e): + raise + print(f"REORDER_SYMM_UNAVAILABLE rank={rank} {e}", flush=True) + dist.destroy_process_group() + raise SystemExit(0) from None for i in range(N_CE_LAYERS): s = arena.take((H, H)) s.normal_(0.0, H**-0.5).add_(0.01 * i) diff --git a/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py b/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py index 9374f18..de3e72d 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): @@ -130,6 +158,37 @@ 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 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 x.meta.get("magi_fsdp_weight_ag")]) == 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_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 61b7c57..d87b6d5 100644 --- a/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py +++ b/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py @@ -177,6 +177,16 @@ def test_reorder_graph_mismatch_slot_mode(): assert "REORDER_PASS" in p.stdout, out[-3000:] +def _skip_if_symm_mem_unusable(out: str) -> None: + """Symmetric memory needs an initialized NVLink fabric on the host (on NVSwitch + machines, a running nvidia-fabricmanager). Where the driver refuses the rendezvous + the copy-engine transport cannot run at all, so there is nothing here to assert on -- + the helper says so explicitly and only for a driver-level refusal.""" + line = next((ln for ln in out.splitlines() if "REORDER_SYMM_UNAVAILABLE" in ln), None) + if line is not None: + pytest.skip(f"symmetric memory unusable on this host: {line.strip()}") + + @requires_cuda @requires_torchrun @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") @@ -190,6 +200,7 @@ def test_reorder_copy_engine_recognition(): """ p = _run(2, "--copy-engine", port="29634") out = p.stdout + p.stderr + _skip_if_symm_mem_unusable(out) 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 diff --git a/tests/feature_tests/symm_helper/verify_symm_e2e.py b/tests/feature_tests/symm_helper/verify_symm_e2e.py index 38de7ad..0436936 100644 --- a/tests/feature_tests/symm_helper/verify_symm_e2e.py +++ b/tests/feature_tests/symm_helper/verify_symm_e2e.py @@ -142,6 +142,16 @@ def _named_shards(model: nn.Module): yield (f"{mod_name}.{p_name}" if mod_name else p_name), p +def _is_driver_refusal(e: Exception) -> bool: + """True when the CUDA driver, not our code, rejected the symmetric-memory setup. + + Anything else -- a shape error, an arena overflow, a missing registration -- is a + real failure and must not be mistaken for an unequipped host. + """ + msg = str(e) + return "CUDA driver error" in msg or "not supported" in msg + + def main() -> None: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--hidden", type=int, default=4096) @@ -199,7 +209,19 @@ def spy_rewrite(graph): 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) + try: + model.to_empty(device=device) + except RuntimeError as e: + # to_empty is where the patched _apply opens the symmetric window, and a window + # needs an initialized NVLink fabric -- on NVSwitch hosts a running + # nvidia-fabricmanager. Where there is none the driver refuses the rendezvous on + # every rank before any of the chain under test runs, so report the host as + # unable rather than the chain as broken. + if not (ce and _is_driver_refusal(e)): + raise + print(f"E2E_SYMM_UNAVAILABLE rank={rank} {e}", flush=True) + dist.destroy_process_group() + raise SystemExit(0) from None _load_into_shards(model, state) # (1) placement: the block's shards are in a window, the head's are not. From 1dcec05e549aee6396bc359de1caedef2fb361c1 Mon Sep 17 00:00:00 2001 From: wtr Date: Fri, 28 Aug 2026 21:54:34 +0800 Subject: [PATCH 07/16] chore --- tests/feature_tests/test_symm_e2e.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/feature_tests/test_symm_e2e.py b/tests/feature_tests/test_symm_e2e.py index 31ef817..12cb3a4 100644 --- a/tests/feature_tests/test_symm_e2e.py +++ b/tests/feature_tests/test_symm_e2e.py @@ -65,6 +65,16 @@ def _run(transport: str, port: str) -> subprocess.CompletedProcess: ) +def _skip_if_symm_mem_unusable(out: str) -> None: + """Symmetric memory needs an initialized NVLink fabric on the host (on NVSwitch + machines, a running nvidia-fabricmanager). Where the driver refuses the rendezvous + the copy-engine transport cannot run at all, so there is nothing here to assert on -- + the helper says so explicitly and only for a driver-level refusal.""" + line = next((ln for ln in out.splitlines() if "E2E_SYMM_UNAVAILABLE" in ln), None) + if line is not None: + pytest.skip(f"symmetric memory unusable on this host: {line.strip()}") + + @requires_2gpu @requires_torchrun def test_copy_engine_end_to_end(): @@ -73,6 +83,7 @@ def test_copy_engine_end_to_end(): is where a missing wait would show up.""" p = _run("copy_engine", "29641") out = p.stdout + p.stderr + _skip_if_symm_mem_unusable(out) 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:] From 3eda0a72f1238c5f5ab5407544c9bd2c155714eb Mon Sep 17 00:00:00 2001 From: wtr Date: Sat, 29 Aug 2026 17:38:02 +0800 Subject: [PATCH 08/16] [Fix] Disable symm-mem multicast in copy-engine torchrun helpers so they run on hosts without an initialized NVLink fabric --- .../fsdp_overlap_helper/reorder_helper.py | 31 +++------------ .../fsdp/test_fsdp_overlap_reorder.py | 11 ------ .../symm_helper/verify_symm_e2e.py | 38 +++++-------------- tests/feature_tests/test_symm_e2e.py | 11 ------ 4 files changed, 15 insertions(+), 76 deletions(-) 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 a3b0477..f98e436 100644 --- a/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py +++ b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py @@ -69,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 @@ -95,16 +97,6 @@ def __init__(self) -> None: self.node = _FakeIR() -def _is_driver_refusal(e: Exception) -> bool: - """True when the CUDA driver, not our code, rejected the symmetric-memory setup. - - Anything else -- a shape error, an arena overflow, a missing registration -- is a - real failure and must not be mistaken for an unequipped host. - """ - msg = str(e) - return "CUDA driver error" in msg or "not supported" in msg - - def _mode_ladder_selfcheck(rank: int) -> bool: """Assert every rung of ``_negotiate_mode``'s ladder, with rank 1 feeding the divergent input. All ranks walk the cases in the same order, so the symmetric @@ -183,18 +175,7 @@ def fn(x, w0, shard): arena = SymmArena(torch.bfloat16, torch.device("cuda", dev), grp) for _ in range(N_CE_LAYERS): arena.reserve(H * H) - try: - arena.commit() - except RuntimeError as e: - # A symmetric window needs an initialized NVLink fabric -- on NVSwitch hosts a - # running nvidia-fabricmanager. Where it is not, the driver refuses the - # rendezvous on every rank before any of the pass under test runs, so report - # the host as unable rather than the pass as broken. - if not _is_driver_refusal(e): - raise - print(f"REORDER_SYMM_UNAVAILABLE rank={rank} {e}", flush=True) - dist.destroy_process_group() - raise SystemExit(0) from None + arena.commit() for i in range(N_CE_LAYERS): s = arena.take((H, H)) s.normal_(0.0, H**-0.5).add_(0.01 * i) diff --git a/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py b/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py index d87b6d5..61b7c57 100644 --- a/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py +++ b/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py @@ -177,16 +177,6 @@ def test_reorder_graph_mismatch_slot_mode(): assert "REORDER_PASS" in p.stdout, out[-3000:] -def _skip_if_symm_mem_unusable(out: str) -> None: - """Symmetric memory needs an initialized NVLink fabric on the host (on NVSwitch - machines, a running nvidia-fabricmanager). Where the driver refuses the rendezvous - the copy-engine transport cannot run at all, so there is nothing here to assert on -- - the helper says so explicitly and only for a driver-level refusal.""" - line = next((ln for ln in out.splitlines() if "REORDER_SYMM_UNAVAILABLE" in ln), None) - if line is not None: - pytest.skip(f"symmetric memory unusable on this host: {line.strip()}") - - @requires_cuda @requires_torchrun @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 GPUs") @@ -200,7 +190,6 @@ def test_reorder_copy_engine_recognition(): """ p = _run(2, "--copy-engine", port="29634") out = p.stdout + p.stderr - _skip_if_symm_mem_unusable(out) 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 diff --git a/tests/feature_tests/symm_helper/verify_symm_e2e.py b/tests/feature_tests/symm_helper/verify_symm_e2e.py index 0436936..142cef3 100644 --- a/tests/feature_tests/symm_helper/verify_symm_e2e.py +++ b/tests/feature_tests/symm_helper/verify_symm_e2e.py @@ -50,13 +50,15 @@ import argparse import os -import torch -import torch.distributed as dist -import torch.nn as nn -from torch.distributed.device_mesh import init_device_mesh +os.environ.setdefault("TORCH_SYMM_MEM_DISABLE_MULTICAST", "1") -from magi_compiler import magi_compile -from magi_compiler.config import CompileMode, CudaGraphMode +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: @@ -142,16 +144,6 @@ def _named_shards(model: nn.Module): yield (f"{mod_name}.{p_name}" if mod_name else p_name), p -def _is_driver_refusal(e: Exception) -> bool: - """True when the CUDA driver, not our code, rejected the symmetric-memory setup. - - Anything else -- a shape error, an arena overflow, a missing registration -- is a - real failure and must not be mistaken for an unequipped host. - """ - msg = str(e) - return "CUDA driver error" in msg or "not supported" in msg - - def main() -> None: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--hidden", type=int, default=4096) @@ -209,19 +201,7 @@ def spy_rewrite(graph): 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") - try: - model.to_empty(device=device) - except RuntimeError as e: - # to_empty is where the patched _apply opens the symmetric window, and a window - # needs an initialized NVLink fabric -- on NVSwitch hosts a running - # nvidia-fabricmanager. Where there is none the driver refuses the rendezvous on - # every rank before any of the chain under test runs, so report the host as - # unable rather than the chain as broken. - if not (ce and _is_driver_refusal(e)): - raise - print(f"E2E_SYMM_UNAVAILABLE rank={rank} {e}", flush=True) - dist.destroy_process_group() - raise SystemExit(0) from None + model.to_empty(device=device) _load_into_shards(model, state) # (1) placement: the block's shards are in a window, the head's are not. diff --git a/tests/feature_tests/test_symm_e2e.py b/tests/feature_tests/test_symm_e2e.py index 12cb3a4..31ef817 100644 --- a/tests/feature_tests/test_symm_e2e.py +++ b/tests/feature_tests/test_symm_e2e.py @@ -65,16 +65,6 @@ def _run(transport: str, port: str) -> subprocess.CompletedProcess: ) -def _skip_if_symm_mem_unusable(out: str) -> None: - """Symmetric memory needs an initialized NVLink fabric on the host (on NVSwitch - machines, a running nvidia-fabricmanager). Where the driver refuses the rendezvous - the copy-engine transport cannot run at all, so there is nothing here to assert on -- - the helper says so explicitly and only for a driver-level refusal.""" - line = next((ln for ln in out.splitlines() if "E2E_SYMM_UNAVAILABLE" in ln), None) - if line is not None: - pytest.skip(f"symmetric memory unusable on this host: {line.strip()}") - - @requires_2gpu @requires_torchrun def test_copy_engine_end_to_end(): @@ -83,7 +73,6 @@ def test_copy_engine_end_to_end(): is where a missing wait would show up.""" p = _run("copy_engine", "29641") out = p.stdout + p.stderr - _skip_if_symm_mem_unusable(out) 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:] From bf04b18ff718db6dc62edbf0fd7cbd902c8d977e Mon Sep 17 00:00:00 2001 From: wtr Date: Thu, 3 Sep 2026 14:45:41 +0800 Subject: [PATCH 09/16] [Fix] Make copy-engine eligibility rank-identical for uneven Shard(0) weights --- .../passes/fsdp_overlap/bucket_all_gather.py | 21 ++++-- .../passes/fsdp_overlap/lower_and_bucket.py | 11 +-- .../fsdp_overlap/redistribute_lowering.py | 6 ++ .../passes/fsdp_overlap/symm_ag_rewrite.py | 10 ++- magi_compiler/symm_mem/arena.py | 19 ++++- .../fsdp/test_fsdp_overlap_bucket.py | 43 +++++++++++- .../fsdp/test_fsdp_overlap_lowering.py | 20 ++++++ tests/feature_tests/test_symm_ag_rewrite.py | 69 ++++++++++++++++++- 8 files changed, 181 insertions(+), 18 deletions(-) diff --git a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py index 24a6dab..5d15852 100644 --- a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py +++ b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py @@ -70,13 +70,23 @@ def _is_weight_all_gather(node: fx.Node) -> bool: 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 -- the gathered size divided + by the world, NOT the local shard's own meta. + + Used to cap coalesced bucket size, so it has to be rank-identical: ranks that cut + a bucket at different members submit coalesced launches with different membership, + which never completes. Reading the gather's INPUT meta happens to be safe for a + graph this repo lowered -- the trailing ranks of an uneven ``Shard(0)`` own fewer + rows, but the pad in front of the gather brings them back to ``chunk`` -- and that + is far too subtle a thing to rest a collective on. A gather matched by + ``_gathers_a_weight`` rather than emitted by the lowering has no such pad. The + gather's own ``example_value`` is ``(world * chunk, ...)`` on every rank + unconditionally.""" + 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( @@ -159,6 +169,7 @@ def _coalesce_one_bucket(graph: fx.GraphModule, node_index: dict[fx.Node, int], coalesced.meta["example_value"] = list(ag_metas) coalesced.meta["magi_fsdp_weight_ag"] = True coalesced.meta["magi_fsdp_weight_ag_coalesced"] = True + coalesced.meta["magi_fsdp_uneven_shard"] = any(ag.meta.get("magi_fsdp_uneven_shard") for ag in ag_nodes) outs = [] for i, am in enumerate(ag_metas): diff --git a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py index b9938f6..cac068d 100644 --- a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py +++ b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py @@ -39,10 +39,13 @@ def lower_and_bucket_full_graph( the byte cap in program order (see ``bucket_weight_all_gather_coalesced``). 0 = no cap (one bucket per (group, dtype) run). - ``transport="copy_engine"`` buckets *first* (only arena-shard gathers, so - cast/pad stays out of the bucket), then retargets both the leftover singles - and the coalesced launches at the copy-engine ops. The wrapper still runs - one gather per member; reorder just sees one comm node per bucket. + ``transport="copy_engine"`` buckets *first* (only arena-shard gathers, so a cast + of the shard and an unevenly split weight stay out of the bucket), then retargets + both the leftover singles and the coalesced launches at the copy-engine ops. The + wrapper still runs one gather per member; reorder just sees one comm node per + bucket. Bucket membership therefore depends on the eligibility predicate, which + is why that predicate has to answer identically on every rank -- see + ``symm_ag_rewrite._is_uneven_shard``. Returns the number of buckets created. """ diff --git a/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py index cd5e6a7..e31f363 100644 --- a/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py +++ b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py @@ -176,6 +176,12 @@ def lower_prim_redistribute_to_collectives(graph: fx.GraphModule) -> int: # 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 + # Whether this Shard(0) divides evenly, recorded from F and world -- both + # the same on every rank. Downstream transport choices must key off THIS + # and never off `L`: L is what makes the pad above appear on the trailing + # ranks only, so a predicate that reads the pad splits one collective into + # copy-engine on some ranks and NCCL on others, which never completes. + ag.meta["magi_fsdp_uneven_shard"] = 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) diff --git a/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py b/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py index 1871cfb..a566029 100644 --- a/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py +++ b/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py @@ -23,8 +23,14 @@ _ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default +def _is_uneven_shard(node: fx.Node) -> bool: + return bool(node.meta.get("magi_fsdp_uneven_shard")) + + def _input_is_arena_shard(node: fx.Node) -> bool: """True when the gather input is ``to_local(placeholder/get_attr)`` with no intervening op.""" + if _is_uneven_shard(node): + return False src = node.args[0] if node.args else None if not isinstance(src, fx.Node) or src.op != "call_method" or src.target != "to_local": return False @@ -33,6 +39,8 @@ def _input_is_arena_shard(node: fx.Node) -> bool: def _coalesced_inputs_are_arena_shards(node: fx.Node) -> bool: + if _is_uneven_shard(node): + return False locs = node.args[0] if node.args else None if not isinstance(locs, (list, tuple)) or not locs: return False @@ -75,7 +83,7 @@ def rewrite_weight_ag_to_copy_engine(graph: fx.GraphModule) -> int: graph.recompile() magi_logger.info( "FSDP copy-engine rewrite: %d weight all-gather(s) retargeted, " - "%d left on NCCL (input is a cast/pad of the shard, not the shard itself)", + "%d left on NCCL (input is a cast of the shard, or the shard is unevenly split)", rewritten, skipped, ) diff --git a/magi_compiler/symm_mem/arena.py b/magi_compiler/symm_mem/arena.py index a8266cb..a991c89 100644 --- a/magi_compiler/symm_mem/arena.py +++ b/magi_compiler/symm_mem/arena.py @@ -172,13 +172,28 @@ def barrier_after_load() -> None: def _is_gatherable_shard(t: object) -> bool: - """A Shard(0) DTensor on a 1-D mesh -- the only placement the copy-engine gather handles.""" + """An EVENLY split Shard(0) DTensor on a 1-D mesh -- the only placement the + copy-engine gather handles. + + Uneven ``Shard(0)`` (``dim0 % world != 0``) is excluded on purpose, for three + reasons that all trace back to the trailing ranks owning fewer rows: the window + would be sized from a rank-dependent local numel, so ``rendezvous`` rejects it + outright once the difference survives ``ALIGN``; below that threshold it is worse + than an error, because every shard after the uneven one lands at a different + offset per rank while ``peer_views`` slices each peer at the LOCAL offset; and the + gather copies one fixed-size slab per peer, which would read past a shorter + peer's rows. ``dim0`` and the mesh size are identical on every rank, so all + ranks drop the same parameters and the offset walk stays symmetric. Dropped + weights keep their ordinary allocation and gather over NCCL. + """ from torch.distributed.tensor import DTensor, Shard if not isinstance(t, DTensor): return False placements = t.placements - return len(placements) == 1 and isinstance(placements[0], Shard) and placements[0].dim == 0 + if not (len(placements) == 1 and isinstance(placements[0], Shard) and placements[0].dim == 0): + return False + return int(t.shape[0]) % int(t.device_mesh.size(0)) == 0 def _group_name_of(t) -> str | None: diff --git a/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py b/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py index 4d3d00a..6272474 100644 --- a/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py +++ b/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py @@ -43,19 +43,24 @@ 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. """ 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 = [] @@ -145,6 +150,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 de3e72d..d7babbf 100644 --- a/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py +++ b/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py @@ -189,6 +189,26 @@ def test_lowering_reads_forward_dtype_from_method_kwargs(dist_1rank): 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 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 "magi_fsdp_uneven_shard" in ag[0].meta + assert ag[0].meta["magi_fsdp_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/test_symm_ag_rewrite.py b/tests/feature_tests/test_symm_ag_rewrite.py index 8d0253e..7cc66f0 100644 --- a/tests/feature_tests/test_symm_ag_rewrite.py +++ b/tests/feature_tests/test_symm_ag_rewrite.py @@ -69,12 +69,17 @@ def _symm_op(): return SYMM_ALL_GATHER -def _graph_with_gathers(mesh, n: int, *, derive: str | None = None, marked: bool = True, shapes: list[tuple] | None = None): +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. + ``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 @@ -101,6 +106,7 @@ def _graph_with_gathers(mesh, n: int, *, derive: str | None = None, marked: bool ag.meta["example_value"] = local._local_tensor.new_empty(local._local_tensor.shape) if marked: ag.meta["magi_fsdp_weight_ag"] = True + ag.meta["magi_fsdp_uneven_shard"] = uneven outs.append(g.call_function(_WAIT, (ag,))) g.output(tuple(outs)) return fx.GraphModule(torch.nn.Module(), g) @@ -142,7 +148,13 @@ def test_group_args_are_preserved(mesh_1rank): @pytest.mark.parametrize("derive", ["cast", "pad"]) def test_derived_shards_stay_on_nccl(mesh_1rank, derive): """A cast or pad output is not in the symmetric window; retargeting it would - be rejected at run time, so it must stay on NCCL.""" + be rejected at run time, so it must stay on NCCL. + + 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_shard_stays_on_nccl_without_a_pad``. + """ from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine gm = _graph_with_gathers(mesh_1rank, 3, derive=derive) @@ -151,6 +163,57 @@ def test_derived_shards_stay_on_nccl(mesh_1rank, derive): assert not _gathers(gm, _symm_op()) +@requires_cuda +def test_uneven_shard_stays_on_nccl_without_a_pad(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 is made from the flag, which is derived from F and world and is + therefore the same everywhere. + """ + from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine + + gm = _graph_with_gathers(mesh_1rank, 3, uneven=True) + 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_uneven_shard_does_not_join_a_copy_engine_bucket(mesh_1rank): + """Bucket membership is a transport decision too: a bucket is one submission, and + an uneven weight inside a copy-engine bucket would take the whole bucket with it.""" + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED + + gm = _graph_with_gathers(mesh_1rank, 3, uneven=True) + assert lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine") == 0 + assert len(_gathers(gm, _AG)) == 3 + assert not _gathers(gm, SYMM_ALL_GATHER_COALESCED) + + +@requires_cuda +def test_uneven_shard_does_not_hold_back_its_even_neighbours(mesh_1rank): + """Only the uneven weight loses the copy engine. Excluding it from the bucket + must not split the even weights around it into separate buckets either, or the + fix would cost throughput on every model with one odd-shaped weight.""" + from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED + + gm = _graph_with_gathers(mesh_1rank, 4) + _gathers(gm, _AG)[1].meta["magi_fsdp_uneven_shard"] = True + + assert lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine") == 1 + coalesced = _gathers(gm, SYMM_ALL_GATHER_COALESCED) + assert len(coalesced) == 1 + assert len(coalesced[0].args[0]) == 3 # the three even weights, in one bucket + assert len(_gathers(gm, _AG)) == 1 # the uneven one, alone, on NCCL + + @requires_cuda def test_unmarked_gathers_are_left_alone(mesh_1rank): """CP / TP / MoE gathers are never marked, and must survive copy-engine mode From 079ef1144a7a6bc03aff8adfb0dc917ad39880e0 Mon Sep 17 00:00:00 2001 From: wtr Date: Thu, 3 Sep 2026 17:28:25 +0800 Subject: [PATCH 10/16] Replace all arena with Buffer --- magi_compiler/_api.py | 8 +- .../passes/fsdp_overlap/lower_and_bucket.py | 6 +- .../passes/fsdp_overlap/symm_ag_rewrite.py | 12 +- magi_compiler/symm_mem/__init__.py | 22 +-- magi_compiler/symm_mem/all_gather.py | 6 +- .../symm_mem/{arena.py => symm_buffer.py} | 102 +++++------ .../fsdp_overlap_helper/reorder_helper.py | 16 +- .../fsdp/test_fsdp_overlap_reorder.py | 2 +- .../fsdp/test_profiling_estimator.py | 14 +- .../symm_helper/verify_symm_e2e.py | 20 +-- tests/feature_tests/test_symm_ag_rewrite.py | 8 +- tests/feature_tests/test_symm_all_gather.py | 34 ++-- ...test_symm_arena.py => test_symm_buffer.py} | 158 +++++++++--------- tests/feature_tests/test_symm_e2e.py | 2 +- 14 files changed, 205 insertions(+), 205 deletions(-) rename magi_compiler/symm_mem/{arena.py => symm_buffer.py} (79%) rename tests/feature_tests/{test_symm_arena.py => test_symm_buffer.py} (79%) diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 2018c2b..1077f0c 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -224,9 +224,9 @@ def _magi_compile_class( _patch_cpu_offload_apply(cls, conf) if issubclass(cls, nn.Module) and conf.fsdp_config.transport == "copy_engine": - from magi_compiler.symm_mem import patch_symm_arena_apply + from magi_compiler.symm_mem import patch_symm_buffer_apply - patch_symm_arena_apply(cls) + patch_symm_buffer_apply(cls) old_init = cls.__init__ @@ -252,9 +252,9 @@ def _magi_compile_bound_method( return instance if conf.fsdp_config.transport == "copy_engine" and isinstance(instance, nn.Module): - from magi_compiler.symm_mem import migrate_to_arenas + from magi_compiler.symm_mem import migrate_to_buffers - migrate_to_arenas(instance) + migrate_to_buffers(instance) old_method = getattr(instance, method_name) diff --git a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py index cac068d..39c7b93 100644 --- a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py +++ b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py @@ -39,7 +39,7 @@ def lower_and_bucket_full_graph( the byte cap in program order (see ``bucket_weight_all_gather_coalesced``). 0 = no cap (one bucket per (group, dtype) run). - ``transport="copy_engine"`` buckets *first* (only arena-shard gathers, so a cast + ``transport="copy_engine"`` buckets *first* (only SymmBuffer-shard gathers, so a cast of the shard and an unevenly split weight stay out of the bucket), then retargets both the leftover singles and the coalesced launches at the copy-engine ops. The wrapper still runs one gather per member; reorder just sees one comm node per @@ -55,9 +55,9 @@ def lower_and_bucket_full_graph( bucket_mode = (bucket_mode or "none").lower() n = 0 if bucket_mode == "coalesced": - from .symm_ag_rewrite import _input_is_arena_shard + from .symm_ag_rewrite import _input_is_symm_shard - eligible = _input_is_arena_shard if transport == "copy_engine" else None + eligible = _input_is_symm_shard if transport == "copy_engine" else None n = bucket_weight_all_gather_coalesced(graph, bucket_size_bytes=bucket_size_bytes, eligible=eligible) magi_logger.info("Whole-graph FSDP bucketing (%s): created %d buckets", bucket_mode, n) elif bucket_mode not in ("none", ""): diff --git a/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py b/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py index a566029..fabb112 100644 --- a/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py +++ b/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py @@ -27,7 +27,7 @@ def _is_uneven_shard(node: fx.Node) -> bool: return bool(node.meta.get("magi_fsdp_uneven_shard")) -def _input_is_arena_shard(node: fx.Node) -> bool: +def _input_is_symm_shard(node: fx.Node) -> bool: """True when the gather input is ``to_local(placeholder/get_attr)`` with no intervening op.""" if _is_uneven_shard(node): return False @@ -38,16 +38,16 @@ def _input_is_arena_shard(node: fx.Node) -> bool: return isinstance(owner, fx.Node) and owner.op in ("placeholder", "get_attr") -def _coalesced_inputs_are_arena_shards(node: fx.Node) -> bool: +def _coalesced_inputs_are_symm_shards(node: fx.Node) -> bool: if _is_uneven_shard(node): return False locs = node.args[0] if node.args else None if not isinstance(locs, (list, tuple)) or not locs: return False - return all(isinstance(loc, fx.Node) and _input_is_arena_shard_from_local(loc) for loc in locs) + return all(isinstance(loc, fx.Node) and _input_is_symm_shard_from_local(loc) for loc in locs) -def _input_is_arena_shard_from_local(src: fx.Node) -> bool: +def _input_is_symm_shard_from_local(src: fx.Node) -> bool: if src.op != "call_method" or src.target != "to_local": return False owner = src.args[0] if src.args else None @@ -66,13 +66,13 @@ def rewrite_weight_ag_to_copy_engine(graph: fx.GraphModule) -> int: if not node.meta.get("magi_fsdp_weight_ag"): continue if node.target is _ALL_GATHER: - if not _input_is_arena_shard(node): + if not _input_is_symm_shard(node): skipped += 1 continue node.target = SYMM_ALL_GATHER rewritten += 1 elif node.target is _ALL_GATHER_COALESCED: - if not _coalesced_inputs_are_arena_shards(node): + if not _coalesced_inputs_are_symm_shards(node): skipped += 1 continue node.target = SYMM_ALL_GATHER_COALESCED diff --git a/magi_compiler/symm_mem/__init__.py b/magi_compiler/symm_mem/__init__.py index 8222d17..5d501dd 100644 --- a/magi_compiler/symm_mem/__init__.py +++ b/magi_compiler/symm_mem/__init__.py @@ -19,30 +19,30 @@ whether the copy-engine transport is available. Import that module by path. """ -from .arena import ( +from .symm_buffer import ( ShardEntry, - SymmArena, + SymmBuffer, barrier_after_load, find_shard_by_layout, lookup_shard, - materialize_into_arenas, - migrate_to_arenas, - patch_symm_arena_apply, + materialize_into_buffers, + migrate_to_buffers, + patch_symm_buffer_apply, register_shard, - registered_arenas, + registered_buffers, reset_registry, ) __all__ = [ "ShardEntry", - "SymmArena", + "SymmBuffer", "barrier_after_load", "find_shard_by_layout", "lookup_shard", - "materialize_into_arenas", - "migrate_to_arenas", - "patch_symm_arena_apply", + "materialize_into_buffers", + "migrate_to_buffers", + "patch_symm_buffer_apply", "register_shard", - "registered_arenas", + "registered_buffers", "reset_registry", ] diff --git a/magi_compiler/symm_mem/all_gather.py b/magi_compiler/symm_mem/all_gather.py index 9ecc712..405ef72 100644 --- a/magi_compiler/symm_mem/all_gather.py +++ b/magi_compiler/symm_mem/all_gather.py @@ -22,12 +22,12 @@ from magi_compiler.utils import magi_logger -from .arena import lookup_shard +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 arena. +# only reason ``group_name`` is here: the copy engine reads peers from the SymmBuffer. _SCHEMA = "symm_all_gather(Tensor local, int group_size, str group_name) -> Tensor" _SCHEMA_COALESCED = "symm_all_gather_coalesced(Tensor[] shards, int group_size, str group_name) -> Tensor[]" @@ -134,7 +134,7 @@ def _shard_peers(local: torch.Tensor, group_size: int) -> tuple[torch.Tensor, .. if entry is None: raise RuntimeError( "magi::symm_all_gather got a tensor that is not a registered symmetric-memory shard. " - "Only weights materialized through the arena can be gathered by the copy engine; " + "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 diff --git a/magi_compiler/symm_mem/arena.py b/magi_compiler/symm_mem/symm_buffer.py similarity index 79% rename from magi_compiler/symm_mem/arena.py rename to magi_compiler/symm_mem/symm_buffer.py index a991c89..8e1bf92 100644 --- a/magi_compiler/symm_mem/arena.py +++ b/magi_compiler/symm_mem/symm_buffer.py @@ -28,7 +28,7 @@ _TO_EMPTY_LAMBDA = "Module.to_empty.." -class SymmArena: +class SymmBuffer: """One symmetric-memory window, suballocated to many weight shards. One window per (decorated block, dtype, process group): a single rendezvous, @@ -71,7 +71,7 @@ def take(self, shape: torch.Size | tuple[int, ...]) -> torch.Tensor: self._cursor += self._round(numel) if self._cursor > self._reserved: raise RuntimeError( - f"symmetric arena overflow: wanted {self._cursor} elems, reserved {self._reserved}. " + f"symmetric buffer overflow: wanted {self._cursor} elems, reserved {self._reserved}. " "The sizing walk and the dispensing walk must visit the same shards in the same order." ) return self.buf[off : off + numel].view(shape) @@ -104,7 +104,7 @@ def _round(cls, numel: int) -> int: class ShardEntry: """What the run-time gather needs to know about one local shard.""" - arena: SymmArena + buffer: SymmBuffer offset: int local: torch.Tensor peer_views: tuple[torch.Tensor, ...] @@ -116,12 +116,12 @@ def shape(self) -> tuple[int, ...]: # Keyed by ``data_ptr()``: the gather op only sees a plain tensor. _SHARD_REGISTRY: dict[int, ShardEntry] = {} -_ARENAS: list[SymmArena] = [] +_BUFFERS: list[SymmBuffer] = [] _BARRIER_DONE = False -def register_shard(local: torch.Tensor, arena: SymmArena) -> ShardEntry: - entry = ShardEntry(arena=arena, offset=arena.offset_of(local), local=local, peer_views=tuple(arena.peer_views(local))) +def register_shard(local: torch.Tensor, buffer: SymmBuffer) -> ShardEntry: + entry = ShardEntry(buffer=buffer, offset=buffer.offset_of(local), local=local, peer_views=tuple(buffer.peer_views(local))) _SHARD_REGISTRY[local.data_ptr()] = entry return entry @@ -130,8 +130,8 @@ def lookup_shard(data_ptr: int) -> ShardEntry | None: return _SHARD_REGISTRY.get(data_ptr) -def registered_arenas() -> list[SymmArena]: - return list(_ARENAS) +def registered_buffers() -> list[SymmBuffer]: + return list(_BUFFERS) def find_shard_by_layout(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor | None: @@ -144,10 +144,10 @@ def find_shard_by_layout(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Te def reset_registry() -> None: - """Test-only: drop every arena so a new model can be built in-process.""" + """Test-only: drop every buffer so a new model can be built in-process.""" global _BARRIER_DONE _SHARD_REGISTRY.clear() - _ARENAS.clear() + _BUFFERS.clear() _BARRIER_DONE = False @@ -157,16 +157,16 @@ def barrier_after_load() -> None: Must run after weights are loaded and before the first peer read. """ global _BARRIER_DONE - if _BARRIER_DONE or not _ARENAS: + if _BARRIER_DONE or not _BUFFERS: return if dist.is_available() and dist.is_initialized(): torch.cuda.synchronize() dist.barrier() _BARRIER_DONE = True magi_logger.info( - "Symmetric arena: published %d arena(s), %.1f MiB, %d shards; steady state is barrier-free", - len(_ARENAS), - sum(a.nbytes for a in _ARENAS) / 2**20, + "SymmBuffer: published %d buffer(s), %.1f MiB, %d shards; steady state is barrier-free", + len(_BUFFERS), + sum(b.nbytes for b in _BUFFERS) / 2**20, len(_SHARD_REGISTRY), ) @@ -203,7 +203,7 @@ def _group_name_of(t) -> str | None: return None -def _arena_key(t) -> tuple[torch.dtype, str]: +def _buffer_key(t) -> tuple[torch.dtype, str]: """One window per (dtype, process group). Same dtype on two meshes (gaga4 FSDP + edp) must not share a window.""" group_name = _group_name_of(t) if group_name is None: @@ -225,40 +225,40 @@ def _apply_order_entries(mod: nn.Module): yield mod, name, p -def _plan_arenas(shards: list, device: torch.device) -> dict[tuple[torch.dtype, str], SymmArena]: +def _plan_buffers(shards: list, device: torch.device) -> dict[tuple[torch.dtype, str], SymmBuffer]: """Size and commit one window per (dtype, group). Dedup by identity so a tied weight reserves a single slot.""" - arenas: dict[tuple[torch.dtype, str], SymmArena] = {} + buffers: dict[tuple[torch.dtype, str], SymmBuffer] = {} seen: set[int] = set() for p in shards: if id(p) in seen: continue seen.add(id(p)) - key = _arena_key(p) - arena = arenas.get(key) - if arena is None: - arena = arenas[key] = SymmArena(p.dtype, device, key[1]) - arena.reserve(p._local_tensor.numel()) + key = _buffer_key(p) + buffer = buffers.get(key) + if buffer is None: + buffer = buffers[key] = SymmBuffer(p.dtype, device, key[1]) + buffer.reserve(p._local_tensor.numel()) - for arena in arenas.values(): - arena.commit() # the only collective, once per window - _ARENAS.extend(arenas.values()) - return arenas + for buffer in buffers.values(): + buffer.commit() # the only collective, once per window + _BUFFERS.extend(buffers.values()) + return buffers -def materialize_into_arenas(mod: nn.Module, device: torch.device) -> dict[tuple[torch.dtype, str], SymmArena]: +def materialize_into_buffers(mod: nn.Module, device: torch.device) -> dict[tuple[torch.dtype, str], SymmBuffer]: """Size windows for ``mod``'s Shard(0) shards while they are still on meta. Non-gatherable params are left to the caller.""" shards = [p for _, _, p in _apply_order_entries(mod) if _is_gatherable_shard(p)] if not shards: return {} - return _plan_arenas(shards, device) + return _plan_buffers(shards, device) -def migrate_to_arenas(root: nn.Module) -> dict[tuple[torch.dtype, str], SymmArena]: +def migrate_to_buffers(root: nn.Module) -> dict[tuple[torch.dtype, str], SymmBuffer]: """Copy already-allocated Shard(0) shards into symmetric memory. Used when ``magi_compile(model, ...)`` is given a live model rather than a meta + ``to_empty`` path. ``load_state_dict(assign=True)`` after this would - replace arena views with ordinary tensors; the gather then rejects them. + replace buffer views with ordinary tensors; the gather then rejects them. """ entries = [(m, n, p) for m, n, p in _apply_order_entries(root) if _is_gatherable_shard(p)] if not entries: @@ -267,7 +267,7 @@ def migrate_to_arenas(root: nn.Module) -> dict[tuple[torch.dtype, str], SymmAren device = entries[0][2]._local_tensor.device if device.type != "cuda": raise RuntimeError(f"symmetric memory needs the shards on cuda, found {device}") - arenas = _plan_arenas([p for _, _, p in entries], device) + buffers = _plan_buffers([p for _, _, p in entries], device) from torch.distributed.tensor import DTensor @@ -275,23 +275,23 @@ def migrate_to_arenas(root: nn.Module) -> dict[tuple[torch.dtype, str], SymmAren for owner, name, p in entries: local = views.get(id(p)) if local is None: - arena = arenas[_arena_key(p)] - local = views[id(p)] = arena.take(p._local_tensor.shape) + buffer = buffers[_buffer_key(p)] + local = views[id(p)] = buffer.take(p._local_tensor.shape) local.copy_(p._local_tensor) - register_shard(local, arena) + register_shard(local, buffer) moved = DTensor.from_local(local, p.device_mesh, p.placements, run_check=False) owner.register_parameter(name, nn.Parameter(moved, requires_grad=p.requires_grad)) magi_logger.info( - "Symmetric arena: migrated %d shard(s) into %.1f MiB across %d window(s)", + "SymmBuffer: migrated %d shard(s) into %.1f MiB across %d window(s)", len(views), - sum(a.nbytes for a in arenas.values()) / 2**20, - len(arenas), + sum(b.nbytes for b in buffers.values()) / 2**20, + len(buffers), ) - return arenas + return buffers -def patch_symm_arena_apply(cls: type[nn.Module]) -> None: +def patch_symm_buffer_apply(cls: type[nn.Module]) -> None: """Install the ``_apply`` interception on a decorated class. Mirrors ``_patch_cpu_offload_apply``: take over for ``to_empty``'s lambda, delegate everything else. @@ -299,19 +299,19 @@ def patch_symm_arena_apply(cls: type[nn.Module]) -> None: if getattr(cls, "_magi_symm_apply_patched", False): return orig_apply = cls._apply - magi_logger.info("Symmetric arena: intercepting %s._apply for copy-engine FSDP", cls.__name__) + magi_logger.info("SymmBuffer: intercepting %s._apply for copy-engine FSDP", cls.__name__) def _symm_apply(self, fn, recurse: bool = True): if getattr(fn, "__qualname__", "") != _TO_EMPTY_LAMBDA: return orig_apply(self, fn, recurse) - if getattr(self, "_magi_symm_arenas", None) is not None: + if getattr(self, "_magi_symm_buffers", None) is not None: return orig_apply(self, fn, recurse) device = torch.device(inspect.getclosurevars(fn).nonlocals["device"]) from torch.distributed.tensor import DTensor - arenas = materialize_into_arenas(self, device) - if not arenas: + buffers = materialize_into_buffers(self, device) + if not buffers: return orig_apply(self, fn, recurse) views: dict[int, torch.Tensor] = {} @@ -322,21 +322,21 @@ def materialize(t: torch.Tensor) -> torch.Tensor: # Tied weight: same view so tying survives materialization. local = views.get(id(t)) if local is None: - arena = arenas[_arena_key(t)] - local = views[id(t)] = arena.take(t._local_tensor.shape) - register_shard(local, arena) + buffer = buffers[_buffer_key(t)] + local = views[id(t)] = buffer.take(t._local_tensor.shape) + register_shard(local, buffer) return DTensor.from_local(local, t.device_mesh, t.placements, run_check=False) # Do not forge to_empty's qualname: a nested decorated block must fail - # the check above and delegate, so its params land in *this* arena. + # the check above and delegate, so its params land in *this* buffer. out = orig_apply(self, materialize, recurse) - self._magi_symm_arenas = arenas + self._magi_symm_buffers = buffers magi_logger.info( - "Symmetric arena: %s materialized %d shard(s) into %.1f MiB across %d window(s)", + "SymmBuffer: %s materialized %d shard(s) into %.1f MiB across %d window(s)", cls.__name__, len(views), - sum(a.nbytes for a in arenas.values()) / 2**20, - len(arenas), + sum(b.nbytes for b in buffers.values()) / 2**20, + len(buffers), ) return out 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 f98e436..720bce6 100644 --- a/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py +++ b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py @@ -39,7 +39,7 @@ the expected mode for each rung of the ladder (identical / slot / pinned / abort). With ``--copy-engine``: the same shape, but the gathers are -``magi::symm_all_gather`` reading a symmetric arena. Two things are checked that +``magi::symm_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 @@ -123,7 +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 arena instead of NCCL") + 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: @@ -169,17 +169,17 @@ def fn(x, w0, shard): ce_shards: list = [] if args.copy_engine: - from magi_compiler.symm_mem import SymmArena, register_shard + from magi_compiler.symm_mem import SymmBuffer, register_shard from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER - arena = SymmArena(torch.bfloat16, torch.device("cuda", dev), grp) + buffer = SymmBuffer(torch.bfloat16, torch.device("cuda", dev), grp) for _ in range(N_CE_LAYERS): - arena.reserve(H * H) - arena.commit() + buffer.reserve(H * H) + buffer.commit() for i in range(N_CE_LAYERS): - s = arena.take((H, H)) + s = buffer.take((H, H)) s.normal_(0.0, H**-0.5).add_(0.01 * i) - register_shard(s, arena) + register_shard(s, buffer) ce_shards.append(s) # A peer read is only legal once that peer has written its shard. torch.cuda.synchronize() diff --git a/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py b/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py index 61b7c57..17994a5 100644 --- a/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py +++ b/tests/feature_tests/fsdp/test_fsdp_overlap_reorder.py @@ -181,7 +181,7 @@ def test_reorder_graph_mismatch_slot_mode(): @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 arena. + """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 diff --git a/tests/feature_tests/fsdp/test_profiling_estimator.py b/tests/feature_tests/fsdp/test_profiling_estimator.py index 275a9f4..833107d 100644 --- a/tests/feature_tests/fsdp/test_profiling_estimator.py +++ b/tests/feature_tests/fsdp/test_profiling_estimator.py @@ -705,18 +705,18 @@ def _ce_group_name() -> str: def _register_shards(shapes, dtype=torch.bfloat16): """Register ``shapes`` as real symmetric-memory shards, filled distinctly.""" - from magi_compiler.symm_mem import SymmArena, register_shard + from magi_compiler.symm_mem import SymmBuffer, register_shard - arena = SymmArena(dtype, torch.device("cuda", 0), _ce_group_name()) + buffer = SymmBuffer(dtype, torch.device("cuda", 0), _ce_group_name()) for shape in shapes: - arena.reserve(shape[0] * shape[1]) - arena.commit() + buffer.reserve(shape[0] * shape[1]) + buffer.commit() shards = [] for i, shape in enumerate(shapes): - s = arena.take(shape) + s = buffer.take(shape) s.fill_(i + 1) - register_shard(s, arena) + register_shard(s, buffer) shards.append(s) torch.cuda.synchronize() return shards @@ -730,7 +730,7 @@ def _ce_replay_snode(shapes, coalesced=False): @requires_cuda def test_symm_ag_replay_gathers_the_registered_shard(pg_1rank, symm_registry): - """The replay must run the real op on a real arena shard: a gather of an + """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)]) diff --git a/tests/feature_tests/symm_helper/verify_symm_e2e.py b/tests/feature_tests/symm_helper/verify_symm_e2e.py index 142cef3..b84ac08 100644 --- a/tests/feature_tests/symm_helper/verify_symm_e2e.py +++ b/tests/feature_tests/symm_helper/verify_symm_e2e.py @@ -23,7 +23,7 @@ The builder here is deliberately ignorant of symmetric memory, because the real one is too. -Then a checkpoint load writes into the arena views (``copy_``, not ``assign``), +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: @@ -120,7 +120,7 @@ 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 arena views out for ordinary tensors and quietly + 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 @@ -205,16 +205,16 @@ def spy_rewrite(graph): _load_into_shards(model, state) # (1) placement: the block's shards are in a window, the head's are not. - from magi_compiler.symm_mem import lookup_shard, registered_arenas + from magi_compiler.symm_mem import lookup_shard, registered_buffers block_names = {n for n, _ in _named_shards(model) if n.startswith("block.")} - in_arena = [n for n, p in _named_shards(model) if lookup_shard(p._local_tensor.data_ptr()) is not None] - head_in_arena = [n for n in in_arena if not n.startswith("block.")] - arena_mib = sum(a.nbytes for a in registered_arenas()) / 2**20 - placement_ok = (set(in_arena) == block_names) if ce else (in_arena == []) + 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_arena)}/{len(block_names)} block shards in {len(registered_arenas())} " - f"window(s) ({arena_mib:.0f} MiB), {len(head_in_arena)} stray -> {'ok' if placement_ok else 'WRONG'}" + 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'}" ) # (2) + (3): compile, then run twice. The second step is the one that @@ -243,7 +243,7 @@ def rel(a, b): log( f"\nCONFIG world={world} hidden={args.hidden} layers={args.n_layers} tokens={args.n_tokens} " f"transport={args.transport}\n" - f"MEMORY arena={arena_mib:.0f}MiB " + f"MEMORY buffer={buffer_mib:.0f}MiB " f"(one gathered layer = {args.hidden ** 2 * dtype.itemsize / 2**20:.0f}MiB)" ) diff --git a/tests/feature_tests/test_symm_ag_rewrite.py b/tests/feature_tests/test_symm_ag_rewrite.py index 7cc66f0..7ddb087 100644 --- a/tests/feature_tests/test_symm_ag_rewrite.py +++ b/tests/feature_tests/test_symm_ag_rewrite.py @@ -263,7 +263,7 @@ def test_end_to_end_lowering_then_rewrite(mesh_1rank): @requires_cuda def test_copy_engine_buckets_then_rewrites_coalesced(mesh_1rank): - """Phase-1 wrap: arena gathers are bucketed first, then the coalesced + """Phase-1 wrap: SymmBuffer gathers are bucketed first, then the coalesced launch is retargeted. Members stay separate dests underneath.""" from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED @@ -281,7 +281,7 @@ def test_copy_engine_buckets_then_rewrites_coalesced(mesh_1rank): @requires_cuda def test_copy_engine_does_not_bucket_cast_gathers(mesh_1rank): - """Cast outputs are not arena shards; they must stay on NCCL and not join a CE bucket.""" + """Cast outputs are not SymmBuffer shards; they must stay on NCCL and not join a CE bucket.""" from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED @@ -327,7 +327,7 @@ def _coalesced_graph(mesh, n: int, *, derive: str | None = None, marked: bool = @requires_cuda -def test_coalesced_of_arena_shards_is_retargeted(mesh_1rank): +def test_coalesced_of_symm_shards_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 SYMM_ALL_GATHER_COALESCED @@ -340,7 +340,7 @@ def test_coalesced_of_arena_shards_is_retargeted(mesh_1rank): @requires_cuda def test_coalesced_is_all_or_nothing(mesh_1rank): """A bucket is one submission, so it can only go to the copy engine if - *every* member is an arena shard -- one cast member has to keep the whole + *every* member is an SymmBuffer shard -- one cast 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 rewrite_weight_ag_to_copy_engine from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED diff --git a/tests/feature_tests/test_symm_all_gather.py b/tests/feature_tests/test_symm_all_gather.py index 0fac81c..b4dcef7 100644 --- a/tests/feature_tests/test_symm_all_gather.py +++ b/tests/feature_tests/test_symm_all_gather.py @@ -18,7 +18,7 @@ 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 arena never saw is rejected loudly rather than gathering +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 @@ -72,12 +72,12 @@ def _clean_state(): reset_registry() -def _arena_model(mesh, hidden: int = 64, n_layers: int = 3, dtype=torch.bfloat16): - """A meta-built, Shard(0)-sharded model materialized into a symmetric arena, +def _symm_model(mesh, hidden: int = 64, n_layers: int = 3, dtype=torch.bfloat16): + """A meta-built, Shard(0)-sharded model materialized into a symmetric buffer, exactly the shape step 1 produces.""" from torch.distributed.tensor import Shard, distribute_tensor - from magi_compiler.symm_mem import materialize_into_arenas + from magi_compiler.symm_mem import materialize_into_buffers class Block(nn.Module): def __init__(self): @@ -93,16 +93,16 @@ def __init__(self): device = torch.device("cuda", 0) from torch.distributed.tensor import DTensor - from magi_compiler.symm_mem.arena import _arena_key, register_shard + from magi_compiler.symm_mem.symm_buffer import _buffer_key, register_shard - arenas = materialize_into_arenas(model, device) + buffers = materialize_into_buffers(model, device) views: dict[int, torch.Tensor] = {} def materialize(t): if isinstance(t, DTensor): - arena = arenas[_arena_key(t)] - local = arena.take(t._local_tensor.shape) - register_shard(local, arena) + buffer = buffers[_buffer_key(t)] + local = buffer.take(t._local_tensor.shape) + register_shard(local, buffer) views[id(t)] = local return DTensor.from_local(local, t.device_mesh, t.placements, run_check=False) return torch.empty_like(t, device=device) @@ -119,7 +119,7 @@ def materialize(t): @requires_cuda def test_gather_matches_nccl_bitwise(mesh_1rank): - _, shards = _arena_model(mesh_1rank) + _, shards = _symm_model(mesh_1rank) for shard in shards: got = _WAIT(torch.ops.magi.symm_all_gather(shard, 1, "")) @@ -132,7 +132,7 @@ def test_gather_matches_nccl_bitwise(mesh_1rank): @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 = _arena_model(mesh_1rank) + _, shards = _symm_model(mesh_1rank) outs = torch.ops.magi.symm_all_gather_coalesced(list(shards), 1, "") assert len(outs) == len(shards) for out, shard in zip(outs, shards): @@ -144,7 +144,7 @@ def test_coalesced_wrap_matches_per_member_gather(mesh_1rank): @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 = _arena_model(mesh_1rank) + _, shards = _symm_model(mesh_1rank) shard = shards[0] out = _WAIT(torch.ops.magi.symm_all_gather(shard, 1, "")) @@ -156,7 +156,7 @@ def test_wait_tensor_picks_up_the_registered_event(mesh_1rank): 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 = _arena_model(mesh_1rank) + _, shards = _symm_model(mesh_1rank) a = torch.ops.magi.symm_all_gather(shards[0], 1, "") b = torch.ops.magi.symm_all_gather(shards[1], 1, "") _WAIT(a) @@ -169,7 +169,7 @@ def test_each_gather_returns_a_fresh_buffer(mesh_1rank): @requires_cuda def test_unregistered_shard_is_rejected(mesh_1rank): - """A weight the arena never claimed has no peer views, so gathering it would + """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"): @@ -178,7 +178,7 @@ def test_unregistered_shard_is_rejected(mesh_1rank): @requires_cuda def test_group_size_mismatch_is_rejected(mesh_1rank): - _, shards = _arena_model(mesh_1rank) + _, shards = _symm_model(mesh_1rank) with pytest.raises(RuntimeError, match="peers but the gather asks for"): torch.ops.magi.symm_all_gather(shards[0], 4, "") @@ -188,7 +188,7 @@ 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 = _arena_model(mesh_1rank) + _, 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.symm_all_gather_coalesced([shards[0], ordinary], 1, "") @@ -271,7 +271,7 @@ def test_per_peer_fallback_matches_the_batched_path(mesh_1rank, monkeypatch): 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 = _arena_model(mesh_1rank) + _, shards = _symm_model(mesh_1rank) batched = _WAIT(torch.ops.magi.symm_all_gather(shards[0], 1, "")) torch.cuda.synchronize() diff --git a/tests/feature_tests/test_symm_arena.py b/tests/feature_tests/test_symm_buffer.py similarity index 79% rename from tests/feature_tests/test_symm_arena.py rename to tests/feature_tests/test_symm_buffer.py index 269c70d..3b40ed7 100644 --- a/tests/feature_tests/test_symm_arena.py +++ b/tests/feature_tests/test_symm_buffer.py @@ -131,26 +131,26 @@ def forward(self, x): @requires_cuda -def test_root_to_empty_puts_block_shards_in_arena(mesh_1rank): +def test_root_to_empty_puts_block_shards_in_buffer(mesh_1rank): """The decorated block claims its own subtree; the rest stays ordinary.""" - from magi_compiler.symm_mem import lookup_shard, registered_arenas + from magi_compiler.symm_mem import lookup_shard, registered_buffers hidden, n_layers = 64, 3 model = _make_root(_make_block_cls("copy_engine"), hidden, n_layers, torch.bfloat16) _shard(model, mesh_1rank) model.to_empty(device=torch.device("cuda", 0)) - arenas = registered_arenas() - assert len(arenas) == 1, "one window per (dtype, group), not one per weight" - arena = arenas[0] + buffers = registered_buffers() + assert len(buffers) == 1, "one window per (dtype, group), not one per weight" + buffer = buffers[0] block_shards = list(model.block.parameters()) assert len(block_shards) == 2 * n_layers - assert all(arena.contains(p._local_tensor) for p in block_shards) + assert all(buffer.contains(p._local_tensor) for p in block_shards) assert all(lookup_shard(p._local_tensor.data_ptr()) is not None for p in block_shards) # Outside the decorated class: untouched by the interception. - assert not arena.contains(model.head.weight._local_tensor) + assert not buffer.contains(model.head.weight._local_tensor) assert lookup_shard(model.head.weight._local_tensor.data_ptr()) is None # Still ordinary, working DTensors on real storage. @@ -161,18 +161,18 @@ def test_root_to_empty_puts_block_shards_in_arena(mesh_1rank): @requires_cuda def test_nccl_transport_leaves_allocation_alone(mesh_1rank): """The default transport must not install the interception at all.""" - from magi_compiler.symm_mem import registered_arenas + from magi_compiler.symm_mem import registered_buffers model = _make_root(_make_block_cls("nccl"), 64, 2, torch.bfloat16) _shard(model, mesh_1rank) model.to_empty(device=torch.device("cuda", 0)) - assert registered_arenas() == [] + assert registered_buffers() == [] assert all(p._local_tensor.device.type == "cuda" for p in model.parameters()) @requires_cuda -def test_peer_view_round_trips_through_arena(mesh_1rank): +def test_peer_view_round_trips_through_buffer(mesh_1rank): """A shard written locally must be visible through its own peer view: this is the addressing the copy-engine gather depends on.""" from magi_compiler.symm_mem import lookup_shard @@ -201,7 +201,7 @@ def test_tied_weights_share_one_slot(mesh_1rank): ``data_parallel`` calls ``distribute_tensor`` per module, which replaces each entry with its own DTensor and unties them on its own. """ - from magi_compiler.symm_mem import registered_arenas + from magi_compiler.symm_mem import registered_buffers hidden = 64 @@ -219,19 +219,19 @@ def forward(self, x): model.block.b.weight = model.block.a.weight model.to_empty(device=torch.device("cuda", 0)) - arena = registered_arenas()[0] + buffer = registered_buffers()[0] assert model.block.a.weight is model.block.b.weight, "tying must survive materialization" - assert arena.contains(model.block.a.weight._local_tensor) + assert buffer.contains(model.block.a.weight._local_tensor) # One slot in the window, not two. - slot = arena.ALIGN * ((hidden * hidden + arena.ALIGN - 1) // arena.ALIGN) - assert arena.nbytes == slot * torch.bfloat16.itemsize + slot = buffer.ALIGN * ((hidden * hidden + buffer.ALIGN - 1) // buffer.ALIGN) + assert buffer.nbytes == slot * torch.bfloat16.itemsize @requires_cuda -def test_nested_decorated_block_shares_the_outer_arena(mesh_1rank): +def test_nested_decorated_block_shares_the_outer_buffer(mesh_1rank): """A decorated block inside a decorated block must not open a second window: the inner one has to fail the lambda check and delegate.""" - from magi_compiler.symm_mem import registered_arenas + from magi_compiler.symm_mem import registered_buffers inner_cls = _decorate( type( @@ -261,25 +261,25 @@ def forward(self, x): _shard(model, mesh_1rank) model.to_empty(device=torch.device("cuda", 0)) - arenas = registered_arenas() - assert len(arenas) == 1, f"nested decoration opened {len(arenas)} windows" - assert arenas[0].contains(model.block.own.weight._local_tensor) - assert arenas[0].contains(model.block.inner.lin.weight._local_tensor) + buffers = registered_buffers() + assert len(buffers) == 1, f"nested decoration opened {len(buffers)} windows" + assert buffers[0].contains(model.block.own.weight._local_tensor) + assert buffers[0].contains(model.block.inner.lin.weight._local_tensor) @requires_cuda def test_non_shard0_placement_falls_back(mesh_1rank): """Replicate weights are not gatherable, so they must be allocated normally - rather than silently placed in the arena.""" + rather than silently placed in the buffer.""" from torch.distributed.tensor import Replicate - from magi_compiler.symm_mem import registered_arenas + from magi_compiler.symm_mem import registered_buffers model = _make_root(_make_block_cls("copy_engine"), 64, 2, torch.bfloat16) _shard(model, mesh_1rank, placement=Replicate()) model.to_empty(device=torch.device("cuda", 0)) - assert registered_arenas() == [] + assert registered_buffers() == [] assert all(p._local_tensor.device.type == "cuda" for p in model.block.parameters()) @@ -294,18 +294,18 @@ def test_two_process_groups_same_dtype_get_two_windows(mesh_1rank, monkeypatch): """ from torch.distributed.tensor import Shard, distribute_tensor - from magi_compiler.symm_mem import arena as sa + from magi_compiler.symm_mem import symm_buffer as sb hidden = 64 a = nn.Parameter(distribute_tensor(torch.empty(hidden, hidden, dtype=torch.bfloat16), mesh_1rank, [Shard(0)])) b = nn.Parameter(distribute_tensor(torch.empty(hidden, hidden, dtype=torch.bfloat16), mesh_1rank, [Shard(0)])) - monkeypatch.setattr(sa, "_group_name_of", lambda p, _a=a: "dense_fsdp" if p is _a else "edp") - monkeypatch.setattr(sa.SymmArena, "commit", lambda self: None) - arenas = sa._plan_arenas([a, b], torch.device("cuda", 0)) - assert set(arenas) == {(torch.bfloat16, "dense_fsdp"), (torch.bfloat16, "edp")} - assert arenas[(torch.bfloat16, "dense_fsdp")].group_name == "dense_fsdp" - assert arenas[(torch.bfloat16, "edp")].group_name == "edp" + monkeypatch.setattr(sb, "_group_name_of", lambda p, _a=a: "dense_fsdp" if p is _a else "edp") + monkeypatch.setattr(sb.SymmBuffer, "commit", lambda self: None) + buffers = sb._plan_buffers([a, b], torch.device("cuda", 0)) + assert set(buffers) == {(torch.bfloat16, "dense_fsdp"), (torch.bfloat16, "edp")} + assert buffers[(torch.bfloat16, "dense_fsdp")].group_name == "dense_fsdp" + assert buffers[(torch.bfloat16, "edp")].group_name == "edp" @requires_cuda @@ -331,14 +331,14 @@ def group_name(): return dist.group.WORLD.group_name -def _committed_arena(group_name, numels, dtype=torch.bfloat16): - from magi_compiler.symm_mem import SymmArena +def _committed_buffer(group_name, numels, dtype=torch.bfloat16): + from magi_compiler.symm_mem import SymmBuffer - arena = SymmArena(dtype, torch.device("cuda", 0), group_name) + buffer = SymmBuffer(dtype, torch.device("cuda", 0), group_name) for n in numels: - arena.reserve(n) - arena.commit() - return arena + buffer.reserve(n) + buffer.commit() + return buffer @requires_cuda @@ -346,16 +346,16 @@ def test_shards_are_dispensed_at_aligned_offsets(mesh_1rank, group_name): """Slots are padded to ``ALIGN`` for copy-engine throughput, so the second shard does not start where the first one ends. ``offset_of`` is what the peer views are built from, so it has to agree with what ``take`` handed out.""" - from magi_compiler.symm_mem import SymmArena + from magi_compiler.symm_mem import SymmBuffer rows, cols = 3, 5 # 15 elems: deliberately not a multiple of ALIGN - arena = _committed_arena(group_name, [rows * cols, rows * cols]) - first = arena.take((rows, cols)) - second = arena.take((rows, cols)) + buffer = _committed_buffer(group_name, [rows * cols, rows * cols]) + first = buffer.take((rows, cols)) + second = buffer.take((rows, cols)) - assert arena.offset_of(first) == 0 - assert arena.offset_of(second) == SymmArena.ALIGN - assert arena.contains(first) and arena.contains(second) + assert buffer.offset_of(first) == 0 + assert buffer.offset_of(second) == SymmBuffer.ALIGN + assert buffer.contains(first) and buffer.contains(second) assert first.shape == (rows, cols) @@ -364,19 +364,19 @@ def test_dispensing_more_than_was_reserved_is_an_error(mesh_1rank, group_name): """The sizing walk and the dispensing walk are two separate traversals; if they ever disagree the shards silently overlap, so the window must run out rather than hand back memory reserved for someone else.""" - arena = _committed_arena(group_name, [64]) - arena.take((8, 8)) - with pytest.raises(RuntimeError, match="symmetric arena overflow"): - arena.take((8, 8)) + buffer = _committed_buffer(group_name, [64]) + buffer.take((8, 8)) + with pytest.raises(RuntimeError, match="symmetric buffer overflow"): + buffer.take((8, 8)) @requires_cuda def test_contains_rejects_memory_outside_the_window(mesh_1rank, group_name): """``contains`` is how the rewrite decides a weight is gatherable; a caching allocator tensor must never pass.""" - arena = _committed_arena(group_name, [64]) - arena.take((8, 8)) - assert not arena.contains(torch.empty(8, 8, device="cuda", dtype=torch.bfloat16)) + buffer = _committed_buffer(group_name, [64]) + buffer.take((8, 8)) + assert not buffer.contains(torch.empty(8, 8, device="cuda", dtype=torch.bfloat16)) @requires_cuda @@ -386,11 +386,11 @@ def test_find_shard_by_layout_matches_on_shape_and_dtype(mesh_1rank, group_name) be a None it can degrade on, not a wrong-dtype shard it would gather.""" from magi_compiler.symm_mem import find_shard_by_layout, register_shard - arena = _committed_arena(group_name, [8 * 4, 16 * 4]) - small = arena.take((8, 4)) - large = arena.take((16, 4)) - register_shard(small, arena) - register_shard(large, arena) + buffer = _committed_buffer(group_name, [8 * 4, 16 * 4]) + small = buffer.take((8, 4)) + large = buffer.take((16, 4)) + register_shard(small, buffer) + register_shard(large, buffer) assert find_shard_by_layout((8, 4), torch.bfloat16) is small assert find_shard_by_layout((16, 4), torch.bfloat16) is large @@ -399,31 +399,31 @@ def test_find_shard_by_layout_matches_on_shape_and_dtype(mesh_1rank, group_name) @requires_cuda -def test_reset_registry_drops_arenas_and_shards(mesh_1rank, group_name): +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 find_shard_by_layout, lookup_shard, register_shard, registered_arenas, reset_registry + from magi_compiler.symm_mem import find_shard_by_layout, lookup_shard, register_shard, registered_buffers, reset_registry - arena = _committed_arena(group_name, [8 * 4]) - shard = arena.take((8, 4)) - register_shard(shard, arena) + buffer = _committed_buffer(group_name, [8 * 4]) + shard = buffer.take((8, 4)) + register_shard(shard, buffer) assert lookup_shard(shard.data_ptr()) is not None reset_registry() - assert registered_arenas() == [] + assert registered_buffers() == [] assert lookup_shard(shard.data_ptr()) is None assert find_shard_by_layout((8, 4), torch.bfloat16) is None # --------------------------------------------------------------------------- -# migrate_to_arenas -- the live-model path (magi_compile on an already +# migrate_to_buffers -- the live-model path (magi_compile on an already # materialized model, where there is no to_empty to intercept). # --------------------------------------------------------------------------- @requires_cuda def test_migrate_moves_live_shards_and_keeps_their_values(mesh_1rank): """Unlike ``to_empty``, this runs on weights that already hold data, so the copy is load-bearing: dropping it would gather uninitialized memory.""" - from magi_compiler.symm_mem import lookup_shard, migrate_to_arenas + from magi_compiler.symm_mem import lookup_shard, migrate_to_buffers hidden = 64 @@ -437,13 +437,13 @@ def __init__(self): _shard(model, mesh_1rank) before = {n: p._local_tensor.clone() for n, p in model.named_parameters()} - arenas = migrate_to_arenas(model) - assert len(arenas) == 1 - arena = next(iter(arenas.values())) + buffers = migrate_to_buffers(model) + assert len(buffers) == 1 + buffer = next(iter(buffers.values())) for name, p in model.named_parameters(): local = p._local_tensor - assert arena.contains(local), f"{name} was not migrated" + assert buffer.contains(local), f"{name} was not migrated" assert lookup_shard(local.data_ptr()) is not None assert torch.equal(local, before[name]), f"{name} lost its values" @@ -453,7 +453,7 @@ def test_migrate_gives_a_tied_weight_one_slot(mesh_1rank): """Migration rebuilds each parameter, so python identity does not survive -- what must survive is the storage, or the tie is gone and the window is overflowed by a walk that sized it once.""" - from magi_compiler.symm_mem import migrate_to_arenas + from magi_compiler.symm_mem import migrate_to_buffers hidden = 64 @@ -467,21 +467,21 @@ def __init__(self): _shard(model, mesh_1rank) model.b.weight = model.a.weight - arenas = migrate_to_arenas(model) - arena = next(iter(arenas.values())) + buffers = migrate_to_buffers(model) + buffer = next(iter(buffers.values())) assert model.a.weight._local_tensor.data_ptr() == model.b.weight._local_tensor.data_ptr() - slot = arena.ALIGN * ((hidden * hidden + arena.ALIGN - 1) // arena.ALIGN) - assert arena.nbytes == slot * torch.bfloat16.itemsize + slot = buffer.ALIGN * ((hidden * hidden + buffer.ALIGN - 1) // buffer.ALIGN) + assert buffer.nbytes == slot * torch.bfloat16.itemsize @requires_cuda def test_migrate_refuses_a_model_that_was_never_materialized(mesh_1rank): - """``migrate_to_arenas`` is the live-model entry point, so being handed a + """``migrate_to_buffers`` is the live-model entry point, so being handed a still-on-meta model is the way it gets misused. Copying from meta silently produces a window of uninitialized weights, so it has to fail instead.""" from torch.distributed.tensor import DTensor, Shard - from magi_compiler.symm_mem import migrate_to_arenas + from magi_compiler.symm_mem import migrate_to_buffers with torch.device("meta"): model = nn.Linear(8, 8, bias=False, dtype=torch.bfloat16) @@ -490,14 +490,14 @@ def test_migrate_refuses_a_model_that_was_never_materialized(mesh_1rank): assert model.weight._local_tensor.is_meta # the state under test with pytest.raises(RuntimeError, match="needs the shards on cuda"): - migrate_to_arenas(model) + migrate_to_buffers(model) @requires_cuda def test_migrate_leaves_a_model_with_no_gatherable_shards_alone(mesh_1rank): """A plain (unsharded) model must not open an empty window.""" - from magi_compiler.symm_mem import migrate_to_arenas, registered_arenas + from magi_compiler.symm_mem import migrate_to_buffers, registered_buffers model = nn.Linear(8, 8, bias=False, dtype=torch.bfloat16).to("cuda") - assert migrate_to_arenas(model) == {} - assert registered_arenas() == [] + assert migrate_to_buffers(model) == {} + assert registered_buffers() == [] diff --git a/tests/feature_tests/test_symm_e2e.py b/tests/feature_tests/test_symm_e2e.py index 31ef817..944aa91 100644 --- a/tests/feature_tests/test_symm_e2e.py +++ b/tests/feature_tests/test_symm_e2e.py @@ -14,7 +14,7 @@ """End-to-end guard for ``fsdp_config.transport="copy_engine"``. -The unit tests each stub something out: the arena tests materialize by calling +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, From 2bc891215e640a83b6d5f9b0da43bd718240f320 Mon Sep 17 00:00:00 2001 From: wtr Date: Mon, 7 Sep 2026 20:52:23 +0800 Subject: [PATCH 11/16] [Refactor] Bind copy-engine FSDP weights from the captured graph and keep uneven Shard(0) gathers rank-identical --- magi_compiler/_api.py | 15 - magi_compiler/config.py | 15 +- magi_compiler/magi_backend/magi_backend.py | 22 +- magi_compiler/passes/fsdp_overlap/__init__.py | 11 +- .../passes/fsdp_overlap/bucket_all_gather.py | 28 +- .../passes/fsdp_overlap/copy_engine.py | 108 ++++ .../passes/fsdp_overlap/lower_and_bucket.py | 34 +- .../passes/fsdp_overlap/node_meta.py | 57 ++ .../fsdp_overlap/redistribute_lowering.py | 12 +- .../passes/fsdp_overlap/symm_ag_rewrite.py | 90 --- magi_compiler/symm_mem/__init__.py | 19 +- magi_compiler/symm_mem/bind.py | 324 ++++++++++ magi_compiler/symm_mem/symm_buffer.py | 301 +++------- .../fsdp_overlap_helper/reorder_helper.py | 12 +- .../uneven_shard_helper.py | 269 +++++++++ .../fsdp/test_fsdp_overlap_bucket.py | 4 +- .../fsdp/test_fsdp_overlap_lowering.py | 14 +- .../fsdp/test_profiling_estimator.py | 12 +- .../test_python_scalar_constant_folding.py | 4 +- .../symm_helper/verify_symm_e2e.py | 24 +- ...symm_ag_rewrite.py => test_copy_engine.py} | 274 ++++----- tests/feature_tests/test_symm_all_gather.py | 32 +- tests/feature_tests/test_symm_bind.py | 561 ++++++++++++++++++ tests/feature_tests/test_symm_buffer.py | 503 ---------------- .../test_uneven_shard_transport.py | 82 +++ 25 files changed, 1753 insertions(+), 1074 deletions(-) create mode 100644 magi_compiler/passes/fsdp_overlap/copy_engine.py create mode 100644 magi_compiler/passes/fsdp_overlap/node_meta.py delete mode 100644 magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py create mode 100644 magi_compiler/symm_mem/bind.py create mode 100644 tests/feature_tests/fsdp/fsdp_overlap_helper/uneven_shard_helper.py rename tests/feature_tests/{test_symm_ag_rewrite.py => test_copy_engine.py} (63%) create mode 100644 tests/feature_tests/test_symm_bind.py delete mode 100644 tests/feature_tests/test_symm_buffer.py create mode 100644 tests/feature_tests/test_uneven_shard_transport.py diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py index 1077f0c..033a020 100644 --- a/magi_compiler/_api.py +++ b/magi_compiler/_api.py @@ -184,11 +184,6 @@ def _lazy_init_magi_state( if getattr(state_holder, state_attr, None) is not None: return - if conf.fsdp_config.transport == "copy_engine": - from magi_compiler.symm_mem import barrier_after_load - - barrier_after_load() - compilation_counter.num_models_seen += 1 setattr( @@ -223,11 +218,6 @@ def _magi_compile_class( if issubclass(cls, nn.Module) and conf.offload_config.model_cpu_offload: _patch_cpu_offload_apply(cls, conf) - if issubclass(cls, nn.Module) and conf.fsdp_config.transport == "copy_engine": - from magi_compiler.symm_mem import patch_symm_buffer_apply - - patch_symm_buffer_apply(cls) - old_init = cls.__init__ @functools.wraps(old_init) @@ -251,11 +241,6 @@ def _magi_compile_bound_method( if getattr(instance, installed_attr, False): return instance - if conf.fsdp_config.transport == "copy_engine" and isinstance(instance, nn.Module): - from magi_compiler.symm_mem import migrate_to_buffers - - migrate_to_buffers(instance) - old_method = getattr(instance, method_name) @torch.compiler.disable() diff --git a/magi_compiler/config.py b/magi_compiler/config.py index c27bff6..b81fd37 100644 --- a/magi_compiler/config.py +++ b/magi_compiler/config.py @@ -267,11 +267,16 @@ class FSDPConfig(BaseModel): transport: Literal["nccl", "copy_engine"] = Field( "nccl", description=( - "How weight all-gathers move bytes. 'nccl': ring kernels on the SMs. " - "'copy_engine': weight shards are allocated in symmetric memory at model build time and " - "gathered by peer copy-engine reads -- zero SM occupancy and no per-step cross-rank barrier, " - "at a lower raw bandwidth. Requires all ranks of the FSDP mesh dim to be NVLink-connected " - "within one node, and static weights (inference)." + "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." ), ) diff --git a/magi_compiler/magi_backend/magi_backend.py b/magi_compiler/magi_backend/magi_backend.py index 1beb364..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 @@ -616,7 +617,12 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None: 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, transport=fsdp_cfg.transport + 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: transport=%s bucket_mode=%s bucket_size=%d MiB created %d buckets", @@ -643,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 ( @@ -661,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 @@ -722,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 49b1ff1..245ca4f 100644 --- a/magi_compiler/passes/fsdp_overlap/__init__.py +++ b/magi_compiler/passes/fsdp_overlap/__init__.py @@ -13,15 +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 -from .symm_ag_rewrite import rewrite_weight_ag_to_copy_engine __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 5d15852..c4e35c8 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,7 +66,7 @@ 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) @@ -167,9 +169,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 - coalesced.meta["magi_fsdp_uneven_shard"] = any(ag.meta.get("magi_fsdp_uneven_shard") for ag in ag_nodes) + 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): @@ -188,7 +190,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, eligible=None) -> 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`` @@ -204,24 +206,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 - if eligible is not None and not eligible(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..a5a446e --- /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 SYMM_ALL_GATHER, SYMM_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 = SYMM_ALL_GATHER if node.target is _ALL_GATHER else SYMM_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 39c7b93..4524f0c 100644 --- a/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py +++ b/magi_compiler/passes/fsdp_overlap/lower_and_bucket.py @@ -14,17 +14,25 @@ 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 -from .symm_ag_rewrite import rewrite_weight_ag_to_copy_engine def lower_and_bucket_full_graph( - graph: fx.GraphModule, bucket_mode: str, bucket_size_bytes: int = 0, transport: str = "nccl" + 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). @@ -39,26 +47,26 @@ def lower_and_bucket_full_graph( the byte cap in program order (see ``bucket_weight_all_gather_coalesced``). 0 = no cap (one bucket per (group, dtype) run). - ``transport="copy_engine"`` buckets *first* (only SymmBuffer-shard gathers, so a cast - of the shard and an unevenly split weight stay out of the bucket), then retargets - both the leftover singles and the coalesced launches at the copy-engine ops. The - wrapper still runs one gather per member; reorder just sees one comm node per - bucket. Bucket membership therefore depends on the eligibility predicate, which - is why that predicate has to answer identically on every rank -- see - ``symm_ag_rewrite._is_uneven_shard``. + ``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) + 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": - from .symm_ag_rewrite import _input_is_symm_shard - - eligible = _input_is_symm_shard if transport == "copy_engine" else None - n = bucket_weight_all_gather_coalesced(graph, bucket_size_bytes=bucket_size_bytes, eligible=eligible) + 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'") 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 e31f363..5958d96 100644 --- a/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py +++ b/magi_compiler/passes/fsdp_overlap/redistribute_lowering.py @@ -22,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 @@ -173,15 +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 - # Whether this Shard(0) divides evenly, recorded from F and world -- both - # the same on every rank. Downstream transport choices must key off THIS - # and never off `L`: L is what makes the pad above appear on the trailing - # ranks only, so a predicate that reads the pad splits one collective into - # copy-engine on some ranks and NCCL on others, which never completes. - ag.meta["magi_fsdp_uneven_shard"] = world * chunk != F + 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) diff --git a/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py b/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py deleted file mode 100644 index fabb112..0000000 --- a/magi_compiler/passes/fsdp_overlap/symm_ag_rewrite.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright (c) 2026 SandAI. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from __future__ import annotations - -import torch -import torch.fx as fx - -from magi_compiler.utils import magi_logger - -_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default -_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default - - -def _is_uneven_shard(node: fx.Node) -> bool: - return bool(node.meta.get("magi_fsdp_uneven_shard")) - - -def _input_is_symm_shard(node: fx.Node) -> bool: - """True when the gather input is ``to_local(placeholder/get_attr)`` with no intervening op.""" - if _is_uneven_shard(node): - return False - src = node.args[0] if node.args else None - if not isinstance(src, fx.Node) or src.op != "call_method" or src.target != "to_local": - return False - owner = src.args[0] if src.args else None - return isinstance(owner, fx.Node) and owner.op in ("placeholder", "get_attr") - - -def _coalesced_inputs_are_symm_shards(node: fx.Node) -> bool: - if _is_uneven_shard(node): - return False - locs = node.args[0] if node.args else None - if not isinstance(locs, (list, tuple)) or not locs: - return False - return all(isinstance(loc, fx.Node) and _input_is_symm_shard_from_local(loc) for loc in locs) - - -def _input_is_symm_shard_from_local(src: fx.Node) -> bool: - if src.op != "call_method" or src.target != "to_local": - return False - owner = src.args[0] if src.args else None - return isinstance(owner, fx.Node) and owner.op in ("placeholder", "get_attr") - - -def rewrite_weight_ag_to_copy_engine(graph: fx.GraphModule) -> int: - """Retarget marked weight gathers to ``magi::symm_all_gather``. Returns count rewritten.""" - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED - - rewritten = 0 - skipped = 0 - for node in graph.graph.nodes: - if node.op != "call_function": - continue - if not node.meta.get("magi_fsdp_weight_ag"): - continue - if node.target is _ALL_GATHER: - if not _input_is_symm_shard(node): - skipped += 1 - continue - node.target = SYMM_ALL_GATHER - rewritten += 1 - elif node.target is _ALL_GATHER_COALESCED: - if not _coalesced_inputs_are_symm_shards(node): - skipped += 1 - continue - node.target = SYMM_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 (input is a cast of the shard, or the shard is unevenly split)", - rewritten, - skipped, - ) - return rewritten diff --git a/magi_compiler/symm_mem/__init__.py b/magi_compiler/symm_mem/__init__.py index 5d501dd..78ae3ec 100644 --- a/magi_compiler/symm_mem/__init__.py +++ b/magi_compiler/symm_mem/__init__.py @@ -19,16 +19,15 @@ whether the copy-engine transport is available. Import that module by path. """ +from .bind import bind_graph_weights, bind_parameters from .symm_buffer import ( ShardEntry, SymmBuffer, - barrier_after_load, + alloc_shard, find_shard_by_layout, + group_name_of, lookup_shard, - materialize_into_buffers, - migrate_to_buffers, - patch_symm_buffer_apply, - register_shard, + publish, registered_buffers, reset_registry, ) @@ -36,13 +35,13 @@ __all__ = [ "ShardEntry", "SymmBuffer", - "barrier_after_load", + "alloc_shard", + "bind_graph_weights", + "bind_parameters", "find_shard_by_layout", + "group_name_of", "lookup_shard", - "materialize_into_buffers", - "migrate_to_buffers", - "patch_symm_buffer_apply", - "register_shard", + "publish", "registered_buffers", "reset_registry", ] diff --git a/magi_compiler/symm_mem/bind.py b/magi_compiler/symm_mem/bind.py new file mode 100644 index 0000000..0676682 --- /dev/null +++ b/magi_compiler/symm_mem/bind.py @@ -0,0 +1,324 @@ +# 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 group_name_of, lookup_shard, open_buffer, publish, 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) + publish() + + 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) + publish() + 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 + + +_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 + + +def _move(plan: list[BindCandidate]) -> int: + """Allocate, fill and repoint 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] + try: + buffer = open_buffer(head.local.dtype, head.local.device, head.group_name, (c.local.numel() for c in members)) + 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 in members: + symm = buffer.take(c.local.shape) + symm.copy_(c.local) + register_shard(symm, buffer) + # In place: Dynamo already guarded these exact objects. + c.local.data = symm + allocated += 1 + + torch.cuda.empty_cache() + 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 index 8e1bf92..f3d362d 100644 --- a/magi_compiler/symm_mem/symm_buffer.py +++ b/magi_compiler/symm_mem/symm_buffer.py @@ -12,28 +12,28 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Symmetric-memory windows for copy-engine weight all-gather. + +Shards are suballocated from pooled windows; runtime gathers look them up by +device pointer. Which weights get an allocation is ``bind``'s job. +""" + from __future__ import annotations -import inspect from dataclasses import dataclass import torch import torch.distributed as dist -import torch.nn as nn from magi_compiler.utils import magi_logger -# Only hijack ``to_empty``'s lambda. ``.cuda()`` / ``.to()`` / ``.float()`` also -# go through ``_apply``; intercepting those would change builder semantics. -_TO_EMPTY_LAMBDA = "Module.to_empty.." - class SymmBuffer: """One symmetric-memory window, suballocated to many weight shards. - One window per (decorated block, dtype, process group): a single rendezvous, - and because every rank walks the module tree in the same order, offset ``k`` - is the same shard on every peer. Two meshes sharing a dtype get two windows. + The driver caps windows at 128 per process regardless of size, so pooling is + required. Every rank walks the same plan, so offset ``k`` is the same shard + on every peer. """ # 256 bf16 elems = 512B, which the copy engine wants for peak throughput. @@ -45,25 +45,27 @@ def __init__(self, dtype: torch.dtype, device: torch.device, group_name: str) -> self.group_name = group_name self.buf: torch.Tensor | None = None self.handle = None - self.peers: list[torch.Tensor] = [] self._reserved = 0 self._cursor = 0 - # -- build phase ------------------------------------------------------ def reserve(self, numel: int) -> None: + """Book space for one shard. Call for every member before ``commit``.""" self._reserved += self._round(numel) def commit(self) -> None: + """Open the window. One ``rendezvous`` for every shard it will hold.""" import torch.distributed._symmetric_memory as symm_mem symm_mem.enable_symm_mem_for_group(self.group_name) self.buf = symm_mem.empty(self._reserved, dtype=self.dtype, device=self.device) self.handle = symm_mem.rendezvous(self.buf, self.group_name) - # Slice each peer's whole window once. VMM maps every window into one - # VA space so these are directly copyable; NCCL-backend peer pointers are not. - self.peers = [self.handle.get_buffer(r, (self._reserved,), self.dtype) for r in range(self.handle.world_size)] def take(self, shape: torch.Size | tuple[int, ...]) -> torch.Tensor: + """Hand out the next slot as a tensor whose storage starts at the slot. + + Not a slice of ``buf``: Dynamo memoized ``storage_offset == 0`` for these + parameters, and a mid-window slice would contradict the shape env. + """ numel = 1 for s in shape: numel *= int(s) @@ -74,9 +76,8 @@ def take(self, shape: torch.Size | tuple[int, ...]) -> torch.Tensor: f"symmetric buffer overflow: wanted {self._cursor} elems, reserved {self._reserved}. " "The sizing walk and the dispensing walk must visit the same shards in the same order." ) - return self.buf[off : off + numel].view(shape) + return self.handle.get_buffer(self.handle.rank, tuple(int(s) for s in shape), self.dtype, off) - # -- query ------------------------------------------------------------ @property def nbytes(self) -> int: return self._reserved * self.dtype.itemsize @@ -85,9 +86,12 @@ def offset_of(self, t: torch.Tensor) -> int: return (t.data_ptr() - self.buf.data_ptr()) // self.buf.element_size() def peer_views(self, t: torch.Tensor) -> list[torch.Tensor]: - """``world_size`` views of the same shard, one per rank, in rank order.""" - off, numel = self.offset_of(t), t.numel() - return [p[off : off + numel].view(t.shape) for p in self.peers] + """``world_size`` views of the same shard, one per rank. + + They borrow the window mapping; ``self.buf`` keeps it alive. + """ + off, shape = self.offset_of(t), tuple(int(s) for s in t.shape) + return [self.handle.get_buffer(r, shape, self.dtype, off) for r in range(self.handle.world_size)] def contains(self, t: torch.Tensor) -> bool: if self.buf is None: @@ -102,7 +106,7 @@ def _round(cls, numel: int) -> int: @dataclass(frozen=True) class ShardEntry: - """What the run-time gather needs to know about one local shard.""" + """What a run-time gather needs to know about one local shard.""" buffer: SymmBuffer offset: int @@ -113,20 +117,65 @@ class ShardEntry: def shape(self) -> tuple[int, ...]: return tuple(self.local.shape) + @property + def dtype(self) -> torch.dtype: + return self.local.dtype + -# Keyed by ``data_ptr()``: the gather op only sees a plain tensor. _SHARD_REGISTRY: dict[int, ShardEntry] = {} _BUFFERS: list[SymmBuffer] = [] -_BARRIER_DONE = False +_UNPUBLISHED = False + + +def open_buffer(dtype: torch.dtype, device: torch.device, group_name: str, numels) -> SymmBuffer: + """Open one window big enough for ``numels``, and track it for ``publish``.""" + global _UNPUBLISHED + buffer = SymmBuffer(dtype, device, group_name) + for numel in numels: + buffer.reserve(int(numel)) + buffer.commit() + _BUFFERS.append(buffer) + _UNPUBLISHED = True + return buffer def register_shard(local: torch.Tensor, buffer: SymmBuffer) -> ShardEntry: + """Record a slot so the run-time gather can find its peer views.""" entry = ShardEntry(buffer=buffer, offset=buffer.offset_of(local), local=local, peer_views=tuple(buffer.peer_views(local))) _SHARD_REGISTRY[local.data_ptr()] = entry return entry +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.take(shape) + register_shard(shard, buffer) + return shard + + +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] + + 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) @@ -134,211 +183,35 @@ def registered_buffers() -> list[SymmBuffer]: return list(_BUFFERS) -def find_shard_by_layout(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor | None: - """Any registered shard with this layout -- the cost model cannot replay a gather on a generic ``empty`` (no peers).""" +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.local.dtype == dtype: + if entry.shape == want and entry.dtype == dtype: return entry.local return None -def reset_registry() -> None: - """Test-only: drop every buffer so a new model can be built in-process.""" - global _BARRIER_DONE - _SHARD_REGISTRY.clear() - _BUFFERS.clear() - _BARRIER_DONE = False - +def publish() -> None: + """Barrier so every rank has copied its shards in before any peer reads. -def barrier_after_load() -> None: - """Publish every rank's freshly written shards, once per process. - - Must run after weights are loaded and before the first peer read. + Once per bind, not per step: the weights never change again. """ - global _BARRIER_DONE - if _BARRIER_DONE or not _BUFFERS: + global _UNPUBLISHED + if not _UNPUBLISHED: return if dist.is_available() and dist.is_initialized(): torch.cuda.synchronize() dist.barrier() - _BARRIER_DONE = True - magi_logger.info( - "SymmBuffer: published %d buffer(s), %.1f MiB, %d shards; steady state is barrier-free", - len(_BUFFERS), - sum(b.nbytes for b in _BUFFERS) / 2**20, - len(_SHARD_REGISTRY), - ) - - -def _is_gatherable_shard(t: object) -> bool: - """An EVENLY split Shard(0) DTensor on a 1-D mesh -- the only placement the - copy-engine gather handles. - - Uneven ``Shard(0)`` (``dim0 % world != 0``) is excluded on purpose, for three - reasons that all trace back to the trailing ranks owning fewer rows: the window - would be sized from a rank-dependent local numel, so ``rendezvous`` rejects it - outright once the difference survives ``ALIGN``; below that threshold it is worse - than an error, because every shard after the uneven one lands at a different - offset per rank while ``peer_views`` slices each peer at the LOCAL offset; and the - gather copies one fixed-size slab per peer, which would read past a shorter - peer's rows. ``dim0`` and the mesh size are identical on every rank, so all - ranks drop the same parameters and the offset walk stays symmetric. Dropped - weights keep their ordinary allocation and gather over NCCL. - """ - from torch.distributed.tensor import DTensor, Shard - - if not isinstance(t, DTensor): - return False - placements = t.placements - if not (len(placements) == 1 and isinstance(placements[0], Shard) and placements[0].dim == 0): - return False - return int(t.shape[0]) % int(t.device_mesh.size(0)) == 0 - - -def _group_name_of(t) -> str | None: - try: - return t.device_mesh._dim_group_names[0] - except Exception: # noqa: BLE001 - return None - - -def _buffer_key(t) -> tuple[torch.dtype, str]: - """One window per (dtype, process group). Same dtype on two meshes (gaga4 FSDP + edp) must not share a window.""" - group_name = _group_name_of(t) - if group_name is None: - raise RuntimeError(f"cannot resolve the process group of a Shard(0) parameter on mesh {t.device_mesh}") - return (t.dtype, group_name) - - -def _apply_order_entries(mod: nn.Module): - """``(owner, name, param)`` in ``_apply`` order: post-order, every ``_parameters`` entry. - - Not ``named_parameters()`` (pre-order, dedups shared tensors). A different - walk would break cross-rank offset symmetry. Walking ``_parameters`` also - finds SimpleFSDP weights in ``parametrizations.weight.original``. - """ - for child in mod.children(): - yield from _apply_order_entries(child) - for name, p in mod._parameters.items(): - if p is not None: - yield mod, name, p - - -def _plan_buffers(shards: list, device: torch.device) -> dict[tuple[torch.dtype, str], SymmBuffer]: - """Size and commit one window per (dtype, group). Dedup by identity so a tied weight reserves a single slot.""" - buffers: dict[tuple[torch.dtype, str], SymmBuffer] = {} - seen: set[int] = set() - for p in shards: - if id(p) in seen: - continue - seen.add(id(p)) - key = _buffer_key(p) - buffer = buffers.get(key) - if buffer is None: - buffer = buffers[key] = SymmBuffer(p.dtype, device, key[1]) - buffer.reserve(p._local_tensor.numel()) - - for buffer in buffers.values(): - buffer.commit() # the only collective, once per window - _BUFFERS.extend(buffers.values()) - return buffers - - -def materialize_into_buffers(mod: nn.Module, device: torch.device) -> dict[tuple[torch.dtype, str], SymmBuffer]: - """Size windows for ``mod``'s Shard(0) shards while they are still on meta. Non-gatherable params are left to the caller.""" - shards = [p for _, _, p in _apply_order_entries(mod) if _is_gatherable_shard(p)] - if not shards: - return {} - return _plan_buffers(shards, device) - - -def migrate_to_buffers(root: nn.Module) -> dict[tuple[torch.dtype, str], SymmBuffer]: - """Copy already-allocated Shard(0) shards into symmetric memory. - - Used when ``magi_compile(model, ...)`` is given a live model rather than a - meta + ``to_empty`` path. ``load_state_dict(assign=True)`` after this would - replace buffer views with ordinary tensors; the gather then rejects them. - """ - entries = [(m, n, p) for m, n, p in _apply_order_entries(root) if _is_gatherable_shard(p)] - if not entries: - return {} - - device = entries[0][2]._local_tensor.device - if device.type != "cuda": - raise RuntimeError(f"symmetric memory needs the shards on cuda, found {device}") - buffers = _plan_buffers([p for _, _, p in entries], device) - - from torch.distributed.tensor import DTensor - - views: dict[int, torch.Tensor] = {} - for owner, name, p in entries: - local = views.get(id(p)) - if local is None: - buffer = buffers[_buffer_key(p)] - local = views[id(p)] = buffer.take(p._local_tensor.shape) - local.copy_(p._local_tensor) - register_shard(local, buffer) - moved = DTensor.from_local(local, p.device_mesh, p.placements, run_check=False) - owner.register_parameter(name, nn.Parameter(moved, requires_grad=p.requires_grad)) - + _UNPUBLISHED = False magi_logger.info( - "SymmBuffer: migrated %d shard(s) into %.1f MiB across %d window(s)", - len(views), - sum(b.nbytes for b in buffers.values()) / 2**20, - len(buffers), + "SymmBuffer: published %d shard(s), %.1f MiB total", len(_BUFFERS), sum(b.nbytes for b in _BUFFERS) / 2**20 ) - return buffers -def patch_symm_buffer_apply(cls: type[nn.Module]) -> None: - """Install the ``_apply`` interception on a decorated class. - - Mirrors ``_patch_cpu_offload_apply``: take over for ``to_empty``'s lambda, delegate everything else. - """ - if getattr(cls, "_magi_symm_apply_patched", False): - return - orig_apply = cls._apply - magi_logger.info("SymmBuffer: intercepting %s._apply for copy-engine FSDP", cls.__name__) - - def _symm_apply(self, fn, recurse: bool = True): - if getattr(fn, "__qualname__", "") != _TO_EMPTY_LAMBDA: - return orig_apply(self, fn, recurse) - if getattr(self, "_magi_symm_buffers", None) is not None: - return orig_apply(self, fn, recurse) - - device = torch.device(inspect.getclosurevars(fn).nonlocals["device"]) - from torch.distributed.tensor import DTensor - - buffers = materialize_into_buffers(self, device) - if not buffers: - return orig_apply(self, fn, recurse) - - views: dict[int, torch.Tensor] = {} - - def materialize(t: torch.Tensor) -> torch.Tensor: - if not _is_gatherable_shard(t): - return torch.empty_like(t, device=device) - # Tied weight: same view so tying survives materialization. - local = views.get(id(t)) - if local is None: - buffer = buffers[_buffer_key(t)] - local = views[id(t)] = buffer.take(t._local_tensor.shape) - register_shard(local, buffer) - return DTensor.from_local(local, t.device_mesh, t.placements, run_check=False) - - # Do not forge to_empty's qualname: a nested decorated block must fail - # the check above and delegate, so its params land in *this* buffer. - out = orig_apply(self, materialize, recurse) - self._magi_symm_buffers = buffers - magi_logger.info( - "SymmBuffer: %s materialized %d shard(s) into %.1f MiB across %d window(s)", - cls.__name__, - len(views), - sum(b.nbytes for b in buffers.values()) / 2**20, - len(buffers), - ) - return out - - cls._apply = _symm_apply - cls._magi_symm_apply_patched = True +def reset_registry() -> None: + """Drop every window and shard. Tests only -- frees the symmetric allocations.""" + global _UNPUBLISHED + _SHARD_REGISTRY.clear() + _BUFFERS.clear() + _UNPUBLISHED = False 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 720bce6..187d04f 100644 --- a/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py +++ b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py @@ -169,21 +169,15 @@ def fn(x, w0, shard): ce_shards: list = [] if args.copy_engine: - from magi_compiler.symm_mem import SymmBuffer, register_shard + from magi_compiler.symm_mem import alloc_shard, publish from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER - buffer = SymmBuffer(torch.bfloat16, torch.device("cuda", dev), grp) - for _ in range(N_CE_LAYERS): - buffer.reserve(H * H) - buffer.commit() for i in range(N_CE_LAYERS): - s = buffer.take((H, H)) + s = alloc_shard((H, H), torch.bfloat16, torch.device("cuda", dev), grp) s.normal_(0.0, H**-0.5).add_(0.01 * i) - register_shard(s, buffer) ce_shards.append(s) # A peer read is only legal once that peer has written its shard. - torch.cuda.synchronize() - dist.barrier() + publish() def fn(x, w0, shards): # noqa: F811 - deliberately replaces the NCCL variant y = (x @ w0).relu() 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..91bcc8f --- /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 SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + + names = { + _NCCL_AG: "nccl", + _NCCL_AG_COALESCED: "nccl_coalesced", + SYMM_ALL_GATHER: "symm", + SYMM_ALL_GATHER_COALESCED: "symm_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 = ({"nccl_coalesced": 1, "symm_coalesced": 1}, [("nccl_coalesced", 2), ("symm_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/test_fsdp_overlap_bucket.py b/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py index 6272474..72e1312 100644 --- a/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py +++ b/tests/feature_tests/fsdp/test_fsdp_overlap_bucket.py @@ -56,6 +56,8 @@ def _build_ag_graph(specs, world=2, group="grp0"): 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): @@ -75,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) diff --git a/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py b/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py index d7babbf..e21977f 100644 --- a/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py +++ b/tests/feature_tests/fsdp/test_fsdp_overlap_lowering.py @@ -141,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 @@ -161,7 +163,7 @@ def test_lowering_skips_non_weight_input(dist_1rank): @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 lower_prim_redistribute_to_collectives + 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 @@ -170,7 +172,7 @@ def test_lowering_rewrites_method_shaped_redistribute(dist_1rank): 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 x.meta.get("magi_fsdp_weight_ag")]) == 1 + assert len([x for x in gm.graph.nodes if is_weight_ag(x)]) == 1 @requires_cuda @@ -199,14 +201,14 @@ def test_lowering_records_shard_evenness_on_the_gather(dist_1rank): 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 lower_prim_redistribute_to_collectives + 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 "magi_fsdp_uneven_shard" in ag[0].meta - assert ag[0].meta["magi_fsdp_uneven_shard"] is False # 8 rows over a 1-rank mesh + assert UNEVEN_SHARD in ag[0].meta + assert ag[0].meta[UNEVEN_SHARD] is False # 8 rows over a 1-rank mesh @requires_cuda diff --git a/tests/feature_tests/fsdp/test_profiling_estimator.py b/tests/feature_tests/fsdp/test_profiling_estimator.py index 833107d..9cdc809 100644 --- a/tests/feature_tests/fsdp/test_profiling_estimator.py +++ b/tests/feature_tests/fsdp/test_profiling_estimator.py @@ -705,20 +705,14 @@ def _ce_group_name() -> str: def _register_shards(shapes, dtype=torch.bfloat16): """Register ``shapes`` as real symmetric-memory shards, filled distinctly.""" - from magi_compiler.symm_mem import SymmBuffer, register_shard - - buffer = SymmBuffer(dtype, torch.device("cuda", 0), _ce_group_name()) - for shape in shapes: - buffer.reserve(shape[0] * shape[1]) - buffer.commit() + from magi_compiler.symm_mem import alloc_shard, publish shards = [] for i, shape in enumerate(shapes): - s = buffer.take(shape) + s = alloc_shard(shape, dtype, torch.device("cuda", 0), _ce_group_name()) s.fill_(i + 1) - register_shard(s, buffer) shards.append(s) - torch.cuda.synchronize() + publish() return shards 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) diff --git a/tests/feature_tests/symm_helper/verify_symm_e2e.py b/tests/feature_tests/symm_helper/verify_symm_e2e.py index b84ac08..bfc0709 100644 --- a/tests/feature_tests/symm_helper/verify_symm_e2e.py +++ b/tests/feature_tests/symm_helper/verify_symm_e2e.py @@ -204,7 +204,19 @@ def spy_rewrite(graph): model.to_empty(device=device) _load_into_shards(model, state) - # (1) placement: the block's shards are in a window, the head's are not. + # (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.")} @@ -217,16 +229,6 @@ def spy_rewrite(graph): f"window(s) ({buffer_mib:.0f} MiB), {len(head_in_buffer)} stray -> {'ok' if placement_ok else 'WRONG'}" ) - # (2) + (3): 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 - expect_rewrites = args.n_layers if ce else 0 rewrite_ok = n_rewritten == expect_rewrites log( diff --git a/tests/feature_tests/test_symm_ag_rewrite.py b/tests/feature_tests/test_copy_engine.py similarity index 63% rename from tests/feature_tests/test_symm_ag_rewrite.py rename to tests/feature_tests/test_copy_engine.py index 7ddb087..07d23bc 100644 --- a/tests/feature_tests/test_symm_ag_rewrite.py +++ b/tests/feature_tests/test_copy_engine.py @@ -12,17 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the copy-engine retarget pass (step 3 of the landing order). - -The pass is deliberately tiny -- swap one node's target -- so what -is worth testing is what it *refuses* to touch. A gather whose input is a cast -or a pad reads a tensor the caching allocator produced, not the symmetric window, -and retargeting it would make the operator reject it at run time. A gather that -was never marked as a SimpleFSDP weight gather (CP, TP, MoE) must stay on NCCL -even in copy-engine mode. - -Built on the real lowering pass so the node shapes are the ones production -produces, with a 1-rank mesh (nothing here touches the transport). +"""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 @@ -83,6 +84,8 @@ def _graph_with_gathers( """ 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() @@ -105,8 +108,7 @@ def _graph_with_gathers( ag = g.call_function(_AG, (cur, 1, "dummy_group")) ag.meta["example_value"] = local._local_tensor.new_empty(local._local_tensor.shape) if marked: - ag.meta["magi_fsdp_weight_ag"] = True - ag.meta["magi_fsdp_uneven_shard"] = uneven + mark_weight_ag(ag, uneven=uneven) outs.append(g.call_function(_WAIT, (ag,))) g.output(tuple(outs)) return fx.GraphModule(torch.nn.Module(), g) @@ -116,55 +118,54 @@ def _gathers(gm, target): return [n for n in gm.graph.nodes if n.op == "call_function" and n.target is target] -@requires_cuda -def test_marked_gathers_are_retargeted_and_waits_untouched(mesh_1rank): - from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine +def _candidates(gm): + from magi_compiler.passes.fsdp_overlap import copy_engine_weight_candidates - gm = _graph_with_gathers(mesh_1rank, 4) - assert rewrite_weight_ag_to_copy_engine(gm) == 4 + return copy_engine_weight_candidates(gm) - 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 +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 -@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 +# --------------------------------------------------------------------------- +# Selection +# --------------------------------------------------------------------------- +@requires_cuda +def test_plain_weight_gathers_are_candidates(mesh_1rank): gm = _graph_with_gathers(mesh_1rank, 3) - rewrite_weight_ag_to_copy_engine(gm) - assert all(n.args[1:3] == (1, "dummy_group") for n in _gathers(gm, _symm_op())) + 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_stay_on_nccl(mesh_1rank, derive): - """A cast or pad output is not in the symmetric window; retargeting it would - be rejected at run time, so it must stay on NCCL. +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_shard_stays_on_nccl_without_a_pad``. + ``test_uneven_shards_are_not_candidates``. """ - from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine - gm = _graph_with_gathers(mesh_1rank, 3, derive=derive) - assert rewrite_weight_ag_to_copy_engine(gm) == 0 - assert len(_gathers(gm, _AG)) == 3 - assert not _gathers(gm, _symm_op()) + assert _candidates(gm) == [] @requires_cuda -def test_uneven_shard_stays_on_nccl_without_a_pad(mesh_1rank): +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 @@ -172,57 +173,65 @@ def test_uneven_shard_stays_on_nccl_without_a_pad(mesh_1rank): 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 is made from the flag, which is derived from F and world and is + So the decision comes off the flag, which is derived from F and world and is therefore the same everywhere. """ - from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine - gm = _graph_with_gathers(mesh_1rank, 3, uneven=True) - assert rewrite_weight_ag_to_copy_engine(gm) == 0 - assert len(_gathers(gm, _AG)) == 3 - assert not _gathers(gm, _symm_op()) + assert _candidates(gm) == [] @requires_cuda -def test_uneven_shard_does_not_join_a_copy_engine_bucket(mesh_1rank): - """Bucket membership is a transport decision too: a bucket is one submission, and - an uneven weight inside a copy-engine bucket would take the whole bucket with it.""" - from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED - - gm = _graph_with_gathers(mesh_1rank, 3, uneven=True) - assert lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine") == 0 - assert len(_gathers(gm, _AG)) == 3 - assert not _gathers(gm, SYMM_ALL_GATHER_COALESCED) +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_uneven_shard_does_not_hold_back_its_even_neighbours(mesh_1rank): - """Only the uneven weight loses the copy engine. Excluding it from the bucket - must not split the even weights around it into separate buckets either, or the - fix would cost throughput on every model with one odd-shaped weight.""" - from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED +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) - _gathers(gm, _AG)[1].meta["magi_fsdp_uneven_shard"] = True + assert _pretend_bound(gm) == 4 + assert rewrite_weight_ag_to_copy_engine(gm) == 4 - assert lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine") == 1 - coalesced = _gathers(gm, SYMM_ALL_GATHER_COALESCED) - assert len(coalesced) == 1 - assert len(coalesced[0].args[0]) == 3 # the three even weights, in one bucket - assert len(_gathers(gm, _AG)) == 1 # the uneven one, alone, on NCCL + 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_unmarked_gathers_are_left_alone(mesh_1rank): - """CP / TP / MoE gathers are never marked, and must survive copy-engine mode - untouched -- the transports coexist in one graph.""" +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, marked=False) + 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 @@ -239,68 +248,94 @@ def test_mixed_graph_splits_by_transport(mesh_1rank): 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_end_to_end_lowering_then_rewrite(mesh_1rank): - """The pass has to match what the lowering pass really emits, not what this - file thinks it emits.""" +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 test_fsdp_overlap_lowering import _build_redistribute_graph - from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph + from magi_compiler.passes.fsdp_overlap import lower_prim_redistribute_to_collectives, rewrite_weight_ag_to_copy_engine gm = _build_redistribute_graph(mesh_1rank, "model_fc1_weight_parameter") - lower_and_bucket_full_graph(gm, "none", transport="copy_engine") + 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 len(symm) == 1 assert _gathers(gm, _WAIT)[0].args[0] is symm[0] +# --------------------------------------------------------------------------- +# Bucketing carries the mark +# --------------------------------------------------------------------------- @requires_cuda -def test_copy_engine_buckets_then_rewrites_coalesced(mesh_1rank): - """Phase-1 wrap: SymmBuffer gathers are bucketed first, then the coalesced - launch is retargeted. Members stay separate dests underneath.""" - from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED +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 SYMM_ALL_GATHER_COALESCED gm = _graph_with_gathers(mesh_1rank, 4) - n = lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine") - assert n == 1 - assert not _gathers(gm, _AG) + _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 not _gathers(gm, SYMM_ALL_GATHER) assert len(_gathers(gm, SYMM_ALL_GATHER_COALESCED)) == 1 - waits = _gathers(gm, _WAIT) - assert len(waits) == 4 + assert len(_gathers(gm, _WAIT)) == 4 @requires_cuda -def test_copy_engine_does_not_bucket_cast_gathers(mesh_1rank): - """Cast outputs are not SymmBuffer shards; they must stay on NCCL and not join a CE bucket.""" - from magi_compiler.passes.fsdp_overlap import lower_and_bucket_full_graph +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 SYMM_ALL_GATHER_COALESCED - gm = _graph_with_gathers(mesh_1rank, 3, derive="cast") - n = lower_and_bucket_full_graph(gm, "coalesced", bucket_size_bytes=0, transport="copy_engine") - assert n == 0 - assert len(_gathers(gm, _AG)) == 3 + 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, SYMM_ALL_GATHER_COALESCED) -def _coalesced_graph(mesh, n: int, *, derive: str | None = None, marked: bool = True): +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. + 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): @@ -309,15 +344,14 @@ def _coalesced_graph(mesh, n: int, *, derive: str | None = None, marked: bool = w.meta["example_value"] = local 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) 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: - ag.meta["magi_fsdp_weight_ag"] = True + 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)) @@ -327,7 +361,7 @@ def _coalesced_graph(mesh, n: int, *, derive: str | None = None, marked: bool = @requires_cuda -def test_coalesced_of_symm_shards_is_retargeted(mesh_1rank): +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 SYMM_ALL_GATHER_COALESCED @@ -337,20 +371,6 @@ def test_coalesced_of_symm_shards_is_retargeted(mesh_1rank): assert len(_gathers(gm, SYMM_ALL_GATHER_COALESCED)) == 1 -@requires_cuda -def test_coalesced_is_all_or_nothing(mesh_1rank): - """A bucket is one submission, so it can only go to the copy engine if - *every* member is an SymmBuffer shard -- one cast 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 rewrite_weight_ag_to_copy_engine - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED - - gm = _coalesced_graph(mesh_1rank, 3, derive="cast") - assert rewrite_weight_ag_to_copy_engine(gm) == 0 - assert len(_gathers(gm, _AG_COALESCED)) == 1 - assert not _gathers(gm, SYMM_ALL_GATHER_COALESCED) - - @requires_cuda def test_unmarked_coalesced_stays_on_nccl(mesh_1rank): from magi_compiler.passes.fsdp_overlap import rewrite_weight_ag_to_copy_engine @@ -358,15 +378,3 @@ def test_unmarked_coalesced_stays_on_nccl(mesh_1rank): gm = _coalesced_graph(mesh_1rank, 2, marked=False) assert rewrite_weight_ag_to_copy_engine(gm) == 0 assert len(_gathers(gm, _AG_COALESCED)) == 1 - - -@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 diff --git a/tests/feature_tests/test_symm_all_gather.py b/tests/feature_tests/test_symm_all_gather.py index b4dcef7..464cb4d 100644 --- a/tests/feature_tests/test_symm_all_gather.py +++ b/tests/feature_tests/test_symm_all_gather.py @@ -73,42 +73,28 @@ def _clean_state(): def _symm_model(mesh, hidden: int = 64, n_layers: int = 3, dtype=torch.bfloat16): - """A meta-built, Shard(0)-sharded model materialized into a symmetric buffer, - exactly the shape step 1 produces.""" + """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 materialize_into_buffers + 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("meta"): + 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)]))) - device = torch.device("cuda", 0) - from torch.distributed.tensor import DTensor - - from magi_compiler.symm_mem.symm_buffer import _buffer_key, register_shard - - buffers = materialize_into_buffers(model, device) - views: dict[int, torch.Tensor] = {} - - def materialize(t): - if isinstance(t, DTensor): - buffer = buffers[_buffer_key(t)] - local = buffer.take(t._local_tensor.shape) - register_shard(local, buffer) - views[id(t)] = local - return DTensor.from_local(local, t.device_mesh, t.placements, run_check=False) - return torch.empty_like(t, device=device) - - materialize.__qualname__ = "Module.to_empty.." - nn.Module._apply(model, materialize) + assert bind_parameters(list(model.parameters())) == n_layers shards = [p._local_tensor for p in model.parameters()] for i, s in enumerate(shards): diff --git a/tests/feature_tests/test_symm_bind.py b/tests/feature_tests/test_symm_bind.py new file mode 100644 index 0000000..4c02a06 --- /dev/null +++ b/tests/feature_tests/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 SYMM_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 SYMM_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 SYMM_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 SYMM_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 SYMM_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 SYMM_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/test_symm_buffer.py b/tests/feature_tests/test_symm_buffer.py deleted file mode 100644 index 3b40ed7..0000000 --- a/tests/feature_tests/test_symm_buffer.py +++ /dev/null @@ -1,503 +0,0 @@ -# Copyright (c) 2026 SandAI. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Materialization tests for ``fsdp_config.transport="copy_engine"``. - -These cover step 1 of the landing order: the weights move into symmetric memory -and nothing else changes. The interesting property is *scope* -- the decorated -class must claim exactly its own subtree, from a ``to_empty`` issued on the root -model, without the builder participating. - -Single rank on purpose: the placement, scoping, aliasing and fallback logic is -rank-independent, and a 1-rank symmetric window exercises the same allocation -path. The multi-rank peer reads are covered by -``example/inference/fsdp_overlap/verify_symm_param_hook.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") - - -@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(autouse=True) -def _clean_registry(): - from magi_compiler.symm_mem import reset_registry - - reset_registry() - yield - reset_registry() - - -def _decorate(cls: type, transport: str) -> type: - """Run the real decorator, so the config wiring is under test too.""" - from magi_compiler import magi_compile - from magi_compiler.config import CompileMode - - def patch(cfg): - cfg.compile_mode = CompileMode.MAGI_COMPILE - cfg.fsdp_config.enable_fsdp = True - cfg.fsdp_config.transport = transport - return cfg - - return magi_compile(cls, config_patch=patch, dynamic_arg_dims={"x": 0}) - - -def _shard(model: nn.Module, mesh, placement=None) -> None: - """Wrap every parameter as a DTensor, like torchtitan ``data_parallel``.""" - from torch.distributed.tensor import Shard, distribute_tensor - - placements = [placement or Shard(0)] - 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, placements))) - - -class Layer(nn.Module): - def __init__(self, hidden: int, dtype: torch.dtype): - super().__init__() - self.wq = nn.Linear(hidden, hidden, bias=False, dtype=dtype) - self.wo = nn.Linear(hidden, hidden, bias=False, dtype=dtype) - - def forward(self, x): - return self.wo(self.wq(x)) - - -def _make_root(block_cls: type, hidden: int, n_layers: int, dtype: torch.dtype) -> nn.Module: - class Root(nn.Module): - """``to_empty`` is called on this, never on the decorated block.""" - - def __init__(self): - 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)) - - with torch.device("meta"): - return Root() - - -def _make_block_cls(transport: str) -> type: - class Block(nn.Module): - def __init__(self, hidden: int, n_layers: int, dtype: torch.dtype): - super().__init__() - self.layers = nn.ModuleList([Layer(hidden, dtype) for _ in range(n_layers)]) - - def forward(self, x): - for layer in self.layers: - x = layer(x) - return x - - return _decorate(Block, transport) - - -@requires_cuda -def test_root_to_empty_puts_block_shards_in_buffer(mesh_1rank): - """The decorated block claims its own subtree; the rest stays ordinary.""" - from magi_compiler.symm_mem import lookup_shard, registered_buffers - - hidden, n_layers = 64, 3 - model = _make_root(_make_block_cls("copy_engine"), hidden, n_layers, torch.bfloat16) - _shard(model, mesh_1rank) - model.to_empty(device=torch.device("cuda", 0)) - - buffers = registered_buffers() - assert len(buffers) == 1, "one window per (dtype, group), not one per weight" - buffer = buffers[0] - - block_shards = list(model.block.parameters()) - assert len(block_shards) == 2 * n_layers - assert all(buffer.contains(p._local_tensor) for p in block_shards) - assert all(lookup_shard(p._local_tensor.data_ptr()) is not None for p in block_shards) - - # Outside the decorated class: untouched by the interception. - assert not buffer.contains(model.head.weight._local_tensor) - assert lookup_shard(model.head.weight._local_tensor.data_ptr()) is None - - # Still ordinary, working DTensors on real storage. - assert all(p._local_tensor.device.type == "cuda" for p in block_shards) - assert not any(p.is_meta for p in block_shards) - - -@requires_cuda -def test_nccl_transport_leaves_allocation_alone(mesh_1rank): - """The default transport must not install the interception at all.""" - from magi_compiler.symm_mem import registered_buffers - - model = _make_root(_make_block_cls("nccl"), 64, 2, torch.bfloat16) - _shard(model, mesh_1rank) - model.to_empty(device=torch.device("cuda", 0)) - - assert registered_buffers() == [] - assert all(p._local_tensor.device.type == "cuda" for p in model.parameters()) - - -@requires_cuda -def test_peer_view_round_trips_through_buffer(mesh_1rank): - """A shard written locally must be visible through its own peer view: this is - the addressing the copy-engine gather depends on.""" - from magi_compiler.symm_mem import lookup_shard - - model = _make_root(_make_block_cls("copy_engine"), 64, 2, torch.bfloat16) - _shard(model, mesh_1rank) - model.to_empty(device=torch.device("cuda", 0)) - - for i, p in enumerate(model.block.parameters()): - p._local_tensor.fill_(i + 1) - torch.cuda.synchronize() - - for i, p in enumerate(model.block.parameters()): - entry = lookup_shard(p._local_tensor.data_ptr()) - assert len(entry.peer_views) == 1 - assert torch.equal(entry.peer_views[0], p._local_tensor), f"shard {i} aliases the wrong bytes" - - -@requires_cuda -def test_tied_weights_share_one_slot(mesh_1rank): - """A shared parameter is visited once per referencing module, so it must get - one slot and keep its identity -- two allocations would silently untie it and - overflow a window sized by a deduped walk. - - Note the tie has to be (re)established *after* sharding: torchtitan's - ``data_parallel`` calls ``distribute_tensor`` per module, which replaces each - entry with its own DTensor and unties them on its own. - """ - from magi_compiler.symm_mem import registered_buffers - - hidden = 64 - - class TiedBlock(nn.Module): - def __init__(self, hidden: int, n_layers: int, dtype: torch.dtype): - super().__init__() - self.a = nn.Linear(hidden, hidden, bias=False, dtype=dtype) - self.b = nn.Linear(hidden, hidden, bias=False, dtype=dtype) - - def forward(self, x): - return self.b(self.a(x)) - - model = _make_root(_decorate(TiedBlock, "copy_engine"), hidden, 1, torch.bfloat16) - _shard(model, mesh_1rank) - model.block.b.weight = model.block.a.weight - model.to_empty(device=torch.device("cuda", 0)) - - buffer = registered_buffers()[0] - assert model.block.a.weight is model.block.b.weight, "tying must survive materialization" - assert buffer.contains(model.block.a.weight._local_tensor) - # One slot in the window, not two. - slot = buffer.ALIGN * ((hidden * hidden + buffer.ALIGN - 1) // buffer.ALIGN) - assert buffer.nbytes == slot * torch.bfloat16.itemsize - - -@requires_cuda -def test_nested_decorated_block_shares_the_outer_buffer(mesh_1rank): - """A decorated block inside a decorated block must not open a second window: - the inner one has to fail the lambda check and delegate.""" - from magi_compiler.symm_mem import registered_buffers - - inner_cls = _decorate( - type( - "Inner", - (nn.Module,), - { - "__init__": lambda self, hidden, dtype: ( - nn.Module.__init__(self), - setattr(self, "lin", nn.Linear(hidden, hidden, bias=False, dtype=dtype)), - )[0], - "forward": lambda self, x: self.lin(x), - }, - ), - "copy_engine", - ) - - class Outer(nn.Module): - def __init__(self, hidden: int, n_layers: int, dtype: torch.dtype): - super().__init__() - self.own = nn.Linear(hidden, hidden, bias=False, dtype=dtype) - self.inner = inner_cls(hidden, dtype) - - def forward(self, x): - return self.inner(self.own(x)) - - model = _make_root(_decorate(Outer, "copy_engine"), 64, 1, torch.bfloat16) - _shard(model, mesh_1rank) - model.to_empty(device=torch.device("cuda", 0)) - - buffers = registered_buffers() - assert len(buffers) == 1, f"nested decoration opened {len(buffers)} windows" - assert buffers[0].contains(model.block.own.weight._local_tensor) - assert buffers[0].contains(model.block.inner.lin.weight._local_tensor) - - -@requires_cuda -def test_non_shard0_placement_falls_back(mesh_1rank): - """Replicate weights are not gatherable, so they must be allocated normally - rather than silently placed in the buffer.""" - from torch.distributed.tensor import Replicate - - from magi_compiler.symm_mem import registered_buffers - - model = _make_root(_make_block_cls("copy_engine"), 64, 2, torch.bfloat16) - _shard(model, mesh_1rank, placement=Replicate()) - model.to_empty(device=torch.device("cuda", 0)) - - assert registered_buffers() == [] - assert all(p._local_tensor.device.type == "cuda" for p in model.block.parameters()) - - -@requires_cuda -def test_two_process_groups_same_dtype_get_two_windows(mesh_1rank, monkeypatch): - """gaga4 shards bf16 dense weights on the FSDP mesh and bf16 experts on the - orthogonal edp mesh. One window per dtype would either rendezvous on the - wrong group or mix offsets that are only meaningful inside one group. - - Fake group names cannot rendezvous, so ``commit`` is stubbed; the assertion - is that planning opens two windows keyed by (dtype, group). - """ - from torch.distributed.tensor import Shard, distribute_tensor - - from magi_compiler.symm_mem import symm_buffer as sb - - hidden = 64 - a = nn.Parameter(distribute_tensor(torch.empty(hidden, hidden, dtype=torch.bfloat16), mesh_1rank, [Shard(0)])) - b = nn.Parameter(distribute_tensor(torch.empty(hidden, hidden, dtype=torch.bfloat16), mesh_1rank, [Shard(0)])) - - monkeypatch.setattr(sb, "_group_name_of", lambda p, _a=a: "dense_fsdp" if p is _a else "edp") - monkeypatch.setattr(sb.SymmBuffer, "commit", lambda self: None) - buffers = sb._plan_buffers([a, b], torch.device("cuda", 0)) - assert set(buffers) == {(torch.bfloat16, "dense_fsdp"), (torch.bfloat16, "edp")} - assert buffers[(torch.bfloat16, "dense_fsdp")].group_name == "dense_fsdp" - assert buffers[(torch.bfloat16, "edp")].group_name == "edp" - - -@requires_cuda -def test_barrier_after_load_is_idempotent(mesh_1rank): - from magi_compiler.symm_mem import barrier_after_load - - model = _make_root(_make_block_cls("copy_engine"), 64, 2, torch.bfloat16) - _shard(model, mesh_1rank) - model.to_empty(device=torch.device("cuda", 0)) - - barrier_after_load() - barrier_after_load() # a second call must not issue a second collective - - -# --------------------------------------------------------------------------- -# Window suballocation. Every rank must agree on which bytes are which shard, -# and nothing at run time re-derives that -- the gather trusts the offsets. -# --------------------------------------------------------------------------- -@pytest.fixture -def group_name(): - import torch.distributed as dist - - return dist.group.WORLD.group_name - - -def _committed_buffer(group_name, numels, dtype=torch.bfloat16): - from magi_compiler.symm_mem import SymmBuffer - - buffer = SymmBuffer(dtype, torch.device("cuda", 0), group_name) - for n in numels: - buffer.reserve(n) - buffer.commit() - return buffer - - -@requires_cuda -def test_shards_are_dispensed_at_aligned_offsets(mesh_1rank, group_name): - """Slots are padded to ``ALIGN`` for copy-engine throughput, so the second - shard does not start where the first one ends. ``offset_of`` is what the - peer views are built from, so it has to agree with what ``take`` handed out.""" - from magi_compiler.symm_mem import SymmBuffer - - rows, cols = 3, 5 # 15 elems: deliberately not a multiple of ALIGN - buffer = _committed_buffer(group_name, [rows * cols, rows * cols]) - first = buffer.take((rows, cols)) - second = buffer.take((rows, cols)) - - assert buffer.offset_of(first) == 0 - assert buffer.offset_of(second) == SymmBuffer.ALIGN - assert buffer.contains(first) and buffer.contains(second) - assert first.shape == (rows, cols) - - -@requires_cuda -def test_dispensing_more_than_was_reserved_is_an_error(mesh_1rank, group_name): - """The sizing walk and the dispensing walk are two separate traversals; if - they ever disagree the shards silently overlap, so the window must run out - rather than hand back memory reserved for someone else.""" - buffer = _committed_buffer(group_name, [64]) - buffer.take((8, 8)) - with pytest.raises(RuntimeError, match="symmetric buffer overflow"): - buffer.take((8, 8)) - - -@requires_cuda -def test_contains_rejects_memory_outside_the_window(mesh_1rank, group_name): - """``contains`` is how the rewrite decides a weight is gatherable; a caching - allocator tensor must never pass.""" - buffer = _committed_buffer(group_name, [64]) - buffer.take((8, 8)) - 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 must - be a None it can degrade on, not a wrong-dtype shard it would gather.""" - from magi_compiler.symm_mem import find_shard_by_layout, register_shard - - buffer = _committed_buffer(group_name, [8 * 4, 16 * 4]) - small = buffer.take((8, 4)) - large = buffer.take((16, 4)) - register_shard(small, buffer) - register_shard(large, buffer) - - 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 find_shard_by_layout, lookup_shard, register_shard, registered_buffers, reset_registry - - buffer = _committed_buffer(group_name, [8 * 4]) - shard = buffer.take((8, 4)) - register_shard(shard, buffer) - 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 - - -# --------------------------------------------------------------------------- -# migrate_to_buffers -- the live-model path (magi_compile on an already -# materialized model, where there is no to_empty to intercept). -# --------------------------------------------------------------------------- -@requires_cuda -def test_migrate_moves_live_shards_and_keeps_their_values(mesh_1rank): - """Unlike ``to_empty``, this runs on weights that already hold data, so the - copy is load-bearing: dropping it would gather uninitialized memory.""" - from magi_compiler.symm_mem import lookup_shard, migrate_to_buffers - - hidden = 64 - - class Live(nn.Module): - def __init__(self): - super().__init__() - self.a = nn.Linear(hidden, hidden, bias=False, dtype=torch.bfloat16) - self.b = nn.Linear(hidden, hidden, bias=False, dtype=torch.bfloat16) - - model = Live().to("cuda") - _shard(model, mesh_1rank) - before = {n: p._local_tensor.clone() for n, p in model.named_parameters()} - - buffers = migrate_to_buffers(model) - assert len(buffers) == 1 - buffer = next(iter(buffers.values())) - - for name, p in model.named_parameters(): - local = p._local_tensor - assert buffer.contains(local), f"{name} was not migrated" - assert lookup_shard(local.data_ptr()) is not None - assert torch.equal(local, before[name]), f"{name} lost its values" - - -@requires_cuda -def test_migrate_gives_a_tied_weight_one_slot(mesh_1rank): - """Migration rebuilds each parameter, so python identity does not survive -- - what must survive is the storage, or the tie is gone and the window is - overflowed by a walk that sized it once.""" - from magi_compiler.symm_mem import migrate_to_buffers - - hidden = 64 - - class Tied(nn.Module): - def __init__(self): - super().__init__() - self.a = nn.Linear(hidden, hidden, bias=False, dtype=torch.bfloat16) - self.b = nn.Linear(hidden, hidden, bias=False, dtype=torch.bfloat16) - - model = Tied().to("cuda") - _shard(model, mesh_1rank) - model.b.weight = model.a.weight - - buffers = migrate_to_buffers(model) - buffer = next(iter(buffers.values())) - assert model.a.weight._local_tensor.data_ptr() == model.b.weight._local_tensor.data_ptr() - slot = buffer.ALIGN * ((hidden * hidden + buffer.ALIGN - 1) // buffer.ALIGN) - assert buffer.nbytes == slot * torch.bfloat16.itemsize - - -@requires_cuda -def test_migrate_refuses_a_model_that_was_never_materialized(mesh_1rank): - """``migrate_to_buffers`` is the live-model entry point, so being handed a - still-on-meta model is the way it gets misused. Copying from meta silently - produces a window of uninitialized weights, so it has to fail instead.""" - from torch.distributed.tensor import DTensor, Shard - - from magi_compiler.symm_mem import migrate_to_buffers - - with torch.device("meta"): - model = nn.Linear(8, 8, bias=False, dtype=torch.bfloat16) - local = DTensor.from_local(model.weight.data, mesh_1rank, [Shard(0)], run_check=False) - model.register_parameter("weight", nn.Parameter(local)) - assert model.weight._local_tensor.is_meta # the state under test - - with pytest.raises(RuntimeError, match="needs the shards on cuda"): - migrate_to_buffers(model) - - -@requires_cuda -def test_migrate_leaves_a_model_with_no_gatherable_shards_alone(mesh_1rank): - """A plain (unsharded) model must not open an empty window.""" - from magi_compiler.symm_mem import migrate_to_buffers, registered_buffers - - model = nn.Linear(8, 8, bias=False, dtype=torch.bfloat16).to("cuda") - assert migrate_to_buffers(model) == {} - assert registered_buffers() == [] diff --git a/tests/feature_tests/test_uneven_shard_transport.py b/tests/feature_tests/test_uneven_shard_transport.py new file mode 100644 index 0000000..5758b07 --- /dev/null +++ b/tests/feature_tests/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={'symm_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={'nccl_coalesced': 1, 'symm_coalesced': 1} " + "sizes=[('nccl_coalesced', 2), ('symm_coalesced', 2)]" in p.stdout + ), out[-4000:] + assert "UNEVEN_PASS" in p.stdout, out[-4000:] From 1070f5e6989182ea682a31f04123c62252a602e2 Mon Sep 17 00:00:00 2001 From: wtr Date: Mon, 7 Sep 2026 21:33:55 +0800 Subject: [PATCH 12/16] chore Co-authored-by: Cursor --- .../passes/fsdp_overlap/bucket_all_gather.py | 13 +------------ .../{ => fsdp}/symm_helper/__init__.py | 0 .../{ => fsdp}/symm_helper/verify_symm_e2e.py | 0 tests/feature_tests/{ => fsdp}/test_copy_engine.py | 0 .../{ => fsdp}/test_symm_all_gather.py | 0 tests/feature_tests/{ => fsdp}/test_symm_bind.py | 0 tests/feature_tests/{ => fsdp}/test_symm_e2e.py | 0 .../{ => fsdp}/test_uneven_shard_transport.py | 0 8 files changed, 1 insertion(+), 12 deletions(-) rename tests/feature_tests/{ => fsdp}/symm_helper/__init__.py (100%) rename tests/feature_tests/{ => fsdp}/symm_helper/verify_symm_e2e.py (100%) rename tests/feature_tests/{ => fsdp}/test_copy_engine.py (100%) rename tests/feature_tests/{ => fsdp}/test_symm_all_gather.py (100%) rename tests/feature_tests/{ => fsdp}/test_symm_bind.py (100%) rename tests/feature_tests/{ => fsdp}/test_symm_e2e.py (100%) rename tests/feature_tests/{ => fsdp}/test_uneven_shard_transport.py (100%) diff --git a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py index c4e35c8..35ea330 100644 --- a/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py +++ b/magi_compiler/passes/fsdp_overlap/bucket_all_gather.py @@ -72,18 +72,7 @@ def _is_weight_all_gather(node: fx.Node) -> bool: def _local_shard_bytes(ag: fx.Node) -> int: - """Bytes one rank contributes to a weight all_gather -- the gathered size divided - by the world, NOT the local shard's own meta. - - Used to cap coalesced bucket size, so it has to be rank-identical: ranks that cut - a bucket at different members submit coalesced launches with different membership, - which never completes. Reading the gather's INPUT meta happens to be safe for a - graph this repo lowered -- the trailing ranks of an uneven ``Shard(0)`` own fewer - rows, but the pad in front of the gather brings them back to ``chunk`` -- and that - is far too subtle a thing to rest a collective on. A gather matched by - ``_gathers_a_weight`` rather than emitted by the lowering has no such pad. The - gather's own ``example_value`` is ``(world * chunk, ...)`` on every rank - unconditionally.""" + """Bytes one rank contributes to a weight all_gather""" m = ag.meta.get("example_value") if m is None: return 0 diff --git a/tests/feature_tests/symm_helper/__init__.py b/tests/feature_tests/fsdp/symm_helper/__init__.py similarity index 100% rename from tests/feature_tests/symm_helper/__init__.py rename to tests/feature_tests/fsdp/symm_helper/__init__.py diff --git a/tests/feature_tests/symm_helper/verify_symm_e2e.py b/tests/feature_tests/fsdp/symm_helper/verify_symm_e2e.py similarity index 100% rename from tests/feature_tests/symm_helper/verify_symm_e2e.py rename to tests/feature_tests/fsdp/symm_helper/verify_symm_e2e.py diff --git a/tests/feature_tests/test_copy_engine.py b/tests/feature_tests/fsdp/test_copy_engine.py similarity index 100% rename from tests/feature_tests/test_copy_engine.py rename to tests/feature_tests/fsdp/test_copy_engine.py diff --git a/tests/feature_tests/test_symm_all_gather.py b/tests/feature_tests/fsdp/test_symm_all_gather.py similarity index 100% rename from tests/feature_tests/test_symm_all_gather.py rename to tests/feature_tests/fsdp/test_symm_all_gather.py diff --git a/tests/feature_tests/test_symm_bind.py b/tests/feature_tests/fsdp/test_symm_bind.py similarity index 100% rename from tests/feature_tests/test_symm_bind.py rename to tests/feature_tests/fsdp/test_symm_bind.py diff --git a/tests/feature_tests/test_symm_e2e.py b/tests/feature_tests/fsdp/test_symm_e2e.py similarity index 100% rename from tests/feature_tests/test_symm_e2e.py rename to tests/feature_tests/fsdp/test_symm_e2e.py diff --git a/tests/feature_tests/test_uneven_shard_transport.py b/tests/feature_tests/fsdp/test_uneven_shard_transport.py similarity index 100% rename from tests/feature_tests/test_uneven_shard_transport.py rename to tests/feature_tests/fsdp/test_uneven_shard_transport.py From e97758c42aa7a4787362988bbcea5fc84b6fb0a9 Mon Sep 17 00:00:00 2001 From: wtr Date: Mon, 7 Sep 2026 21:58:24 +0800 Subject: [PATCH 13/16] Fix CI --- tests/feature_tests/fsdp/test_copy_engine.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/feature_tests/fsdp/test_copy_engine.py b/tests/feature_tests/fsdp/test_copy_engine.py index 07d23bc..0b874eb 100644 --- a/tests/feature_tests/fsdp/test_copy_engine.py +++ b/tests/feature_tests/fsdp/test_copy_engine.py @@ -270,10 +270,10 @@ def test_rewriting_nothing_leaves_the_graph_untouched(mesh_1rank): 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 test_fsdp_overlap_lowering import _build_redistribute_graph - 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) From 2832fcc66b995605bd9ad70477074ffada6163e7 Mon Sep 17 00:00:00 2001 From: wtr Date: Mon, 7 Sep 2026 22:32:09 +0800 Subject: [PATCH 14/16] [Refactor] Name the copy-engine gather op after its transport, not its memory ``symm_all_gather`` named the precondition -- the input must live in a symmetric window -- rather than the mechanism, and symmetric memory does not imply the copy engine: PyTorch's own symm_mem all-gathers run on SM kernels. Not doing that is this op's entire reason to exist, and it is what the reorder pass and the cost model key off, so the name should say it. Rename the op to ``magi::ce_all_gather`` / ``_coalesced``, matching the vocabulary already in use at ``transport="copy_engine"`` and ``node_meta.CE_BOUND``. The buffer, binding and registry keep the ``symm`` name: those really are about the memory model, and the split now marks the boundary instead of blurring it. The uneven-transport helper compared sorted results against a literal whose order encoded the old label's alphabetical position, so its expectation moves with the rename. Co-authored-by: Cursor --- .../passes/fsdp_overlap/copy_engine.py | 4 +- magi_compiler/passes/fsdp_overlap/reorder.py | 30 ++--- magi_compiler/profiling/runtime_estimator.py | 62 +++++----- magi_compiler/symm_mem/__init__.py | 2 +- magi_compiler/symm_mem/all_gather.py | 26 ++--- .../fsdp_overlap_helper/reorder_helper.py | 8 +- .../uneven_shard_helper.py | 8 +- .../fsdp/symm_helper/verify_symm_e2e.py | 6 +- ...mm_all_gather.py => test_ce_all_gather.py} | 32 +++--- tests/feature_tests/fsdp/test_copy_engine.py | 16 +-- .../fsdp/test_profiling_estimator.py | 106 +++++++++--------- tests/feature_tests/fsdp/test_symm_bind.py | 12 +- .../fsdp/test_uneven_shard_transport.py | 6 +- 13 files changed, 158 insertions(+), 160 deletions(-) rename tests/feature_tests/fsdp/{test_symm_all_gather.py => test_ce_all_gather.py} (90%) diff --git a/magi_compiler/passes/fsdp_overlap/copy_engine.py b/magi_compiler/passes/fsdp_overlap/copy_engine.py index a5a446e..7be2467 100644 --- a/magi_compiler/passes/fsdp_overlap/copy_engine.py +++ b/magi_compiler/passes/fsdp_overlap/copy_engine.py @@ -84,7 +84,7 @@ def bind_weights_for_copy_engine(graph: fx.GraphModule, example_inputs: Sequence 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 SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED rewritten = 0 skipped = 0 @@ -96,7 +96,7 @@ def rewrite_weight_ag_to_copy_engine(graph: fx.GraphModule) -> int: if not is_ce_bound(node): skipped += 1 continue - node.target = SYMM_ALL_GATHER if node.target is _ALL_GATHER else SYMM_ALL_GATHER_COALESCED + node.target = CE_ALL_GATHER if node.target is _ALL_GATHER else CE_ALL_GATHER_COALESCED rewritten += 1 if rewritten: diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index c7521f1..73ecc28 100644 --- a/magi_compiler/passes/fsdp_overlap/reorder.py +++ b/magi_compiler/passes/fsdp_overlap/reorder.py @@ -53,19 +53,19 @@ _AG_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default -def _symm_ag_ops(): +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 SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED - return tuple(op for op in (SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED) if op is not None) + return tuple(op for op in (CE_ALL_GATHER, CE_ALL_GATHER_COALESCED) if op is not None) except Exception: # noqa: BLE001 return () -_SYMM_AG_OPS = _symm_ag_ops() -_WEIGHT_AG_OPS = tuple(op for op in (_AG, _AG_COALESCED, *_SYMM_AG_OPS) if op is not None) +_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 @@ -73,24 +73,24 @@ def _symm_ag_ops(): _DEFAULT_WINDOW_MARGIN_NS = 5_000.0 -def _is_symm_ag_ir(node) -> bool: +def _is_ce_ag_ir(node) -> bool: """ - ``magi::symm_all_gather`` lowers to an ordinary FallbackKernel, so + ``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 _SYMM_AG_OPS + return getattr(node, "op_overload", None) in _CE_AG_OPS -def _is_symm_ag_coalesced(node) -> bool: +def _is_ce_ag_coalesced(node) -> bool: try: - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED except Exception: # noqa: BLE001 return False - return SYMM_ALL_GATHER_COALESCED is not None and getattr(node, "op_overload", None) is SYMM_ALL_GATHER_COALESCED + 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_symm_ag_ir(node)) + return node is not None and (is_collective(node) or _is_ce_ag_ir(node)) def _leaf_collective_node(snode: BaseSchedulerNode): @@ -542,18 +542,18 @@ def _launch_group(self, launch, order, buf_to_snode, users) -> list[BaseSchedule group = [launch] node = _leaf_collective_node(launch) produced = set(launch.get_buffer_names()) - if node is not None and (getattr(node, "op_overload", None) is _AG_COALESCED or _is_symm_ag_coalesced(node)): + 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_symm_ag_coalesced(node): + 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_symm_ag_ir(node): + 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 diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index c3fa09f..b82aed1 100644 --- a/magi_compiler/profiling/runtime_estimator.py +++ b/magi_compiler/profiling/runtime_estimator.py @@ -343,23 +343,23 @@ def _collective_spec(node): return op, group_name, group_size, specs -def _symm_ag_ops(): +def _ce_ag_ops(): """Copy-engine gather ops, or empty when the runtime is unavailable.""" try: - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED - return tuple(op for op in (SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED) if op is not None) + 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_symm_ag(snode: BaseSchedulerNode): +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 = _symm_ag_ops() + 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 [])): @@ -368,15 +368,15 @@ def _leaf_symm_ag(snode: BaseSchedulerNode): return None -def _is_symm_ag_coalesced_ir(node) -> bool: +def _is_ce_ag_coalesced_ir(node) -> bool: try: - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED except Exception: # noqa: BLE001 return False - return SYMM_ALL_GATHER_COALESCED is not None and getattr(node, "op_overload", None) is SYMM_ALL_GATHER_COALESCED + return CE_ALL_GATHER_COALESCED is not None and getattr(node, "op_overload", None) is CE_ALL_GATHER_COALESCED -def _symm_ag_spec(node): +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: @@ -389,7 +389,7 @@ def _symm_ag_spec(node): return shapes, ins[0].layout.dtype, int(group_size), str(group_name) -def _symm_ag_launch_wait(snode: BaseSchedulerNode): +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 @@ -397,8 +397,8 @@ def _symm_ag_launch_wait(snode: BaseSchedulerNode): """ from magi_compiler.symm_mem import find_shard_by_layout - node = _leaf_symm_ag(snode) - spec = _symm_ag_spec(node) if node is not None else None + 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 @@ -411,33 +411,33 @@ def _symm_ag_launch_wait(snode: BaseSchedulerNode): dtype, ) return None - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED - if _is_symm_ag_coalesced_ir(node): - op = SYMM_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: SYMM_ALL_GATHER(shards[0], group_size, group_name)), _WAIT + return (lambda: CE_ALL_GATHER(shards[0], group_size, group_name)), _WAIT -def _measure_symm_ag(snode: BaseSchedulerNode) -> float: +def _measure_ce_ag(snode: BaseSchedulerNode) -> float: """Time ``wait(launch())``. Launch-only is ~3us CPU issue; copies run on a side stream.""" - pair = _symm_ag_launch_wait(snode) + pair = _ce_ag_launch_wait(snode) if pair is None: return 0.0 launch, wait = pair return _time_fixed(lambda: wait(launch())) -def _symm_ag_label(snode: BaseSchedulerNode) -> str: - node = _leaf_symm_ag(snode) - spec = _symm_ag_spec(node) if node is not None else None +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"symm_all_gather_coalesced(ws={group_size},n={len(shapes)},{shape0})" - return f"symm_all_gather(ws={group_size},{shape0})" + 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: @@ -605,8 +605,8 @@ def _measure_one(self, snode: BaseSchedulerNode) -> float: """Lockstep-safe single measurement (fixed iters for anything containing a collective); never raises -- falls back to the analytical estimate.""" try: - if _leaf_symm_ag(snode) is not None: - return _measure_symm_ag(snode) + if _leaf_ce_ag(snode) is not None: + return _measure_ce_ag(snode) if contains_collective(snode): return _measure_collective_op(snode) if isinstance(snode, ExternKernelSchedulerNode): @@ -647,24 +647,24 @@ def __call__(self, snode: BaseSchedulerNode) -> float: if _is_multi_output_unpack(snode): return 0.0 - if _leaf_symm_ag(snode) is not None: - node = _leaf_symm_ag(snode) - spec = _symm_ag_spec(node) + 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 = ("symm_ag", group_size, shapes, str(dtype)) + 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="symm_ag", label=_symm_ag_label(snode), measured=False) + 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_symm_ag(snode) + ns = _measure_ce_ag(snode) self._table[ckey].ns = ns self._table[ckey].measured = True self.n_measured += 1 diff --git a/magi_compiler/symm_mem/__init__.py b/magi_compiler/symm_mem/__init__.py index 78ae3ec..ace63f9 100644 --- a/magi_compiler/symm_mem/__init__.py +++ b/magi_compiler/symm_mem/__init__.py @@ -15,7 +15,7 @@ """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::symm_all_gather`` ops, and callers use the import itself as the probe for +``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. """ diff --git a/magi_compiler/symm_mem/all_gather.py b/magi_compiler/symm_mem/all_gather.py index 405ef72..4fec639 100644 --- a/magi_compiler/symm_mem/all_gather.py +++ b/magi_compiler/symm_mem/all_gather.py @@ -28,8 +28,8 @@ # 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 = "symm_all_gather(Tensor local, int group_size, str group_name) -> Tensor" -_SCHEMA_COALESCED = "symm_all_gather_coalesced(Tensor[] shards, int group_size, str group_name) -> Tensor[]" +_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, ...]] @@ -133,7 +133,7 @@ def _shard_peers(local: torch.Tensor, group_size: int) -> tuple[torch.Tensor, .. entry = lookup_shard(local.data_ptr()) if entry is None: raise RuntimeError( - "magi::symm_all_gather got a tensor that is not a registered symmetric-memory shard. " + "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." ) @@ -185,7 +185,7 @@ def _gather_dest(local: torch.Tensor, group_size: int) -> torch.Tensor: return local.new_empty((local.shape[0] * group_size, *local.shape[1:])) -def _symm_all_gather(local: torch.Tensor, group_size: int, group_name: str) -> torch.Tensor: +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)]) @@ -193,11 +193,11 @@ def _symm_all_gather(local: torch.Tensor, group_size: int, group_name: str) -> t return out -def _symm_all_gather_meta(local: torch.Tensor, group_size: int, group_name: str) -> torch.Tensor: +def _ce_all_gather_meta(local: torch.Tensor, group_size: int, group_name: str) -> torch.Tensor: return _gather_dest(local, group_size) -def _symm_all_gather_coalesced(shards: list[torch.Tensor], group_size: int, group_name: str) -> list[torch.Tensor]: +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: @@ -211,23 +211,23 @@ def _symm_all_gather_coalesced(shards: list[torch.Tensor], group_size: int, grou return outs -def _symm_all_gather_coalesced_meta(shards: list[torch.Tensor], group_size: int, group_name: str) -> list[torch.Tensor]: +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("symm_all_gather", _symm_all_gather, "CUDA") - _LIB.impl("symm_all_gather", _symm_all_gather_meta, "Meta") + _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("symm_all_gather_coalesced", _symm_all_gather_coalesced, "CUDA") - _LIB.impl("symm_all_gather_coalesced", _symm_all_gather_coalesced_meta, "Meta") + _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. -SYMM_ALL_GATHER = torch.ops.magi.symm_all_gather.default -SYMM_ALL_GATHER_COALESCED = torch.ops.magi.symm_all_gather_coalesced.default +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/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py index 187d04f..ac50f12 100644 --- a/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py +++ b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py @@ -39,7 +39,7 @@ the expected mode for each rung of the ladder (identical / slot / pinned / abort). With ``--copy-engine``: the same shape, but the gathers are -``magi::symm_all_gather`` reading a symmetric buffer. Two things are checked that +``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 @@ -170,7 +170,7 @@ def fn(x, w0, shard): ce_shards: list = [] if args.copy_engine: from magi_compiler.symm_mem import alloc_shard, publish - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER + 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) @@ -184,7 +184,7 @@ def fn(x, w0, shards): # noqa: F811 - deliberately replaces the NCCL variant y = _WAIT(_AR(y, "sum", grp)) acc = None for i, sh in enumerate(shards): - g = _WAIT(SYMM_ALL_GATHER(sh, world, grp)) + 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 @@ -239,7 +239,7 @@ def greedy_cost(snode) -> float: fallback kernel at 0us, so without this the launches barely move.""" from torch._inductor.comms import estimate_op_runtime - if _ro._is_symm_ag_ir(_ro._leaf_collective_node(snode)): + 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) 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 index 91bcc8f..13b8a23 100644 --- a/tests/feature_tests/fsdp/fsdp_overlap_helper/uneven_shard_helper.py +++ b/tests/feature_tests/fsdp/fsdp_overlap_helper/uneven_shard_helper.py @@ -119,13 +119,13 @@ def _example_inputs(gm) -> list[object]: 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 SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED names = { _NCCL_AG: "nccl", _NCCL_AG_COALESCED: "nccl_coalesced", - SYMM_ALL_GATHER: "symm", - SYMM_ALL_GATHER_COALESCED: "symm_coalesced", + CE_ALL_GATHER: "ce", + CE_ALL_GATHER_COALESCED: "ce_coalesced", } counts: Counter[str] = Counter() sizes: list[tuple[str, int]] = [] @@ -225,7 +225,7 @@ def _check_mixed(mesh, *, even_rows: int, uneven_rows: int, cols: int, say) -> b targets, sizes = _gather_targets(gm) targets = dict(sorted(targets.items())) sizes = sorted(sizes) - expected = ({"nccl_coalesced": 1, "symm_coalesced": 1}, [("nccl_coalesced", 2), ("symm_coalesced", 2)]) + 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 diff --git a/tests/feature_tests/fsdp/symm_helper/verify_symm_e2e.py b/tests/feature_tests/fsdp/symm_helper/verify_symm_e2e.py index bfc0709..90f448d 100644 --- a/tests/feature_tests/fsdp/symm_helper/verify_symm_e2e.py +++ b/tests/feature_tests/fsdp/symm_helper/verify_symm_e2e.py @@ -29,7 +29,7 @@ 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::symm_all_gather``. A pass + 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 @@ -231,9 +231,7 @@ def spy_rewrite(graph): 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::symm_all_gather -> {'ok' if rewrite_ok else 'WRONG'}" - ) + 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() diff --git a/tests/feature_tests/fsdp/test_symm_all_gather.py b/tests/feature_tests/fsdp/test_ce_all_gather.py similarity index 90% rename from tests/feature_tests/fsdp/test_symm_all_gather.py rename to tests/feature_tests/fsdp/test_ce_all_gather.py index 464cb4d..240e2a3 100644 --- a/tests/feature_tests/fsdp/test_symm_all_gather.py +++ b/tests/feature_tests/fsdp/test_ce_all_gather.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Mechanics of ``magi::symm_all_gather`` (step 2 of the landing order). +"""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 @@ -23,7 +23,7 @@ The transport itself (peer reads, overlap, ordering under load) needs several NVLink-connected ranks and lives in -``example/inference/fsdp_overlap/verify_symm_ag_op.py``. +``example/inference/fsdp_overlap/verify_ce_ag_op.py``. """ from __future__ import annotations @@ -61,7 +61,7 @@ def mesh_1rank(): @pytest.fixture(autouse=True) def _clean_state(): - # Importing the module is what defines ``magi::symm_all_gather``. Nothing + # 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 @@ -108,7 +108,7 @@ def test_gather_matches_nccl_bitwise(mesh_1rank): _, shards = _symm_model(mesh_1rank) for shard in shards: - got = _WAIT(torch.ops.magi.symm_all_gather(shard, 1, "")) + 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() @@ -119,11 +119,11 @@ def test_gather_matches_nccl_bitwise(mesh_1rank): 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.symm_all_gather_coalesced(list(shards), 1, "") + 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.symm_all_gather(shard, 1, "")) + ref = _WAIT(torch.ops.magi.ce_all_gather(shard, 1, "")) assert torch.equal(got, ref) @@ -133,7 +133,7 @@ def test_wait_tensor_picks_up_the_registered_event(mesh_1rank): _, shards = _symm_model(mesh_1rank) shard = shards[0] - out = _WAIT(torch.ops.magi.symm_all_gather(shard, 1, "")) + 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" @@ -143,8 +143,8 @@ 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.symm_all_gather(shards[0], 1, "") - b = torch.ops.magi.symm_all_gather(shards[1], 1, "") + 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() @@ -159,14 +159,14 @@ def test_unregistered_shard_is_rejected(mesh_1rank): 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.symm_all_gather(ordinary, 1, "") + 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.symm_all_gather(shards[0], 4, "") + torch.ops.magi.ce_all_gather(shards[0], 4, "") @requires_cuda @@ -177,14 +177,14 @@ def test_coalesced_validates_every_member_before_allocating(mesh_1rank): _, 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.symm_all_gather_coalesced([shards[0], ordinary], 1, "") + 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.symm_all_gather(local, 8, "") + out = torch.ops.magi.ce_all_gather(local, 8, "") assert out.shape == (32, 6) and out.is_meta @@ -194,7 +194,7 @@ def test_coalesced_meta_kernel_shapes(mesh_1rank): 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.symm_all_gather_coalesced(shards, 8, "") + 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) @@ -258,11 +258,11 @@ def test_per_peer_fallback_matches_the_batched_path(mesh_1rank, monkeypatch): from magi_compiler.symm_mem import all_gather as ag_mod _, shards = _symm_model(mesh_1rank) - batched = _WAIT(torch.ops.magi.symm_all_gather(shards[0], 1, "")) + 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.symm_all_gather(shards[0], 1, "")) + 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 index 0b874eb..eb886da 100644 --- a/tests/feature_tests/fsdp/test_copy_engine.py +++ b/tests/feature_tests/fsdp/test_copy_engine.py @@ -65,9 +65,9 @@ def mesh_1rank(): def _symm_op(): - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER - return SYMM_ALL_GATHER + return CE_ALL_GATHER def _graph_with_gathers( @@ -292,7 +292,7 @@ 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 SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED gm = _graph_with_gathers(mesh_1rank, 4) _pretend_bound(gm) @@ -300,7 +300,7 @@ def test_a_bucket_of_bound_gathers_is_retargeted(mesh_1rank): assert rewrite_weight_ag_to_copy_engine(gm) == 1 assert not _gathers(gm, _AG_COALESCED) - assert len(_gathers(gm, SYMM_ALL_GATHER_COALESCED)) == 1 + assert len(_gathers(gm, CE_ALL_GATHER_COALESCED)) == 1 assert len(_gathers(gm, _WAIT)) == 4 @@ -311,7 +311,7 @@ def test_a_bucket_is_all_or_nothing(mesh_1rank): 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 SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED gm = _graph_with_gathers(mesh_1rank, 3) _pretend_bound(gm) @@ -322,7 +322,7 @@ def test_a_bucket_is_all_or_nothing(mesh_1rank): 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, SYMM_ALL_GATHER_COALESCED) + assert not _gathers(gm, CE_ALL_GATHER_COALESCED) def _coalesced_graph(mesh, n: int, *, marked: bool = True, bound: bool = True): @@ -363,12 +363,12 @@ def _coalesced_graph(mesh, n: int, *, marked: bool = True, bound: bool = True): @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 SYMM_ALL_GATHER_COALESCED + 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, SYMM_ALL_GATHER_COALESCED)) == 1 + assert len(_gathers(gm, CE_ALL_GATHER_COALESCED)) == 1 @requires_cuda diff --git a/tests/feature_tests/fsdp/test_profiling_estimator.py b/tests/feature_tests/fsdp/test_profiling_estimator.py index 9cdc809..eba2321 100644 --- a/tests/feature_tests/fsdp/test_profiling_estimator.py +++ b/tests/feature_tests/fsdp/test_profiling_estimator.py @@ -30,15 +30,15 @@ from magi_compiler.profiling import runtime_estimator as re_mod from magi_compiler.profiling.runtime_estimator import ( ProfileEntry, - _is_symm_ag_coalesced_ir, - _leaf_symm_ag, + _ce_ag_label, + _ce_ag_launch_wait, + _ce_ag_spec, + _is_ce_ag_coalesced_ir, + _leaf_ce_ag, + _measure_ce_ag, _measure_extern, - _measure_symm_ag, _realize_arg, _static, - _symm_ag_label, - _symm_ag_launch_wait, - _symm_ag_spec, ) from magi_compiler.utils.envs import TORCH_VERSION @@ -562,9 +562,9 @@ def test_collective_profile_accuracy_multi_rank(): def _ce_ops(): - from magi_compiler.symm_mem.all_gather import SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED - return SYMM_ALL_GATHER, SYMM_ALL_GATHER_COALESCED + return CE_ALL_GATHER, CE_ALL_GATHER_COALESCED class _FakeLayout: @@ -578,8 +578,8 @@ def __init__(self, shape, dtype): self.layout = _FakeLayout(tuple(shape), dtype) -class _FakeSymmAgIR: - """A ``magi::symm_all_gather`` FallbackKernel. +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``. @@ -609,79 +609,79 @@ def __init__(self, node=None, snodes=()): def _ce_snode(shape=(_CE_ROWS, 64), group_size=1): ag, _ = _ce_ops() - return _FakeGatherSnode(node=_FakeSymmAgIR(ag, [shape], group_size=group_size)) + return _FakeGatherSnode(node=_FakeCeAgIR(ag, [shape], group_size=group_size)) # --- spec / recognition / label: structural, no window needed --------------- -def test_symm_ag_spec_is_read_off_constant_args_not_the_origin_node(): +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 = _FakeSymmAgIR(ag, [(_CE_ROWS, 64)], group_size=8, group_name="fsdp") + 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 = _symm_ag_spec(node) + 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_symm_ag_spec_of_a_coalesced_gather_has_one_shape_per_member(): +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 = _symm_ag_spec(_FakeSymmAgIR(coalesced, [(8, 4), (16, 4), (32, 4)], group_size=4)) + 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_symm_ag_spec_is_none_when_the_group_args_are_missing(constant_args): +def test_ce_ag_spec_is_none_when_the_group_args_are_missing(constant_args): ag, _ = _ce_ops() - node = _FakeSymmAgIR(ag, [(8, 4)]) + node = _FakeCeAgIR(ag, [(8, 4)]) node.constant_args = constant_args - assert _symm_ag_spec(node) is None + assert _ce_ag_spec(node) is None -def test_symm_ag_spec_is_none_when_the_node_has_no_inputs(): +def test_ce_ag_spec_is_none_when_the_node_has_no_inputs(): ag, _ = _ce_ops() - assert _symm_ag_spec(_FakeSymmAgIR(ag, [])) is None + assert _ce_ag_spec(_FakeCeAgIR(ag, [])) is None -def test_leaf_symm_ag_finds_the_gather_in_a_plain_and_in_a_fused_snode(): +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 = _FakeSymmAgIR(ag, [(8, 4)]) + node = _FakeCeAgIR(ag, [(8, 4)]) - assert _leaf_symm_ag(_FakeGatherSnode(node=node)) is node + assert _leaf_ce_ag(_FakeGatherSnode(node=node)) is node fused = _FakeGatherSnode(node=None, snodes=[_FakeGatherSnode(node=_FakeMmIR()), _FakeGatherSnode(node=node)]) - assert _leaf_symm_ag(fused) is node + assert _leaf_ce_ag(fused) is node -def test_leaf_symm_ag_is_none_for_an_ordinary_kernel(): +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_symm_ag(_FakeGatherSnode(node=_FakeMmIR())) is None - assert _leaf_symm_ag(_FakeGatherSnode()) is None + 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_symm_ag_coalesced_ir(_FakeSymmAgIR(coalesced, [(8, 4)])) - assert not _is_symm_ag_coalesced_ir(_FakeSymmAgIR(ag, [(8, 4)])) + assert _is_ce_ag_coalesced_ir(_FakeCeAgIR(coalesced, [(8, 4)])) + assert not _is_ce_ag_coalesced_ir(_FakeCeAgIR(ag, [(8, 4)])) -def test_symm_ag_label_reports_transport_world_size_and_member_count(): +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 = _symm_ag_label(_FakeGatherSnode(node=_FakeSymmAgIR(ag, [(8, 4)], group_size=2))) - batch = _symm_ag_label(_FakeGatherSnode(node=_FakeSymmAgIR(coalesced, [(8, 4), (16, 4)], group_size=2))) + 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 == "symm_all_gather(ws=2,8x4)" - assert batch == "symm_all_gather_coalesced(ws=2,n=2,8x4)" + 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 ---------------------------------- @@ -719,17 +719,17 @@ def _register_shards(shapes, dtype=torch.bfloat16): def _ce_replay_snode(shapes, coalesced=False): ag, ag_coalesced = _ce_ops() op = ag_coalesced if coalesced else ag - return _FakeGatherSnode(node=_FakeSymmAgIR(op, shapes, group_size=1, group_name=_ce_group_name())) + return _FakeGatherSnode(node=_FakeCeAgIR(op, shapes, group_size=1, group_name=_ce_group_name())) @requires_cuda -def test_symm_ag_replay_gathers_the_registered_shard(pg_1rank, symm_registry): +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 = _symm_ag_launch_wait(_ce_replay_snode([(_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. @@ -737,11 +737,11 @@ def test_symm_ag_replay_gathers_the_registered_shard(pg_1rank, symm_registry): @requires_cuda -def test_symm_ag_coalesced_replay_covers_every_member(pg_1rank, symm_registry): +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 = _symm_ag_launch_wait(_ce_replay_snode(shapes, coalesced=True)) + 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): @@ -749,25 +749,25 @@ def test_symm_ag_coalesced_replay_covers_every_member(pg_1rank, symm_registry): @requires_cuda -def test_measure_symm_ag_prices_the_gather_above_zero(pg_1rank, symm_registry): +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_symm_ag(_ce_replay_snode([(1024, 1024)])) > 0.0 + assert _measure_ce_ag(_ce_replay_snode([(1024, 1024)])) > 0.0 @requires_cuda -def test_measure_symm_ag_degrades_quietly_when_no_shard_has_that_layout(pg_1rank, symm_registry): +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 _symm_ag_launch_wait(snode) is None - assert _measure_symm_ag(snode) == 0.0 + assert _ce_ag_launch_wait(snode) is None + assert _measure_ce_ag(snode) == 0.0 -# --- __call__: the ("symm_ag", ...) cache ----------------------------------- +# --- __call__: the ("ce_ag", ...) cache ----------------------------------- @pytest.fixture def ce_measure_calls(monkeypatch): """Drive ``__call__`` straight into the copy-engine branch and count measures.""" @@ -780,11 +780,11 @@ def fake_measure(snode): 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_symm_ag", fake_measure) + monkeypatch.setattr(re_mod, "_measure_ce_ag", fake_measure) return calls -def test_isomorphic_symm_ag_gathers_are_measured_once(ce_measure_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() @@ -798,7 +798,7 @@ def test_isomorphic_symm_ag_gathers_are_measured_once(ce_measure_calls): assert len(est.table) == 1 -def test_symm_ag_cache_key_separates_shape_and_world_size(ce_measure_calls): +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() @@ -811,18 +811,18 @@ def test_symm_ag_cache_key_separates_shape_and_world_size(ce_measure_calls): assert est.n_cache_hits == 0 -def test_symm_ag_entry_is_tagged_as_a_copy_engine_gather(ce_measure_calls): +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 == "symm_ag" + assert entry.kind == "ce_ag" assert entry.measured assert entry.ns == 222.0 - assert "symm_all_gather" in entry.label + assert "ce_all_gather" in entry.label -def test_symm_ag_sync_mode_defers_the_measurement_to_warm_and_sync(ce_measure_calls): +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.""" diff --git a/tests/feature_tests/fsdp/test_symm_bind.py b/tests/feature_tests/fsdp/test_symm_bind.py index 4c02a06..27d0454 100644 --- a/tests/feature_tests/fsdp/test_symm_bind.py +++ b/tests/feature_tests/fsdp/test_symm_bind.py @@ -486,13 +486,13 @@ 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 SYMM_ALL_GATHER_COALESCED + 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 SYMM_ALL_GATHER_COALESCED] + 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 @@ -504,7 +504,7 @@ def test_an_unbound_weight_keeps_its_neighbours_on_the_copy_engine(mesh_1rank): 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 SYMM_ALL_GATHER_COALESCED + 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] @@ -512,7 +512,7 @@ def test_an_unbound_weight_keeps_its_neighbours_on_the_copy_engine(mesh_1rank): 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 SYMM_ALL_GATHER_COALESCED] + 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 @@ -529,7 +529,7 @@ def test_unbound_weights_are_still_bucketed_as_nccl(mesh_1rank): """ 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 SYMM_ALL_GATHER_COALESCED + from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER_COALESCED _AG_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default @@ -542,7 +542,7 @@ def test_unbound_weights_are_still_bucketed_as_nccl(mesh_1rank): == 2 ) - (ce,) = [n for n in gm.graph.nodes if n.target is SYMM_ALL_GATHER_COALESCED] + (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 diff --git a/tests/feature_tests/fsdp/test_uneven_shard_transport.py b/tests/feature_tests/fsdp/test_uneven_shard_transport.py index 5758b07..e1d6275 100644 --- a/tests/feature_tests/fsdp/test_uneven_shard_transport.py +++ b/tests/feature_tests/fsdp/test_uneven_shard_transport.py @@ -68,7 +68,7 @@ def test_uneven_shard_transport_is_rank_identical(): 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={'symm_coalesced': 1}" in p.stdout, out[-4000:] + 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. @@ -76,7 +76,7 @@ def test_uneven_shard_transport_is_rank_identical(): 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={'nccl_coalesced': 1, 'symm_coalesced': 1} " - "sizes=[('nccl_coalesced', 2), ('symm_coalesced', 2)]" in p.stdout + "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:] From 5abd8c258f5cf0ff2592fb30b004940f34420cea Mon Sep 17 00:00:00 2001 From: wtr Date: Wed, 9 Sep 2026 13:36:19 +0800 Subject: [PATCH 15/16] [Fix] Carry a gather's unspent compute forward and stop pricing unreplayable kernels at zero --- magi_compiler/passes/fsdp_overlap/reorder.py | 26 +++-- magi_compiler/profiling/runtime_estimator.py | 117 ++++++++++++++----- 2 files changed, 107 insertions(+), 36 deletions(-) diff --git a/magi_compiler/passes/fsdp_overlap/reorder.py b/magi_compiler/passes/fsdp_overlap/reorder.py index 73ecc28..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). @@ -318,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": @@ -345,13 +349,16 @@ 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 %s(%s) cur=%d -> target=%d fc=%d lower=%d | " - "comm=%.1fus acc_upstream=%.1fus need=%.1fus %s", + "comm=%.1fus carry_in=%.1fus acc_upstream=%.1fus need=%.1fus %s", launch.get_name(), getattr(_leaf_collective_node(launch), "op_overload", "?"), cur, @@ -359,6 +366,7 @@ def __call__(self, snodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]: fc_idx, lower, comm_runtime / 1e3, + carry_in / 1e3, acc / 1e3, need / 1e3, "hidden" if acc >= need else "COMPUTE-LIMITED", diff --git a/magi_compiler/profiling/runtime_estimator.py b/magi_compiler/profiling/runtime_estimator.py index b82aed1..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) @@ -222,7 +280,7 @@ def _extern_replay_fn(snode: ExternKernelSchedulerNode): 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() + fx_node = _fx_node_of(snode.node) if fx_node is None: return None target = fx_node.target @@ -256,7 +314,9 @@ def _measure_extern(snode: ExternKernelSchedulerNode, fixed_iters: bool = False) -> NCCL count mismatch -> deadlock.""" fn = _extern_replay_fn(snode) if fn is None: - return 0.0 + # 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 @@ -562,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 @@ -601,24 +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) + 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 From 6a2d87540f652b5a232e5a55a1e624c811c47055 Mon Sep 17 00:00:00 2001 From: wtr Date: Wed, 9 Sep 2026 22:00:29 +0800 Subject: [PATCH 16/16] refactor symm_buffer --- magi_compiler/symm_mem/__init__.py | 5 +- magi_compiler/symm_mem/bind.py | 62 ++++++- magi_compiler/symm_mem/symm_buffer.py | 175 +++++++----------- .../fsdp_overlap_helper/reorder_helper.py | 8 +- .../fsdp/test_profiling_estimator.py | 8 +- 5 files changed, 130 insertions(+), 128 deletions(-) diff --git a/magi_compiler/symm_mem/__init__.py b/magi_compiler/symm_mem/__init__.py index ace63f9..3874358 100644 --- a/magi_compiler/symm_mem/__init__.py +++ b/magi_compiler/symm_mem/__init__.py @@ -19,15 +19,13 @@ whether the copy-engine transport is available. Import that module by path. """ -from .bind import bind_graph_weights, bind_parameters +from .bind import bind_graph_weights, bind_parameters, group_name_of from .symm_buffer import ( ShardEntry, SymmBuffer, alloc_shard, find_shard_by_layout, - group_name_of, lookup_shard, - publish, registered_buffers, reset_registry, ) @@ -41,7 +39,6 @@ "find_shard_by_layout", "group_name_of", "lookup_shard", - "publish", "registered_buffers", "reset_registry", ] diff --git a/magi_compiler/symm_mem/bind.py b/magi_compiler/symm_mem/bind.py index 0676682..ab6f357 100644 --- a/magi_compiler/symm_mem/bind.py +++ b/magi_compiler/symm_mem/bind.py @@ -30,7 +30,7 @@ from magi_compiler.utils import magi_logger -from .symm_buffer import group_name_of, lookup_shard, open_buffer, publish, register_shard +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)``.""" @@ -73,7 +73,6 @@ def bind_graph_weights( return set() allocated = _move(plan) - publish() served = {c.gather for c in plan if c.gather is not None} magi_logger.info( @@ -108,7 +107,6 @@ def bind_parameters(params: Iterable[Any], min_shard_bytes: int = 0) -> int: return 0 _move(plan) - publish() return len(plan) @@ -183,6 +181,21 @@ def _unbindable(param: Any, min_shard_bytes: int) -> str | None: 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). @@ -224,8 +237,36 @@ def _windows(plan: list[BindCandidate]) -> list[list[BindCandidate]]: 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 and repoint every shard. Returns the new allocation count. + """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 @@ -235,8 +276,9 @@ def _move(plan: list[BindCandidate]) -> int: 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, (c.local.numel() for c in members)) + 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( @@ -255,15 +297,19 @@ def _move(plan: list[BindCandidate]) -> int: ) raise - for c in members: - symm = buffer.take(c.local.shape) + for c, offset in zip(members, offsets): + symm = buffer.get_tensor(offset, c.local.shape) symm.copy_(c.local) - register_shard(symm, buffer) + 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 diff --git a/magi_compiler/symm_mem/symm_buffer.py b/magi_compiler/symm_mem/symm_buffer.py index f3d362d..9bf8b40 100644 --- a/magi_compiler/symm_mem/symm_buffer.py +++ b/magi_compiler/symm_mem/symm_buffer.py @@ -14,8 +14,9 @@ """Symmetric-memory windows for copy-engine weight all-gather. -Shards are suballocated from pooled windows; runtime gathers look them up by -device pointer. Which weights get an allocation is ``bind``'s job. +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 @@ -23,86 +24,74 @@ from dataclasses import dataclass import torch -import torch.distributed as dist - -from magi_compiler.utils import magi_logger class SymmBuffer: - """One symmetric-memory window, suballocated to many weight shards. + """One symmetric-memory window, addressed by element offset. - The driver caps windows at 128 per process regardless of size, so pooling is - required. Every rank walks the same plan, so offset ``k`` is the same shard - on every peer. + 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. """ - # 256 bf16 elems = 512B, which the copy engine wants for peak throughput. - ALIGN = 256 - - def __init__(self, dtype: torch.dtype, device: torch.device, group_name: str) -> None: + 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.buf: torch.Tensor | None = None - self.handle = None - self._reserved = 0 - self._cursor = 0 + self.numel = numel + self.buf = buf + self.handle = handle - def reserve(self, numel: int) -> None: - """Book space for one shard. Call for every member before ``commit``.""" - self._reserved += self._round(numel) + @classmethod + def open(cls, dtype: torch.dtype, device: torch.device, group_name: str, numel: int) -> SymmBuffer: + """Allocate and rendezvous a window of ``numel`` elements. - def commit(self) -> None: - """Open the window. One ``rendezvous`` for every shard it will hold.""" + 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(self.group_name) - self.buf = symm_mem.empty(self._reserved, dtype=self.dtype, device=self.device) - self.handle = symm_mem.rendezvous(self.buf, self.group_name) + 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 take(self, shape: torch.Size | tuple[int, ...]) -> torch.Tensor: - """Hand out the next slot as a tensor whose storage starts at the slot. + 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``: Dynamo memoized ``storage_offset == 0`` for these - parameters, and a mid-window slice would contradict the shape env. + 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 *= int(s) - off = self._cursor - self._cursor += self._round(numel) - if self._cursor > self._reserved: + numel *= s + if offset < 0 or offset + numel > self.numel: raise RuntimeError( - f"symmetric buffer overflow: wanted {self._cursor} elems, reserved {self._reserved}. " - "The sizing walk and the dispensing walk must visit the same shards in the same order." + 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, tuple(int(s) for s in shape), self.dtype, off) + return self.handle.get_buffer(self.handle.rank if rank is None else rank, shape, self.dtype, offset) - @property - def nbytes(self) -> int: - return self._reserved * self.dtype.itemsize - - def offset_of(self, t: torch.Tensor) -> int: - return (t.data_ptr() - self.buf.data_ptr()) // self.buf.element_size() - - def peer_views(self, t: torch.Tensor) -> list[torch.Tensor]: - """``world_size`` views of the same shard, one per rank. + 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. """ - off, shape = self.offset_of(t), tuple(int(s) for s in t.shape) - return [self.handle.get_buffer(r, shape, self.dtype, off) for r in range(self.handle.world_size)] + 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: - if self.buf is None: - return False base = self.buf.data_ptr() return base <= t.data_ptr() < base + self.nbytes - @classmethod - def _round(cls, numel: int) -> int: - return (numel + cls.ALIGN - 1) // cls.ALIGN * cls.ALIGN - @dataclass(frozen=True) class ShardEntry: @@ -124,56 +113,26 @@ def dtype(self) -> torch.dtype: _SHARD_REGISTRY: dict[int, ShardEntry] = {} _BUFFERS: list[SymmBuffer] = [] -_UNPUBLISHED = False -def open_buffer(dtype: torch.dtype, device: torch.device, group_name: str, numels) -> SymmBuffer: - """Open one window big enough for ``numels``, and track it for ``publish``.""" - global _UNPUBLISHED - buffer = SymmBuffer(dtype, device, group_name) - for numel in numels: - buffer.reserve(int(numel)) - buffer.commit() +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) - _UNPUBLISHED = True return buffer -def register_shard(local: torch.Tensor, buffer: SymmBuffer) -> ShardEntry: +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=buffer.offset_of(local), local=local, peer_views=tuple(buffer.peer_views(local))) + 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 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.take(shape) - register_shard(shard, buffer) - return shard - - -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] - - 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) @@ -192,26 +151,20 @@ def find_shard_by_layout(shape, dtype: torch.dtype) -> torch.Tensor | None: return None -def publish() -> None: - """Barrier so every rank has copied its shards in before any peer reads. - - Once per bind, not per step: the weights never change again. - """ - global _UNPUBLISHED - if not _UNPUBLISHED: - return - if dist.is_available() and dist.is_initialized(): - torch.cuda.synchronize() - dist.barrier() - _UNPUBLISHED = False - magi_logger.info( - "SymmBuffer: published %d shard(s), %.1f MiB total", len(_BUFFERS), sum(b.nbytes for b in _BUFFERS) / 2**20 - ) - - def reset_registry() -> None: """Drop every window and shard. Tests only -- frees the symmetric allocations.""" - global _UNPUBLISHED _SHARD_REGISTRY.clear() _BUFFERS.clear() - _UNPUBLISHED = False + + +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/fsdp/fsdp_overlap_helper/reorder_helper.py b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py index ac50f12..4367926 100644 --- a/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py +++ b/tests/feature_tests/fsdp/fsdp_overlap_helper/reorder_helper.py @@ -169,15 +169,17 @@ def fn(x, w0, shard): ce_shards: list = [] if args.copy_engine: - from magi_compiler.symm_mem import alloc_shard, publish + 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. - publish() + # 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() diff --git a/tests/feature_tests/fsdp/test_profiling_estimator.py b/tests/feature_tests/fsdp/test_profiling_estimator.py index eba2321..2420684 100644 --- a/tests/feature_tests/fsdp/test_profiling_estimator.py +++ b/tests/feature_tests/fsdp/test_profiling_estimator.py @@ -705,14 +705,18 @@ def _ce_group_name() -> str: def _register_shards(shapes, dtype=torch.bfloat16): """Register ``shapes`` as real symmetric-memory shards, filled distinctly.""" - from magi_compiler.symm_mem import alloc_shard, publish + 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) - publish() + # A peer read is only legal once every rank has filled its shard. + torch.cuda.synchronize() + dist.barrier() return shards