diff --git a/autoparallel/_testing/models/llama3.py b/autoparallel/_testing/models/llama3.py index 90db6653..a4f91d5f 100644 --- a/autoparallel/_testing/models/llama3.py +++ b/autoparallel/_testing/models/llama3.py @@ -600,6 +600,14 @@ def _custom_policy(ctx, func, *args, **kwargs): to_save = func in _OP_SAC_SAVE_LIST and not ( func == torch.ops.aten.mm.default and meta[mm_count_key] % 2 == 0 ) + # Also save the first residual add output (h = x + attn(...)) + # to match the reference's min-cut which saves both residual + # activations per block. + if func == torch.ops.aten.add.Tensor: + add_count_key = f"{mode}_add_count" + meta[add_count_key] += 1 + if meta[add_count_key] == 1: + to_save = True return ( CheckpointPolicy.MUST_SAVE if to_save diff --git a/autoparallel/api.py b/autoparallel/api.py index 1670d509..cda994e9 100644 --- a/autoparallel/api.py +++ b/autoparallel/api.py @@ -5,6 +5,7 @@ import copy import logging +import operator import time from contextlib import ExitStack, contextmanager from dataclasses import dataclass @@ -26,11 +27,13 @@ from .apply_sharding import apply_sharding_to_model from .cast_parametrization import apply_dtype_cast, canonicalize_mp, set_dtype_cast from .graph_passes.activation_checkpointing import mark_fsdp_all_gather_recomputation +from .graph_passes.fuse_allgather import fuse_chained_allgathers from .graph_passes.graph_utils import ( _add_alias, _replace_view_mm_view_with_einsum, assert_has_no_collectives, cleanup_graph, + eliminate_alias_round_trips, fix_scatter_on_aliased_inputs, functionalize_fresh_index_put_mutations, update_joint_with_descriptors, @@ -57,7 +60,36 @@ logger = logging.getLogger(__name__) -def _boxed_nop_preserve_node_meta(fx_g, example_inputs): +def _boxed_nop_preserve_node_meta( + fx_g, example_inputs, pre_pass=None, tag_forward=False +): + if pre_pass is not None: + pre_pass(fx_g.graph) + + if tag_forward: + # Tag the forward graph's OUTPUT values as "must save". These are + # the tensors the first min_cut decided to save for backward — + # only these should be saved in the second compilation. + # Uses the "custom" meta field (not "recompute") to avoid + # interfering with ac_joint_pass which uses "recompute" for + # activation checkpointing decisions. + output_node = next(n for n in fx_g.graph.nodes if n.op == "output") + for out in output_node.args[0]: + if not isinstance(out, torch.fx.Node) or out.op != "call_function": + continue + if out.target == operator.getitem: + # getitem metadata doesn't survive preserve_node_meta + # (Python builtin, not dispatched). Tag the parent + # multi-output op instead — DCE in _extract_fwd_bwd_modules + # removes any unused getitem children. + parent = out.args[0] + if isinstance(parent, torch.fx.Node): + parent.meta.setdefault("custom", {}) + parent.meta["custom"]["ap_must_save"] = True + else: + out.meta.setdefault("custom", {}) + out.meta["custom"]["ap_must_save"] = True + def run(args): with torch.fx.traceback.preserve_node_meta(): return torch.fx.Interpreter(fx_g).boxed_run(args) @@ -361,6 +393,10 @@ def optimize_placement(self, verbose=True): self.sharding_placement = self.sharding_optimizer.get_solution(verbose=False) + eliminated = eliminate_alias_round_trips(self.gm, self.sharding_placement) + if eliminated: + logger.info("Eliminated %d alias redistribution round-trips", eliminated) + if verbose: logger.info(self.sharding_optimizer.get_log(verbose=True)) @@ -473,6 +509,45 @@ def _apply_placement_common(self, sharding_placement): sharded_buffer_dict, ) + def _make_fuse_allgather_pass(self): + flat_mesh = self.mesh._flatten() if self.mesh.ndim > 1 else self.mesh + pg = flat_mesh.get_group() + full_group_size = flat_mesh.size() + full_group_name = pg.group_name + + subgroup_order = { + self.mesh.get_group(mesh_dim=dim).group_name: dim + for dim in range(self.mesh.ndim) + } + + # Column-major process group for ascending (dp→tp) chains from + # reverse shard order. Transposing the mesh tensor and flattening + # gives ranks in column-major order; sort_ranks=False preserves + # that order so allgather concatenates dp-fast. + reversed_full_group_name = None + if self.mesh.ndim == 2: + if not hasattr(self, "_reversed_flat_pg"): + import torch.distributed as dist + from torch._subclasses.fake_tensor import unset_fake_temporarily + + with unset_fake_temporarily(): + col_major_ranks = self.mesh.mesh.T.contiguous().view(-1).tolist() + self._reversed_flat_pg = dist.new_group( + ranks=col_major_ranks, sort_ranks=False + ) + reversed_full_group_name = self._reversed_flat_pg.group_name + + def pre_pass(graph): + fuse_chained_allgathers( + graph, + full_group_size, + full_group_name, + subgroup_order=subgroup_order, + reversed_full_group_name=reversed_full_group_name, + ) + + return pre_pass + def apply_placement(self, sharding_placement): sharded_param_dict, sharded_buffer_dict = self._apply_placement_common( sharding_placement @@ -482,10 +557,23 @@ def apply_placement(self, sharding_placement): self.parallel_gm.graph, self.reshard_after_forward ) + compiler_fn = self.compiler_fn + if self.mesh.ndim > 1: + from functools import partial + + compiler_fn = partial( + compiler_fn, + pre_pass=self._make_fuse_allgather_pass(), + ) + else: + from functools import partial + + fw_compiler_fn = partial(compiler_fn, tag_forward=True) + self.parallel_model_fn = parallel_model_fn = aot_compile_joint_with_descriptors( self.joint_with_descriptors, - fw_compiler=self.compiler_fn, - bw_compiler=self.compiler_fn, + fw_compiler=fw_compiler_fn, + bw_compiler=compiler_fn, ) # Build a forward-only graph for inference (no backward, no @@ -500,6 +588,13 @@ def apply_placement(self, sharding_placement): self.parallel_gm, num_fwd_outputs, num_primals ) compiler_fn = self.compiler_fn + if self.mesh.ndim > 1: + from functools import partial + + compiler_fn = partial( + compiler_fn, + pre_pass=self._make_fuse_allgather_pass(), + ) aot_config = self.joint_with_descriptors._aot_state.aot_config out_spec = self.joint_with_descriptors.out_spec diff --git a/autoparallel/compile.py b/autoparallel/compile.py index 3c73a66a..eb603504 100644 --- a/autoparallel/compile.py +++ b/autoparallel/compile.py @@ -3,13 +3,17 @@ # This source code is licensed under the BSD license found in the # LICENSE file in the root directory of this source tree. -from typing import Optional, Union +import operator +from contextlib import contextmanager +from typing import Any, Optional, Sequence, Union import torch import torch._functorch.config import torch._inductor.config from torch._inductor.compile_fx import compile_fx +from torch._inductor.custom_graph_pass import CustomPartitionerFn +from .api import _suppress_wait_tensor_side_effect from .graph_passes.activation_checkpointing import ac_joint_pass _INDUCTOR_OVERLAP_PATCHES = { @@ -20,6 +24,189 @@ } +class _SaveAllPartitioner(CustomPartitionerFn): + """Tag-driven partitioner: save all forward tensors except MUST_RECOMPUTE. + + The first compilation (inside AutoParallel) already makes recomputation + decisions via MUST_RECOMPUTE tags on FSDP allgather chains. Those tags + survive into the second compilation via preserve_node_meta. + + Unlike default_partition (which uses positional fwd/bwd boundary detection + and fails on interleaved backward ops like gradient reduce_scatters), this + partitioner uses classify_nodes for topology-based boundary detection. And + unlike min_cut_rematerialization_partition, it doesn't run a second min-cut + that would diverge from the first compilation's decisions. + """ + + def __call__( + self, + gm: torch.fx.GraphModule, + joint_inputs: Sequence[object], + **kwargs: Any, + ) -> tuple[torch.fx.GraphModule, torch.fx.GraphModule]: + num_fwd_outputs: int = kwargs.pop("num_fwd_outputs") # type: ignore[assignment] + static_lifetime_input_indices: list[int] | None = kwargs.pop( # type: ignore[assignment] + "static_lifetime_input_indices", None + ) + from torch._functorch.partitioners import ( + _extract_fwd_bwd_modules, + _is_assert_only_symbool, + classify_nodes, + cleanup_recompute_tags, + default_partition, + force_save_bw_mutation_src, + force_save_collectives, + force_save_effectful_ops, + functionalize_rng_ops, + has_recomputable_ops, + has_recomputable_rng_ops, + is_opaque_node, + is_sym_node, + must_recompute, + raise_getitems, + reordering_to_mimic_autograd_engine, + thread_graphsafe_rng_from_hops, + ) + + gm.graph.eliminate_dead_code() + gm.recompile() + + # CSE merges duplicate allgather chains: the forward's + # MUST_RECOMPUTE allgathers and the baked-in backward copies + # compute the same values from the same primals. Without CSE, + # both appear in the backward (duplication). + if torch._functorch.config.cse: + from torch._functorch.partitioners import fx_graph_cse + + cse_graph = fx_graph_cse(gm.graph) + gm.graph = cse_graph + + graph_has_recomputable_ops = has_recomputable_ops(gm) + graph_has_recomputable_rng_ops = has_recomputable_rng_ops(gm) + if graph_has_recomputable_ops: + gm = cleanup_recompute_tags(gm, is_default_partition=False) + + if not torch._functorch.config.unsafe_allow_optimization_of_collectives: + force_save_collectives(gm) + force_save_effectful_ops(gm) + force_save_bw_mutation_src(gm) + + if static_lifetime_input_indices is None: + static_lifetime_input_indices = [] + node_info = classify_nodes(gm, static_lifetime_input_indices, num_fwd_outputs) + + if len(node_info.required_bw_nodes) == 0: + return default_partition( + gm, + joint_inputs, + num_fwd_outputs=num_fwd_outputs, + static_lifetime_input_indices=static_lifetime_input_indices, + static_lifetime_input_nodes=node_info.static_lifetime_input_nodes, + ) + + saved_values = [] + saved_sym_nodes = [] + saved_opaque_nodes = [] + + def _is_multi_output(node: torch.fx.Node) -> bool: + return ( + all(user.target == operator.getitem for user in node.users) + and len(node.users) > 0 + ) + + def _maybe_save(node: torch.fx.Node) -> None: + if is_sym_node(node): + if not _is_assert_only_symbool(node): + saved_sym_nodes.append(node) + return + if _is_multi_output(node): + # Multi-output ops tagged ap_must_save: save all their + # getitem children (DCE removes unused ones later). + if node.meta.get("custom", {}).get("ap_must_save"): + for user in node.users: + if user.target == operator.getitem: + saved_values.append(user) + return + if is_opaque_node(node): + saved_opaque_nodes.append(node) + return + if must_recompute(node): + return + # Save nodes tagged ap_must_save by the first compilation. + # These are the forward graph's output tensors from the first + # partitioner — reproducing its save/recompute decisions. + # ac_joint_pass may override some with PREFER_RECOMPUTE for + # better memory; the must_recompute check above respects this. + # Placeholders (primals) are always saved when needed. + if node.op == "placeholder": + saved_values.append(node) + elif node.meta.get("custom", {}).get("ap_must_save"): + saved_values.append(node) + + for node in node_info.required_fw_nodes: + _maybe_save(node) + + # Unclaimed nodes (neither strictly forward nor backward) may be + # needed by backward outputs — e.g. mutable ops like index_put that + # are _must_be_in_forward. Save them so they're available as backward + # inputs in _extract_fwd_bwd_modules. + for node in node_info.unclaimed_nodes: + _maybe_save(node) + + fw_module, bw_module = _extract_fwd_bwd_modules( + gm, + saved_values, + saved_sym_nodes=saved_sym_nodes, + saved_opaque_nodes=saved_opaque_nodes, + num_fwd_outputs=num_fwd_outputs, + static_lifetime_input_nodes=node_info.static_lifetime_input_nodes, + ) + + if graph_has_recomputable_ops and graph_has_recomputable_rng_ops: + fw_module, bw_module = functionalize_rng_ops( + gm, fw_module, bw_module, len(saved_sym_nodes) + ) + bw_module = reordering_to_mimic_autograd_engine(bw_module) + + fw_module = raise_getitems(fw_module) + bw_module = raise_getitems(bw_module) + + fw_module = thread_graphsafe_rng_from_hops(fw_module, is_backward=False) + bw_module = thread_graphsafe_rng_from_hops(bw_module, is_backward=True) + + return fw_module, bw_module + + def uuid(self) -> Any: + return None + + +@contextmanager +def _patch_partitioner_dce(): + """Patch the partitioner's DCE to allow wait_tensor to be eliminated. + + The partitioner uses its own is_not_collective callback that treats all + _c10d_functional ops as impure, overriding _suppress_wait_tensor_side_effect. + We patch it to let wait_tensor through so unused collectives get DCE'd. + """ + import torch._functorch.partitioners as partitioners + + original = partitioners.is_not_collective + + def patched_is_not_collective(node): + if node.target in ( + torch.ops._c10d_functional.wait_tensor, + torch.ops._c10d_functional.wait_tensor.default, + ): + return False + return original(node) + + partitioners.is_not_collective = patched_is_not_collective + try: + yield + finally: + partitioners.is_not_collective = original + + def _make_ac_joint_pass( ac_stage_size_in_GiB: Optional[Union[float, str]] = "auto", ): @@ -50,16 +237,34 @@ def autoparallel_backend( sqrt(total_recomputable_memory). overlap_scheduling: Enable comm/compute overlap scheduling. """ - functorch_patches = {} - inductor_patches = _INDUCTOR_OVERLAP_PATCHES if overlap_scheduling else {} + functorch_patches: dict[str, Any] = {} if enable_ac: functorch_patches["joint_custom_pass"] = _make_ac_joint_pass( ac_stage_size_in_GiB ) + # Overlap scheduling configs must be set globally (not via config.patch + # context manager) because backward compilation is lazy — it happens on + # the first .backward() call, after compile_fx returns and the context + # manager has exited. + if overlap_scheduling: + cfg = torch._inductor.config.aten_distributed_optimizations + cfg.enable_overlap_scheduling = True + cfg.collective_bucketing = True + cfg.insert_overlap_deps = True + cfg.max_compute_pre_fetch = 10 + + # custom_partitioner_fn only applies to forward partitioning, so it can + # stay in the context manager. + inductor_patches: dict[str, Any] = { + "custom_partitioner_fn": _SaveAllPartitioner(), + } + def backend(gm, example_inputs): with ( + _suppress_wait_tensor_side_effect(), + _patch_partitioner_dce(), torch._functorch.config.patch(functorch_patches), torch._inductor.config.patch(inductor_patches), ): diff --git a/autoparallel/graph_passes/auto_bucketing.py b/autoparallel/graph_passes/auto_bucketing.py index fa01ff70..6aafcb70 100644 --- a/autoparallel/graph_passes/auto_bucketing.py +++ b/autoparallel/graph_passes/auto_bucketing.py @@ -4,16 +4,256 @@ # LICENSE file in the root directory of this source tree. import logging +from collections import Counter, defaultdict from functools import partial import torch from torch._inductor.fx_passes.overlap_scheduling import schedule_overlap_bucketing +from torch.utils._ordered_set import OrderedSet from .autobucketing_inductor import bucket_func, bucket_plan, bucket_utils, reorder logger = logging.getLogger(__name__) +def _patch_fsdp_bucketing(): + """Patch PyTorch's FSDP bucketing for better multi-group handling. + + Two fixes: + 1. Primary-group-only: only include the group with the most FSDP + all-gathers in fsdp_groups, preventing minority groups (tp, combined) + from limiting dp bucketing aggressiveness. + 2. Non-adjacent bucketing: allow collectives to be bucketed even when + interleaved with non-FSDP collectives on other groups. Only + descendant conflicts prevent bucketing, not graph position. + """ + import torch._inductor.fx_passes.bucketing as bucketing_mod + import torch._inductor.fx_passes.fsdp as fsdp_mod + from torch._inductor.fx_passes.bucketing import ( + collect_node_descendants, + is_wait_tensor, + ) + from torch._inductor.fx_passes.fsdp import ( + _find_all_gathers, + _get_group_name, + _get_group_size_from_node, + is_fsdp_all_gather, + ) + + def _patched_identify_fsdp_groups(gm): + fsdp_counts_by_group = Counter() + group_size = None + for n in _find_all_gathers(gm.graph): + if is_fsdp_all_gather(n): + gn = _get_group_name(n) + fsdp_counts_by_group[gn] += 1 + if group_size is None: + group_size = _get_group_size_from_node(n) + + if fsdp_counts_by_group: + primary_group = fsdp_counts_by_group.most_common(1)[0][0] + fsdp_groups = OrderedSet([primary_group]) + else: + fsdp_groups = OrderedSet() + + logger.debug( + "identify_fsdp_groups (patched): fsdp_groups=%s, all_counts=%s", + list(fsdp_groups), + dict(fsdp_counts_by_group), + ) + return fsdp_groups, group_size + + def _patched_greedy_bucket( + gm, + bucket_cap_mb_by_bucket_idx, + filter_node, + node_group_key, + filter_wait_node=None, + ): + g = gm.graph + groups = defaultdict(list) + for node in g.nodes: + if is_wait_tensor(node) and filter_node(node.args[0]): + if (filter_wait_node is None) or filter_wait_node(node): + coll_node = node.args[0] + key = node_group_key(coll_node) + groups[key].append(coll_node) + + if not groups: + return [] + + node_descendents = collect_node_descendants(g) + + buckets = [] + for key, nodes in groups.items(): + cur_bucket = [] + cur_bucket_descendents = OrderedSet() + cur_bucket_size_bytes = 0 + cur_bucket_id = 0 + bucket_size_bytes = int( + bucket_cap_mb_by_bucket_idx(cur_bucket_id) * 1024 * 1024 + ) + for node in nodes: + if node in cur_bucket_descendents: + continue + n_val = node.meta["val"] + out_size_bytes = n_val.numel() * n_val.element_size() + n_input_val = node.all_input_nodes[0].meta["val"] + in_size_bytes = n_input_val.numel() * n_input_val.element_size() + size_bytes = max(out_size_bytes, in_size_bytes) + if ( + cur_bucket_size_bytes + size_bytes > bucket_size_bytes + and cur_bucket + ): + if len(cur_bucket) > 1: + buckets.append(cur_bucket) + cur_bucket = [] + cur_bucket_size_bytes = 0 + cur_bucket_id += 1 + cur_bucket_descendents = OrderedSet() + bucket_size_bytes = int( + bucket_cap_mb_by_bucket_idx(cur_bucket_id) * 1024 * 1024 + ) + cur_bucket_size_bytes += size_bytes + cur_bucket.append(node) + cur_bucket_descendents |= node_descendents[node] + if len(cur_bucket) > 1: + buckets.append(cur_bucket) + return buckets + + fsdp_mod.identify_fsdp_groups = _patched_identify_fsdp_groups + bucketing_mod.greedy_bucket_collective_by_mb = _patched_greedy_bucket + + +_patch_fsdp_bucketing() + + +def _cap_compute_batch_size( + graph, original_compute_names, original_rs_after_compute, max_consecutive=8 +): + """Break up long compute segments between ReduceScatter operations. + + After overlap scheduling + FSDP bucketing, compute nodes may be reordered + so that many layers' matmuls execute before any ReduceScatter fires. This + inflates peak memory because all layers' activations are alive + simultaneously. + + Instead of restoring the full original order (which kills comm/compute + overlap), this function only intervenes when the number of compute nodes + between consecutive ReduceScatter ops exceeds max_consecutive. For each + oversized segment, it sorts compute nodes by original order, splits into + chunks, then: + 1. Chains consecutive compute nodes within each chunk (so they move + together during topological sort). + 2. Adds a dep from the first compute node of each chunk to an RS node + that originally appeared between the previous chunk and this one. + + Args: + graph: The post-scheduled FX graph. + original_compute_names: List of compute node names in original order. + original_rs_after_compute: Dict mapping compute node name to the list + of RS node names that appeared between it and the next compute node + in the original (pre-scheduling) graph. + max_consecutive: Maximum compute nodes allowed between RS ops. + """ + from torch._dynamo.graph_deduplication import _stable_topological_sort + from torch._inductor.fx_passes.overlap_scheduling import is_compute_node + + def _is_rs(node): + if node.op != "call_function": + return False + name = str(node.target) + return "reduce_scatter" in name and "wait" not in name + + original_rank = {name: rank for rank, name in enumerate(original_compute_names)} + node_by_name = {n.name: n for n in graph.nodes} + + scheduled_rs_names = { + n.name for n in graph.nodes if n.op == "call_function" and _is_rs(n) + } + + # Collect compute nodes between RS ops in the post-scheduled graph. + segments: list[tuple[list[torch.fx.Node], torch.fx.Node | None]] = [] + current_compute: list[torch.fx.Node] = [] + for node in graph.nodes: + if node.op != "call_function": + continue + if is_compute_node(node): + current_compute.append(node) + elif _is_rs(node): + segments.append((current_compute, node)) + current_compute = [] + if current_compute: + segments.append((current_compute, None)) + + additional_deps: dict[torch.fx.Node, OrderedSet] = defaultdict(OrderedSet) + + for compute_nodes, _seg_rs_node in segments: + if len(compute_nodes) <= max_consecutive: + continue + + sorted_nodes = sorted( + compute_nodes, + key=lambda n: original_rank.get(n.name, float("inf")), + ) + + # Split into chunks of max_consecutive. + chunks = [] + for i in range(0, len(sorted_nodes), max_consecutive): + chunks.append(sorted_nodes[i : i + max_consecutive]) + + # Chain consecutive compute nodes within each chunk. + for chunk in chunks: + for j in range(1, len(chunk)): + additional_deps[chunk[j]].add(chunk[j - 1]) + + # At each chunk boundary, find an RS that originally appeared between + # the last compute of the previous chunk and the first compute of + # the current chunk, then add dep: first_of_chunk after RS. + for ci in range(1, len(chunks)): + prev_chunk = chunks[ci - 1] + curr_chunk = chunks[ci] + + last_in_prev = prev_chunk[-1] + first_in_curr = curr_chunk[0] + last_rank = original_rank.get(last_in_prev.name, -1) + first_rank = original_rank.get( + first_in_curr.name, len(original_compute_names) + ) + + found_rs_node = None + for r in range(last_rank, first_rank): + cname = ( + original_compute_names[r] + if r < len(original_compute_names) + else None + ) + if cname is None: + continue + for rs_name in original_rs_after_compute.get(cname, []): + if rs_name in scheduled_rs_names: + rs_obj = node_by_name.get(rs_name) + if rs_obj is not None: + found_rs_node = rs_obj + break + if found_rs_node is not None: + break + + if found_rs_node is not None: + additional_deps[first_in_curr].add(found_rs_node) + + if not additional_deps: + return + + try: + _stable_topological_sort(graph, additional_deps) + except AssertionError: + logger.warning( + "Failed to cap compute batch size (cycle detected), " + "falling back to uncapped ordering" + ) + + class simplefsdp_autobucketing_config: """ Config for simplefsdp's autobucketing pass, which by default would give good performance. @@ -107,7 +347,7 @@ class aten_autobucketing_config: compute_overlap_multipler = 1.0 max_coll_distance = 100 custom_runtime_estimation = None - max_compute_pre_fetch = 5 + max_compute_pre_fetch = 50 collective_bucketing = False save_trace = True _counter = 0 @@ -117,6 +357,28 @@ def aten_autobucketing_reordering_pass( gm: torch.fx.Graph, configs: "aten_autobucketing_config" ) -> torch.fx.GraphModule: assert gm.owning_module is not None + + # Record compute + RS interleaving before bucketing + overlap scheduling. + from torch._inductor.fx_passes.overlap_scheduling import is_compute_node + + def _is_rs_node(node): + if node.op != "call_function": + return False + name = str(node.target) + return "reduce_scatter" in name and "wait" not in name + + original_compute_names = [] + original_rs_after_compute: dict[str, list[str]] = {} + last_compute_name = None + for n in gm.owning_module.graph.nodes: + if n.op != "call_function": + continue + if is_compute_node(n): + original_compute_names.append(n.name) + last_compute_name = n.name + elif _is_rs_node(n) and last_compute_name is not None: + original_rs_after_compute.setdefault(last_compute_name, []).append(n.name) + new_gm = schedule_overlap_bucketing( gm.owning_module, collective_bucketing=configs.collective_bucketing, @@ -126,6 +388,10 @@ def aten_autobucketing_reordering_pass( max_in_flight_gb=configs.max_in_flight_gb, max_coll_distance=configs.max_coll_distance, ) + + _cap_compute_batch_size( + new_gm.graph, original_compute_names, original_rs_after_compute + ) new_gm.recompile() if configs.save_trace: diff --git a/autoparallel/graph_passes/fuse_allgather.py b/autoparallel/graph_passes/fuse_allgather.py new file mode 100644 index 00000000..c47f9f3b --- /dev/null +++ b/autoparallel/graph_passes/fuse_allgather.py @@ -0,0 +1,285 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. +# +# This source code is licensed under the BSD license found in the +# LICENSE file in the root directory of this source tree. + +import logging + +import torch + +logger: logging.Logger = logging.getLogger(__name__) + + +def _is_all_gather(node: torch.fx.Node) -> bool: + return ( + node.op == "call_function" + and node.target == torch.ops._c10d_functional.all_gather_into_tensor.default + ) + + +def _is_wait_tensor(node: torch.fx.Node) -> bool: + return ( + node.op == "call_function" + and node.target == torch.ops._c10d_functional.wait_tensor.default + ) + + +def _is_nontrivial_dim_reorder(node: torch.fx.Node) -> bool: + if node.op != "call_function": + return False + if node.target == torch.ops.aten.t.default: + return True + if node.target == torch.ops.aten.transpose.int: + return node.args[1] != node.args[2] + if node.target == torch.ops.aten.permute.default and isinstance( + node.args[1], (list, tuple) + ): + dims = list(node.args[1]) + return dims != list(range(len(dims))) + return False + + +def _is_identity_view_chain(start: torch.fx.Node, end: torch.fx.Node) -> bool: + """Check that the view-op chain from start to end composes to the identity. + + Walks forward from ``start`` through single-user view ops and verifies + that the composed transformation doesn't change the data layout. + Uses FakeTensor metadata: if the output of ``start`` and the input of + ``end`` have the same shape and stride, the chain is an identity + (no data rearrangement, just metadata changes that cancel). + + Only allows ops that are true views (no data copy, no element removal): + permute, transpose, t, view, reshape, expand, unsqueeze, squeeze. + Rejects slice (can drop elements) and any non-view op. + + Returns False for empty chains or chains with no non-trivial dimension + reorder, since consecutive allgathers on different subgroups have + incompatible rank orderings without explicit layout reconciliation. + """ + _ALLOWED_VIEW_OPS = frozenset( + { + torch.ops.aten.permute.default, + torch.ops.aten.transpose.int, + torch.ops.aten.t.default, + torch.ops.aten.view.default, + torch.ops.aten.reshape.default, + torch.ops.aten.expand.default, + torch.ops.aten.unsqueeze.default, + torch.ops.aten.squeeze.default, + torch.ops.aten.squeeze.dim, + } + ) + + # Reject empty chains: no view ops means no layout reconciliation. + users = list(start.users.keys()) + if len(users) == 1 and users[0] is end: + return False + + start_val = start.meta.get("val") + if start_val is None: + return False + start_stride = start_val.stride() + + # Walk forward from start to end, verifying all intermediate ops are + # allowed views and that some op actually reorders dimensions. + node = start + saw_dim_reorder = False + while node is not end: + users = list(node.users.keys()) + if len(users) != 1: + return False + node = users[0] + if node is end: + break + if node.op != "call_function" or node.target not in _ALLOWED_VIEW_OPS: + return False + if _is_nontrivial_dim_reorder(node): + saw_dim_reorder = True + + if not saw_dim_reorder: + return False + + # Verify the composed transformation is identity via FakeTensor metadata. + ag2_input = end.args[0] + end_val = ( + ag2_input.meta.get("val") if isinstance(ag2_input, torch.fx.Node) else None + ) + + if end_val is None: + return False + if start_val.shape != end_val.shape: + return False + if start_stride != end_val.stride(): + return False + return True + + +def fuse_chained_allgathers( + graph: torch.fx.Graph, + full_group_size: int, + full_group_name: str, + subgroup_order: dict[str, int] | None = None, + reversed_full_group_name: str | None = None, +) -> int: + """Fuse consecutive allgather chains on different subgroups into a single allgather. + + Detects chains of two allgathers on different process groups connected + either directly (wait1 feeds into ag2) or through single-user view ops + that compose to the identity:: + + ag1 = all_gather(x, size1, pg1) + wait1 = wait_tensor(ag1) + ... = [optional identity_view_ops(wait1)] + ag2 = all_gather(..., size2, pg2) + wait2 = wait_tensor(ag2) + + and replaces them with:: + + full_ag = all_gather(x, size1 * size2, full_pg) + full_wait = wait_tensor(full_ag) + + The correct flattened process group depends on the chain direction: + - Descending mesh-dim order (tp before dp): ``full_group_name`` + (row-major flat mesh). + - Ascending mesh-dim order (dp before tp): ``reversed_full_group_name`` + (column-major flat mesh). + + Direct chains (no view ops between AGs) are accepted when + ``subgroup_order`` validates the direction. Without ``subgroup_order``, + an identity view chain is still required for safety. + + Returns the number of fusions performed. + """ + fusions = 0 + all_nodes = list(graph.nodes) + + for ag2 in all_nodes: + if not _is_all_gather(ag2): + continue + + # Walk ag2's input backward through single-user nodes to find wait1. + node = ag2.args[0] + if not isinstance(node, torch.fx.Node): + continue + + # Find the wait_tensor that starts the chain. + wait1 = node + while not _is_wait_tensor(wait1): + if len(wait1.users) != 1: + break + if len(wait1.args) == 0: + break + inp = wait1.args[0] + if not isinstance(inp, torch.fx.Node): + break + wait1 = inp + + if not _is_wait_tensor(wait1): + continue + if len(wait1.users) != 1: + continue + + ag1 = wait1.args[0] + if not isinstance(ag1, torch.fx.Node) or not _is_all_gather(ag1): + continue + if len(ag1.users) != 1: + continue + + # Validate that the view chain between wait1 and ag2 is identity, + # or that it's a direct chain (wait1 feeds ag2 with no intermediate ops). + is_direct_chain = ag2.args[0] is wait1 + if not is_direct_chain and not _is_identity_view_chain(wait1, ag2): + continue + + # Validate group sizes. + ag1_group_size = ag1.args[1] + ag2_group_size = ag2.args[1] + if ag1_group_size * ag2_group_size != full_group_size: + continue + + # Validate group names. + ag1_group = ag1.args[2] + ag2_group = ag2.args[2] + assert isinstance(ag1_group, str) + assert isinstance(ag2_group, str) + if ag1_group == ag2_group: + continue + + # Determine the correct flat mesh based on chain direction. + if subgroup_order is not None: + if ag1_group not in subgroup_order or ag2_group not in subgroup_order: + continue + ag1_dim = subgroup_order[ag1_group] + ag2_dim = subgroup_order[ag2_group] + if ag1_dim < ag2_dim: + # Ascending (dp→tp): needs column-major flat mesh + if reversed_full_group_name is None: + continue + target_group_name = reversed_full_group_name + elif ag1_dim > ag2_dim: + # Descending (tp→dp): needs row-major flat mesh + target_group_name = full_group_name + else: + continue + else: + # Without subgroup_order, require identity view chain for safety. + if is_direct_chain: + continue + target_group_name = full_group_name + + # Validate matching dtype. + ag1_val = ag1.meta.get("val") + ag2_val = ag2.meta.get("val") + if ( + ag1_val is not None + and ag2_val is not None + and ag1_val.dtype != ag2_val.dtype + ): + continue + + # Find wait2. + wait2 = None + for user in ag2.users: + if _is_wait_tensor(user): + wait2 = user + break + if wait2 is None: + continue + + # Build the fused allgather. + original_input = ag1.args[0] + + with graph.inserting_before(ag2): + full_ag = graph.call_function( + torch.ops._c10d_functional.all_gather_into_tensor.default, + args=(original_input, full_group_size, target_group_name), + ) + full_ag.meta.update(ag2.meta) + + full_wait = graph.call_function( + torch.ops._c10d_functional.wait_tensor.default, + args=(full_ag,), + ) + full_wait.meta.update(wait2.meta) + + wait2.replace_all_uses_with(full_wait) + fusions += 1 + + logger.debug( + "Fused ag(%s, gs=%d, pg=%s) + ag(gs=%d, pg=%s) -> ag(gs=%d, pg=%s)", + original_input, + ag1_group_size, + ag1_group, + ag2_group_size, + ag2_group, + full_group_size, + target_group_name, + ) + + if fusions > 0: + graph.eliminate_dead_code() + logger.info( + "Fused %d chained allgather pairs into full-mesh allgathers", fusions + ) + + return fusions diff --git a/autoparallel/graph_passes/graph_utils.py b/autoparallel/graph_passes/graph_utils.py index 434c65c6..e9c6c0a1 100644 --- a/autoparallel/graph_passes/graph_utils.py +++ b/autoparallel/graph_passes/graph_utils.py @@ -183,6 +183,59 @@ def delete_user_cb(n): return gm +def eliminate_alias_round_trips(gm: torch.fx.GraphModule, solution: dict) -> int: + """Bypass aliases for consumers that would redistribute back to the producer's placement. + + When an alias node A's chosen placement differs from its producer X's placement, + a consumer C of A whose input_spec matches X's placement would redistribute + A -> X-placement at runtime. Rewire such consumers to read X directly so the + larger intermediate (e.g. all-gathered) tensor isn't kept alive across them. + """ + eliminated = 0 + changed = False + for alias in list(gm.graph.nodes): + if alias.op != "call_function" or alias.target != torch.ops.aten.alias.default: + continue + if alias not in solution: + continue + producer = alias.args[0] + if not isinstance(producer, torch.fx.Node) or producer not in solution: + continue + x_placements = solution[producer].output_specs.placements + a_placements = solution[alias].output_specs.placements + if x_placements == a_placements: + continue + + for consumer in list(alias.users): + if consumer not in solution: + continue + c_input_specs = solution[consumer].input_specs + if c_input_specs is None: + continue + positions = [ + i for i, n in enumerate(all_input_nodes(consumer)) if n is alias + ] + if not positions: + continue + specs_at_positions = [c_input_specs[i] for i in positions] + if any(s is None for s in specs_at_positions): + continue + if not all(s.placements == x_placements for s in specs_at_positions): + continue + consumer.replace_input_with(alias, producer) + eliminated += 1 + changed = True + + if not alias.users: + del solution[alias] + gm.graph.erase_node(alias) + changed = True + + if changed: + gm.recompile() + return eliminated + + def is_collective(node: torch.fx.Node) -> bool: return ( node.op == "call_function" diff --git a/tests/test_activation_checkpointing.py b/tests/test_activation_checkpointing.py index 7da1ee1f..62faf8e4 100644 --- a/tests/test_activation_checkpointing.py +++ b/tests/test_activation_checkpointing.py @@ -1026,3 +1026,35 @@ def input_fn(): } assert expected_custom.items() <= fwd_sdpa.meta["custom"].items() assert expected_custom.items() <= bwd_sdpa.meta["custom"].items() + + +# --------------------------------------------------------------------------- +# Family 4: autoparallel_backend torch.compile round-trip +# --------------------------------------------------------------------------- + + +def test_autoparallel_backend_compile(device_mesh_2d): + """autoparallel_backend() compiles successfully with default_partition + and produces a backward that recomputes FSDP allgathers.""" + import torch._dynamo + + from autoparallel.compile import autoparallel_backend + + parallel_mod, bs, seq_len, dim = _build_parallel_module( + AttentionBlockNoAC, device_mesh_2d + ) + _materialize_parallel_module(parallel_mod) + + local_bs = bs // device_mesh_2d.shape[0] + x = torch.rand(local_bs, seq_len, dim, device="cuda", requires_grad=False) + + try: + compiled = torch.compile( + parallel_mod, + backend=autoparallel_backend(overlap_scheduling=False), + fullgraph=True, + ) + out = compiled(x) + out.sum().backward() + finally: + torch._dynamo.reset() diff --git a/tests/test_fsdp_all_gather_tagging.py b/tests/test_fsdp_all_gather_tagging.py index ab971c76..0af9ce44 100644 --- a/tests/test_fsdp_all_gather_tagging.py +++ b/tests/test_fsdp_all_gather_tagging.py @@ -158,6 +158,26 @@ def test_force_recompute_tags_dtype_cast_before_ag(): assert node.meta["ac_graph_id"] == AP_AC_GRAPH_ID +def test_force_recompute_tags_permute_after_wait(): + """A permute([1, 0]) after wait_tensor gets MUST_RECOMPUTE (single-input chain).""" + graph = _new_graph() + param = _add_placeholder(graph, "param") + activation = _add_placeholder(graph, "activation") + ag = _add_all_gather(graph, param) + wait = _add_wait_tensor(graph, ag) + permute = graph.call_function(torch.ops.aten.permute.default, args=(wait, [1, 0])) + permute.meta["val"] = torch.empty(64) + out = _add_mm(graph, permute, activation) + _add_output(graph, [out]) + + force_recompute_fsdp_all_gather(graph) + + for node in graph.nodes: + if node.target == torch.ops.aten.permute.default: + assert node.meta["recompute"] is CheckpointPolicy.MUST_RECOMPUTE + assert node.meta["ac_graph_id"] == AP_AC_GRAPH_ID + + def test_force_recompute_ignores_non_fsdp_all_gather(): """all_gather nodes that don't trace back to a placeholder are skipped.""" graph = _new_graph() diff --git a/tests/test_fuse_allgather.py b/tests/test_fuse_allgather.py new file mode 100644 index 00000000..ff3e73a7 --- /dev/null +++ b/tests/test_fuse_allgather.py @@ -0,0 +1,528 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. +# +# This source code is licensed under the BSD license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for chained allgather fusion pass. + +These tests build minimal FX graphs that mimic chained allgather patterns +without running the full AutoParallel pipeline or needing real process groups. +""" + +import pytest +import torch +import torch.fx + +from autoparallel.api import _suppress_wait_tensor_side_effect +from autoparallel.graph_passes.fuse_allgather import fuse_chained_allgathers + +AG = torch.ops._c10d_functional.all_gather_into_tensor.default +WAIT = torch.ops._c10d_functional.wait_tensor.default +PERMUTE = torch.ops.aten.permute.default + + +@pytest.fixture(autouse=True) +def suppress_wait_side_effect(): + """Allow DCE to remove wait_tensor nodes, matching the runtime environment.""" + with _suppress_wait_tensor_side_effect(): + yield + + +def _count_ops(graph, target): + return len(graph.find_nodes(op="call_function", target=target)) + + +def _add_placeholder(graph, name, shape): + node = graph.placeholder(name) + node.meta["val"] = torch.empty(*shape) + return node + + +def _add_all_gather(graph, input_node, group_size, group_name): + in_shape = input_node.meta["val"].shape + out_shape = (in_shape[0] * group_size, *in_shape[1:]) + node = graph.call_function(AG, args=(input_node, group_size, group_name)) + node.meta["val"] = torch.empty(*out_shape, dtype=input_node.meta["val"].dtype) + return node + + +def _add_wait_tensor(graph, input_node): + node = graph.call_function(WAIT, args=(input_node,)) + node.meta["val"] = input_node.meta["val"].clone() + return node + + +def _add_permute(graph, input_node, dims): + node = graph.call_function(PERMUTE, args=(input_node, dims)) + in_val = input_node.meta["val"] + node.meta["val"] = in_val.permute(dims) + return node + + +def _add_chained_allgather(graph, input_node, size1=16, pg1="dp", size2=8, pg2="tp"): + """Build: input -> ag1 -> wait -> permute([1,0]) -> permute([1,0]) -> ag2 -> wait.""" + ag1 = _add_all_gather(graph, input_node, size1, pg1) + wait1 = _add_wait_tensor(graph, ag1) + p1 = _add_permute(graph, wait1, [1, 0]) + p2 = _add_permute(graph, p1, [1, 0]) + ag2 = _add_all_gather(graph, p2, size2, pg2) + wait2 = _add_wait_tensor(graph, ag2) + return wait2 + + +def test_basic_fusion(): + """Detects chained allgathers and fuses into a single full-mesh allgather.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + result = _add_chained_allgather(graph, x) + graph.output((result,)) + + assert _count_ops(graph, AG) == 2 + assert _count_ops(graph, WAIT) == 2 + assert _count_ops(graph, PERMUTE) == 2 + + fusions = fuse_chained_allgathers( + graph, full_group_size=128, full_group_name="full" + ) + + assert fusions == 1 + assert _count_ops(graph, AG) == 1 + assert _count_ops(graph, WAIT) == 1 + assert _count_ops(graph, PERMUTE) == 0 + + ag_node = graph.find_nodes(op="call_function", target=AG)[0] + assert ag_node.args[1] == 128 + assert ag_node.args[2] == "full" + + +def test_multiple_chains(): + """Multiple independent chains in the same graph are all fused.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + y = _add_placeholder(graph, "y", (8, 4096)) + r1 = _add_chained_allgather(graph, x) + r2 = _add_chained_allgather(graph, y) + graph.output((r1, r2)) + + assert _count_ops(graph, AG) == 4 + + fusions = fuse_chained_allgathers( + graph, full_group_size=128, full_group_name="full" + ) + + assert fusions == 2 + assert _count_ops(graph, AG) == 2 + assert _count_ops(graph, WAIT) == 2 + + +def test_no_fusion_wrong_permute(): + """Non-[1,0] permutes prevent fusion (view chain doesn't trace through).""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + ag1 = _add_all_gather(graph, x, 16, "dp") + wait1 = _add_wait_tensor(graph, ag1) + p1 = _add_permute(graph, wait1, [1, 0]) + p2 = _add_permute(graph, p1, [0, 1]) # identity, not [1, 0] + ag2 = _add_all_gather(graph, p2, 8, "tp") + wait2 = _add_wait_tensor(graph, ag2) + graph.output((wait2,)) + + fusions = fuse_chained_allgathers( + graph, full_group_size=128, full_group_name="full" + ) + + assert fusions == 0 + assert _count_ops(graph, AG) == 2 + + +def test_no_fusion_wait_multiple_users(): + """If the first wait has other users, the intermediate result is consumed elsewhere.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + ag1 = _add_all_gather(graph, x, 16, "dp") + wait1 = _add_wait_tensor(graph, ag1) + p1 = _add_permute(graph, wait1, [1, 0]) + p2 = _add_permute(graph, p1, [1, 0]) + ag2 = _add_all_gather(graph, p2, 8, "tp") + wait2 = _add_wait_tensor(graph, ag2) + # wait1 also used directly in output + graph.output((wait2, wait1)) + + fusions = fuse_chained_allgathers( + graph, full_group_size=128, full_group_name="full" + ) + + assert fusions == 0 + assert _count_ops(graph, AG) == 2 + + +def test_no_fusion_group_size_mismatch(): + """If size1 * size2 != full_group_size, no fusion occurs.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + result = _add_chained_allgather(graph, x, size1=16, size2=8) + graph.output((result,)) + + # Wrong full_group_size + fusions = fuse_chained_allgathers(graph, full_group_size=64, full_group_name="full") + + assert fusions == 0 + assert _count_ops(graph, AG) == 2 + + +def test_no_fusion_same_group(): + """Two allgathers on the same process group are not fused.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + result = _add_chained_allgather(graph, x, size1=4, pg1="dp", size2=4, pg2="dp") + graph.output((result,)) + + fusions = fuse_chained_allgathers(graph, full_group_size=16, full_group_name="full") + + assert fusions == 0 + + +def test_subgroup_order_validation(): + """When subgroup_order is provided, only matching groups in valid order fuse.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + result = _add_chained_allgather(graph, x, size1=16, pg1="dp", size2=8, pg2="tp") + graph.output((result,)) + + # Unknown subgroup names — should not fuse + fusions = fuse_chained_allgathers( + graph, + full_group_size=128, + full_group_name="full", + subgroup_order={"other1": 0, "other2": 1}, + ) + assert fusions == 0 + assert _count_ops(graph, AG) == 2 + + # dp(0)→tp(1) ascending chain: needs reversed_full_group_name + fusions = fuse_chained_allgathers( + graph, + full_group_size=128, + full_group_name="full", + subgroup_order={"dp": 0, "tp": 1}, + reversed_full_group_name="full_rev", + ) + assert fusions == 1 + assert _count_ops(graph, AG) == 1 + + # Verify the fused AG uses the reversed group + ag_node = graph.find_nodes(op="call_function", target=AG)[0] + assert ag_node.args[2] == "full_rev" + + +def test_descending_subgroup_order_fuses(): + """Descending mesh-dim order (tp→dp) fuses with the default flat mesh.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + result = _add_chained_allgather(graph, x, size1=8, pg1="tp", size2=16, pg2="dp") + graph.output((result,)) + + fusions = fuse_chained_allgathers( + graph, + full_group_size=128, + full_group_name="full", + subgroup_order={"dp": 0, "tp": 1}, + ) + + assert fusions == 1 + assert _count_ops(graph, AG) == 1 + + ag_node = graph.find_nodes(op="call_function", target=AG)[0] + assert ag_node.args[2] == "full" + + +def test_with_cast_before_allgather(): + """The cast before the first allgather is preserved and becomes the fused ag's input.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + # Simulate dtype cast + cast = graph.call_function( + torch.ops.prims.convert_element_type.default, args=(x, torch.bfloat16) + ) + cast.meta["val"] = x.meta["val"].to(torch.bfloat16) + result = _add_chained_allgather(graph, cast) + graph.output((result,)) + + fusions = fuse_chained_allgathers( + graph, full_group_size=128, full_group_name="full" + ) + + assert fusions == 1 + assert _count_ops(graph, AG) == 1 + + # Cast should still be present + cast_count = _count_ops(graph, torch.ops.prims.convert_element_type.default) + assert cast_count == 1 + + # The allgather input should be the cast output + ag_node = graph.find_nodes(op="call_function", target=AG)[0] + assert ag_node.args[0].target == torch.ops.prims.convert_element_type.default + + +def test_standalone_allgather_untouched(): + """A single allgather without a chain is not affected.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + ag = _add_all_gather(graph, x, 8, "tp") + wait = _add_wait_tensor(graph, ag) + graph.output((wait,)) + + fusions = fuse_chained_allgathers( + graph, full_group_size=128, full_group_name="full" + ) + + assert fusions == 0 + assert _count_ops(graph, AG) == 1 + + +def test_direct_chain_no_views(): + """Two allgathers directly chained (ag1 -> wait -> ag2) with no view ops. + + Without subgroup_order, this is NOT fuseable (can't validate direction). + """ + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + ag1 = _add_all_gather(graph, x, 16, "dp") + wait1 = _add_wait_tensor(graph, ag1) + ag2 = _add_all_gather(graph, wait1, 8, "tp") + wait2 = _add_wait_tensor(graph, ag2) + graph.output((wait2,)) + + fusions = fuse_chained_allgathers( + graph, full_group_size=128, full_group_name="full" + ) + + assert fusions == 0 + assert _count_ops(graph, AG) == 2 + assert _count_ops(graph, WAIT) == 2 + + +def test_direct_chain_ascending_with_reversed_group(): + """Direct dp→tp chain fuses when subgroup_order and reversed group are provided.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + ag1 = _add_all_gather(graph, x, 16, "dp") + wait1 = _add_wait_tensor(graph, ag1) + ag2 = _add_all_gather(graph, wait1, 8, "tp") + wait2 = _add_wait_tensor(graph, ag2) + graph.output((wait2,)) + + fusions = fuse_chained_allgathers( + graph, + full_group_size=128, + full_group_name="full", + subgroup_order={"dp": 0, "tp": 1}, + reversed_full_group_name="full_rev", + ) + + assert fusions == 1 + assert _count_ops(graph, AG) == 1 + assert _count_ops(graph, WAIT) == 1 + + ag_node = graph.find_nodes(op="call_function", target=AG)[0] + assert ag_node.args[1] == 128 + assert ag_node.args[2] == "full_rev" + + +def test_direct_chain_descending_with_subgroup_order(): + """Direct tp→dp chain fuses with the default flat mesh.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + ag1 = _add_all_gather(graph, x, 8, "tp") + wait1 = _add_wait_tensor(graph, ag1) + ag2 = _add_all_gather(graph, wait1, 16, "dp") + wait2 = _add_wait_tensor(graph, ag2) + graph.output((wait2,)) + + fusions = fuse_chained_allgathers( + graph, + full_group_size=128, + full_group_name="full", + subgroup_order={"dp": 0, "tp": 1}, + ) + + assert fusions == 1 + assert _count_ops(graph, AG) == 1 + + ag_node = graph.find_nodes(op="call_function", target=AG)[0] + assert ag_node.args[2] == "full" + + +def test_ascending_without_reversed_group_does_not_fuse(): + """dp→tp chain without reversed_full_group_name cannot fuse.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + ag1 = _add_all_gather(graph, x, 16, "dp") + wait1 = _add_wait_tensor(graph, ag1) + ag2 = _add_all_gather(graph, wait1, 8, "tp") + wait2 = _add_wait_tensor(graph, ag2) + graph.output((wait2,)) + + fusions = fuse_chained_allgathers( + graph, + full_group_size=128, + full_group_name="full", + subgroup_order={"dp": 0, "tp": 1}, + # no reversed_full_group_name + ) + + assert fusions == 0 + assert _count_ops(graph, AG) == 2 + + +def test_noop_view_chain(): + """A no-op view/reshape between allgathers does not reconcile rank ordering. + + Even though a view op is present, if strides never change the chain is + semantically equivalent to a direct chain and must not be fused. + """ + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + ag1 = _add_all_gather(graph, x, 16, "dp") + wait1 = _add_wait_tensor(graph, ag1) + # No-op view: same shape, same strides + view = graph.call_function(torch.ops.aten.view.default, args=(wait1, [128, 4096])) + view.meta["val"] = wait1.meta["val"].clone() + ag2 = _add_all_gather(graph, view, 8, "tp") + wait2 = _add_wait_tensor(graph, ag2) + graph.output((wait2,)) + + fusions = fuse_chained_allgathers( + graph, full_group_size=128, full_group_name="full" + ) + + assert fusions == 0 + assert _count_ops(graph, AG) == 2 + + +def test_unsqueeze_squeeze_chain(): + """A temporary stride change without dimension reorder is not fuseable.""" + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (8, 4096)) + ag1 = _add_all_gather(graph, x, 16, "dp") + wait1 = _add_wait_tensor(graph, ag1) + unsqueeze = graph.call_function(torch.ops.aten.unsqueeze.default, args=(wait1, 0)) + unsqueeze.meta["val"] = wait1.meta["val"].unsqueeze(0) + squeeze = graph.call_function(torch.ops.aten.squeeze.dim, args=(unsqueeze, 0)) + squeeze.meta["val"] = unsqueeze.meta["val"].squeeze(0) + ag2 = _add_all_gather(graph, squeeze, 8, "tp") + wait2 = _add_wait_tensor(graph, ag2) + graph.output((wait2,)) + + fusions = fuse_chained_allgathers( + graph, full_group_size=128, full_group_name="full" + ) + + assert fusions == 0 + assert _count_ops(graph, AG) == 2 + + +def test_fuse_with_device_mesh(device_mesh_2d): + """Integration test: build fuse pass from a real 2D DeviceMesh and verify + that direct dp→tp chains fuse using the column-major process group.""" + import torch.distributed as dist + + mesh = device_mesh_2d + dp_size = mesh.size(0) + tp_size = mesh.size(1) + full_size = dp_size * tp_size + + # Row-major flat mesh (default) + flat_mesh = mesh._flatten() + full_group_name = flat_mesh.get_group().group_name + + # Column-major process group via sort_ranks=False + col_major_ranks = mesh.mesh.T.contiguous().view(-1).tolist() + reversed_pg = dist.new_group(ranks=col_major_ranks, sort_ranks=False) + reversed_full_group_name = reversed_pg.group_name + + assert full_group_name != reversed_full_group_name + + subgroup_order = { + mesh.get_group(mesh_dim=dim).group_name: dim for dim in range(mesh.ndim) + } + dp_pg = mesh.get_group(mesh_dim=0).group_name + tp_pg = mesh.get_group(mesh_dim=1).group_name + assert dp_pg != tp_pg + + # Build a direct dp→tp chain using real PG names. + graph = torch.fx.Graph() + x = _add_placeholder(graph, "x", (1, 4096)) + ag1 = _add_all_gather(graph, x, dp_size, dp_pg) + wait1 = _add_wait_tensor(graph, ag1) + ag2 = _add_all_gather(graph, wait1, tp_size, tp_pg) + wait2 = _add_wait_tensor(graph, ag2) + graph.output((wait2,)) + + fusions = fuse_chained_allgathers( + graph, + full_group_size=full_size, + full_group_name=full_group_name, + subgroup_order=subgroup_order, + reversed_full_group_name=reversed_full_group_name, + ) + + assert fusions == 1 + assert _count_ops(graph, AG) == 1 + + ag_node = graph.find_nodes(op="call_function", target=AG)[0] + assert ag_node.args[1] == full_size + assert ag_node.args[2] == reversed_full_group_name + + +def test_col_major_rank_ordering_correctness(): + """Verify that allgather on the column-major process group produces the + same result as two sequential allgathers (dp then tp) for reverse shard + order data. + + For reverse shard order S(0)S(0) on a (dp, tp) mesh, tp is the outer + shard and dp is the inner shard. Rank (dp=i, tp=j) holds the shard at + position j*dp_size + i in the full tensor. + + Sequential AG(dp) then AG(tp) recovers the full tensor in correct order. + A single AG on the column-major group must produce the same ordering. + """ + dp_size, tp_size = 4, 2 + full_size = dp_size * tp_size + chunk = 3 # rows per shard + + # Per-rank shards: rank r at (dp=r//tp, tp=r%tp) holds shard j*dp+i + rank_to_shard = {} + for r in range(full_size): + dp_idx, tp_idx = r // tp_size, r % tp_size + shard_pos = tp_idx * dp_size + dp_idx + rank_to_shard[r] = torch.arange( + shard_pos * chunk, (shard_pos + 1) * chunk, dtype=torch.float + ) + + # --- Sequential: AG(dp) then AG(tp) --- + # AG(dp): for each tp group, gather dp shards in rank order. + after_ag_dp = {} + for tp_idx in range(tp_size): + dp_group = sorted(r for r in range(full_size) if r % tp_size == tp_idx) + gathered = torch.cat([rank_to_shard[r] for r in dp_group]) + for r in dp_group: + after_ag_dp[r] = gathered + + # AG(tp): for each dp group, gather tp results in rank order. + for dp_idx in range(dp_size): + tp_group = sorted(r for r in range(full_size) if r // tp_size == dp_idx) + sequential_result = torch.cat([after_ag_dp[r] for r in tp_group]) + break # all ranks get the same result + + # --- Fused: single AG on column-major group --- + col_major_ranks = [] + for tp_idx in range(tp_size): + for dp_idx in range(dp_size): + col_major_ranks.append(dp_idx * tp_size + tp_idx) + fused_result = torch.cat([rank_to_shard[r] for r in col_major_ranks]) + + # Both must produce the original tensor in order [0, 1, 2, ..., N*chunk-1] + expected = torch.arange(full_size * chunk, dtype=torch.float) + torch.testing.assert_close(sequential_result, expected) + torch.testing.assert_close(fused_result, expected) diff --git a/tests/test_graph_utils.py b/tests/test_graph_utils.py index a5125d1a..7541919d 100644 --- a/tests/test_graph_utils.py +++ b/tests/test_graph_utils.py @@ -4,10 +4,14 @@ # LICENSE file in the root directory of this source tree. import torch +from torch.distributed._tensor.placement_types import DTensorSpec, TensorMeta +from torch.distributed.tensor._op_schema import OpSpec +from torch.distributed.tensor.placement_types import Replicate, Shard from torch.fx.experimental.proxy_tensor import make_fx from autoparallel.graph_passes.graph_utils import ( _replace_view_mm_view_with_einsum, + eliminate_alias_round_trips, functionalize_fresh_index_put_mutations, ) @@ -340,3 +344,203 @@ def test_extract_forward_deepcopy_with_tensor_constants(): # The real tensor constant should survive the deepcopy assert hasattr(result, "_tensor_constant0") assert result._tensor_constant0.device.type == "cuda" + + +# --- eliminate_alias_round_trips tests --- + + +def _make_spec(mesh, placements, shape, dtype=torch.float32): + t = torch.empty(shape, dtype=dtype, device="meta") + return DTensorSpec( + mesh, tuple(placements), tensor_meta=TensorMeta(t.shape, t.stride(), t.dtype) + ) + + +def _build_alias_graph(num_consumers): + """Build a GraphModule: x -> alias -> add(alias, alias, ..., y) per consumer. + + Returns (gm, x, alias, consumers). Each consumer is an add taking the + alias twice (so we cover the repeated-input case) plus a distinct + placeholder, so we can independently control its input_specs. + """ + gm = torch.fx.GraphModule({}, torch.fx.Graph()) + g = gm.graph + x = g.placeholder("x") + x.meta["val"] = torch.empty(4, device="meta") + alias = g.call_function(torch.ops.aten.alias.default, args=(x,)) + alias.meta["val"] = torch.empty(4, device="meta") + consumers = [] + for i in range(num_consumers): + y = g.placeholder(f"y{i}") + y.meta["val"] = torch.empty(4, device="meta") + c = g.call_function(torch.ops.aten.add.Tensor, args=(alias, y)) + c.meta["val"] = torch.empty(4, device="meta") + consumers.append(c) + g.output(tuple(consumers)) + gm.recompile() + return gm, x, alias, consumers + + +def test_eliminate_alias_round_trips_rewires_matching_consumer(device_mesh_2d): + mesh = device_mesh_2d + shape = (4,) + x_pl = (Shard(0), Shard(0)) + a_pl = (Shard(0), Replicate()) # alias all-gathered along tp + + gm, x, alias, [c_round_trip, c_use_ag] = _build_alias_graph(2) + sol = { + x: OpSpec(_make_spec(mesh, x_pl, shape)), + alias: OpSpec( + _make_spec(mesh, a_pl, shape), + input_specs=(_make_spec(mesh, a_pl, shape),), + ), + c_round_trip: OpSpec( + _make_spec(mesh, x_pl, shape), + input_specs=( + _make_spec(mesh, x_pl, shape), # would redistribute alias R -> S(1) + _make_spec(mesh, x_pl, shape), + ), + ), + c_use_ag: OpSpec( + _make_spec(mesh, a_pl, shape), + input_specs=( + _make_spec(mesh, a_pl, shape), # uses alias directly at S(0)R + _make_spec(mesh, a_pl, shape), + ), + ), + } + + eliminated = eliminate_alias_round_trips(gm, sol) + + assert eliminated == 1 + assert list(c_round_trip.all_input_nodes)[0] is x + assert list(c_use_ag.all_input_nodes)[0] is alias + assert alias in sol # still referenced by c_use_ag + + +def test_eliminate_alias_round_trips_erases_when_no_users(device_mesh_2d): + mesh = device_mesh_2d + shape = (4,) + x_pl = (Shard(0), Shard(0)) + a_pl = (Shard(0), Replicate()) + + gm, x, alias, [c1, c2] = _build_alias_graph(2) + sol = { + x: OpSpec(_make_spec(mesh, x_pl, shape)), + alias: OpSpec( + _make_spec(mesh, a_pl, shape), + input_specs=(_make_spec(mesh, a_pl, shape),), + ), + c1: OpSpec( + _make_spec(mesh, x_pl, shape), + input_specs=(_make_spec(mesh, x_pl, shape), _make_spec(mesh, x_pl, shape)), + ), + c2: OpSpec( + _make_spec(mesh, x_pl, shape), + input_specs=(_make_spec(mesh, x_pl, shape), _make_spec(mesh, x_pl, shape)), + ), + } + + eliminated = eliminate_alias_round_trips(gm, sol) + + assert eliminated == 2 + assert alias not in sol + alias_nodes = gm.graph.find_nodes( + op="call_function", target=torch.ops.aten.alias.default + ) + assert alias_nodes == [] + + +def test_eliminate_alias_round_trips_noop_when_placements_match(device_mesh_2d): + mesh = device_mesh_2d + shape = (4,) + pl = (Shard(0), Replicate()) + + gm, x, alias, [c] = _build_alias_graph(1) + sol = { + x: OpSpec(_make_spec(mesh, pl, shape)), + alias: OpSpec( + _make_spec(mesh, pl, shape), + input_specs=(_make_spec(mesh, pl, shape),), + ), + c: OpSpec( + _make_spec(mesh, pl, shape), + input_specs=(_make_spec(mesh, pl, shape), _make_spec(mesh, pl, shape)), + ), + } + + eliminated = eliminate_alias_round_trips(gm, sol) + + assert eliminated == 0 + assert alias in sol + assert list(c.all_input_nodes)[0] is alias + + +def test_eliminate_alias_round_trips_skips_intermediate_redistribution(device_mesh_2d): + """Consumer that doesn't want either x's or alias's placement is left alone.""" + mesh = device_mesh_2d + shape = (4,) + x_pl = (Shard(0), Shard(0)) + a_pl = (Shard(0), Replicate()) + other_pl = (Replicate(), Replicate()) + + gm, x, alias, [c] = _build_alias_graph(1) + sol = { + x: OpSpec(_make_spec(mesh, x_pl, shape)), + alias: OpSpec( + _make_spec(mesh, a_pl, shape), + input_specs=(_make_spec(mesh, a_pl, shape),), + ), + c: OpSpec( + _make_spec(mesh, other_pl, shape), + input_specs=( + _make_spec(mesh, other_pl, shape), + _make_spec(mesh, other_pl, shape), + ), + ), + } + + eliminated = eliminate_alias_round_trips(gm, sol) + + assert eliminated == 0 + assert list(c.all_input_nodes)[0] is alias + + +def test_eliminate_alias_round_trips_handles_repeated_input(device_mesh_2d): + """A consumer using alias in multiple positions is rewired exactly once.""" + mesh = device_mesh_2d + shape = (4,) + x_pl = (Shard(0), Shard(0)) + a_pl = (Shard(0), Replicate()) + + gm = torch.fx.GraphModule({}, torch.fx.Graph()) + g = gm.graph + x = g.placeholder("x") + x.meta["val"] = torch.empty(4, device="meta") + alias = g.call_function(torch.ops.aten.alias.default, args=(x,)) + alias.meta["val"] = torch.empty(4, device="meta") + # Consumer takes alias as both args (e.g. x + x) + c = g.call_function(torch.ops.aten.add.Tensor, args=(alias, alias)) + c.meta["val"] = torch.empty(4, device="meta") + g.output((c,)) + gm.recompile() + + sol = { + x: OpSpec(_make_spec(mesh, x_pl, shape)), + alias: OpSpec( + _make_spec(mesh, a_pl, shape), + input_specs=(_make_spec(mesh, a_pl, shape),), + ), + c: OpSpec( + _make_spec(mesh, x_pl, shape), + input_specs=(_make_spec(mesh, x_pl, shape), _make_spec(mesh, x_pl, shape)), + ), + } + + eliminated = eliminate_alias_round_trips(gm, sol) + + # Exactly one consumer was rewired (positions for the same consumer fold + # into a single replace_input_with call). + assert eliminated == 1 + assert all(n is x for n in c.all_input_nodes) + assert alias not in sol # alias has no remaining users → erased