From 434cadebd18a77659b17d1e71da1ec443989af1c Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Wed, 20 May 2026 07:34:58 +0000
Subject: [PATCH 01/12] Fix double recomputation in torch.compile round-trip

When `torch.compile(parallel_mod, backend=autoparallel_backend())` compiles an
already-partitioned module, the second AOT autograd partitioner was making
different recomputation decisions from the first, causing redundant computation
and wasted memory.

Three fixes, best reviewed in order:

**1. Tag the full downstream chain in `force_recompute_fsdp_all_gather`**
(`activation_checkpointing.py`). The previous code only tagged allgather,
wait_tensor, slice, and dtype_cast with MUST_RECOMPUTE, missing the
`permute([1,0])` that follows. The first partitioner saved the permute output
(the transposed allgathered weight, 2.62 GiB total), defeating MUST_RECOMPUTE.
The fix walks the single-input chain downstream from wait_tensor and tags all
layout ops.

**2. Replace the second partitioner with `_SaveAllPartitioner`**
(`compile.py`). Instead of running a second `min_cut_rematerialization_partition`
that diverges, use a custom tag-driven partitioner that reproduces the first
min_cut's save/recompute decisions. It uses `classify_nodes` for correct
fwd/bwd boundary detection (handles interleaved backward ops that
`default_partition` can't), CSE to deduplicate baked-in backward allgather
chains, and saves only nodes tagged `ap_must_save` by the first compilation.

**3. Tag forward/backward nodes and saved tensors in the first compilation**
(`api.py`). The first compilation's `fw_compiler` tags forward nodes with
`custom: {ap_graph_part: "forward"}` and marks min_cut's saved-for-backward
tensors with `ap_must_save: True` (or `ap_save_getitems` for multi-output
ops like SDPA). The `bw_compiler` tags backward nodes with
`custom: {ap_graph_part: "backward"}`. Tags survive into the second
compilation via `preserve_node_meta` and `_COPY_META_FIELDS`.

Result (LLaMA-3 8B, 32 layers, 128 GPUs): second compilation's forward
allgather count exactly matches the first (357 = 357). Forward peak memory
6.40 GiB, backward 8.21 GiB.

Authored with Claude.
---
 autoparallel/api.py                           |  72 +++++-
 autoparallel/compile.py                       | 215 +++++++++++++++++-
 .../graph_passes/activation_checkpointing.py  |  19 +-
 autoparallel/graph_passes/fuse_allgather.py   | 156 +++++++++++++
 tests/test_activation_checkpointing.py        |  35 +++
 tests/test_fsdp_all_gather_tagging.py         |  20 ++
 6 files changed, 504 insertions(+), 13 deletions(-)
 create mode 100644 autoparallel/graph_passes/fuse_allgather.py

diff --git a/autoparallel/api.py b/autoparallel/api.py
index 1670d509..fe378241 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,6 +27,7 @@
 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_dp_tp_allgather
 from .graph_passes.graph_utils import (
     _add_alias,
     _replace_view_mm_view_with_einsum,
@@ -57,7 +59,46 @@
 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, tag_backward=False
+):
+    if pre_pass is not None:
+        pre_pass(fx_g.graph)
+
+    if tag_forward:
+        # Tag all forward call_function nodes so the second compilation
+        # knows they belong to the original forward graph.
+        for node in fx_g.graph.nodes:
+            if node.op == "call_function":
+                node.meta.setdefault("custom", {})
+                node.meta["custom"]["ap_graph_part"] = "forward"
+        # Additionally 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.
+        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 nodes don't get their custom field preserved by
+                # preserve_node_meta (operator.getitem is a Python builtin,
+                # not dispatched through TorchDispatchMode). Tag the parent
+                # multi-output op with the specific getitem indices to save.
+                parent = out.args[0]
+                if isinstance(parent, torch.fx.Node):
+                    parent.meta.setdefault("custom", {})
+                    parent.meta["custom"].setdefault("ap_save_getitems", [])
+                    parent.meta["custom"]["ap_save_getitems"].append(out.args[1])
+            else:
+                out.meta.setdefault("custom", {})
+                out.meta["custom"]["ap_must_save"] = True
+
+    if tag_backward:
+        for node in fx_g.graph.nodes:
+            if node.op == "call_function":
+                node.meta.setdefault("custom", {})
+                node.meta["custom"]["ap_graph_part"] = "backward"
+
     def run(args):
         with torch.fx.traceback.preserve_node_meta():
             return torch.fx.Interpreter(fx_g).boxed_run(args)
@@ -473,6 +514,17 @@ 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
+
+        def pre_pass(graph):
+            fuse_dp_tp_allgather(graph, full_group_size, 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 +534,24 @@ 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)
+        bw_compiler_fn = partial(compiler_fn, tag_backward=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=bw_compiler_fn,
         )
 
         # Build a forward-only graph for inference (no backward, no
diff --git a/autoparallel/compile.py b/autoparallel/compile.py
index 3c73a66a..981838da 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,205 @@
 }
 
 
+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)
+
+        # Nodes from the first compilation's backward and untagged nodes
+        # (fresh autograd backward ops) should not appear in the forward
+        # graph. Mark them must_be_in_backward so classify_nodes excludes
+        # them from the forward.
+        for node in gm.graph.nodes:
+            if node.op != "call_function":
+                continue
+            if node.meta.get("custom", {}).get("ap_graph_part") != "forward":
+                node.meta["partitioner_tag"] = "must_be_in_backward"
+
+        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):
+                return
+            if is_opaque_node(node):
+                saved_opaque_nodes.append(node)
+                return
+            if must_recompute(node):
+                return
+            # Only save nodes that the first compilation's min_cut decided
+            # to save (tagged ap_must_save via the "custom" meta field).
+            # These are exactly the forward graph's output tensors from the
+            # first partitioner — reproducing its save/recompute decisions.
+            # Placeholders (primals) don't carry the tag but must always be
+            # saved when needed.
+            #
+            # For getitem nodes (outputs of multi-output ops like SDPA),
+            # preserve_node_meta doesn't copy the custom field (Python
+            # builtin, not dispatched). Instead, the parent op carries
+            # ap_save_getitems with the specific indices to save.
+            custom = node.meta.get("custom", {})
+            if node.op == "placeholder":
+                saved_values.append(node)
+            elif custom.get("ap_must_save", False):
+                saved_values.append(node)
+            elif node.target == operator.getitem and isinstance(
+                node.args[0], torch.fx.Node
+            ):
+                parent_custom = node.args[0].meta.get("custom", {})
+                if node.args[1] in parent_custom.get("ap_save_getitems", []):
+                    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 +253,22 @@ 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] = {}
+    inductor_patches: dict[str, Any] = (
+        dict(_INDUCTOR_OVERLAP_PATCHES) if overlap_scheduling else {}
+    )
 
     if enable_ac:
         functorch_patches["joint_custom_pass"] = _make_ac_joint_pass(
             ac_stage_size_in_GiB
         )
 
+    inductor_patches["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/activation_checkpointing.py b/autoparallel/graph_passes/activation_checkpointing.py
index 6c9e7f31..0b7be1bd 100644
--- a/autoparallel/graph_passes/activation_checkpointing.py
+++ b/autoparallel/graph_passes/activation_checkpointing.py
@@ -250,13 +250,18 @@ def force_recompute_node(node):
             ag_node = node.args[0]
             force_recompute_node(ag_node)  # all_gather
             force_recompute_node(node)  # wait_tensor
-            # Force-recompute slice that comes after wait
-            for user in node.users:
-                if (
-                    user.op == "call_function"
-                    and user.target == torch.ops.aten.slice.Tensor
-                ):
-                    force_recompute_node(user)
+            # Force-recompute the downstream single-input chain from
+            # wait_tensor (permute, slice, reshape, etc.). Without this,
+            # the partitioner saves the first untagged downstream node,
+            # which holds the same allgathered data and defeats the
+            # MUST_RECOMPUTE on the allgather itself.
+            w = node
+            while len(w.users) == 1:
+                user = next(iter(w.users))
+                if len(user.all_input_nodes) > 1:
+                    break
+                force_recompute_node(user)
+                w = user
             # Force-recompute potential dtype casts from all_gather
             if (
                 ag_node.all_input_nodes[0].op == "call_function"
diff --git a/autoparallel/graph_passes/fuse_allgather.py b/autoparallel/graph_passes/fuse_allgather.py
new file mode 100644
index 00000000..c4580c03
--- /dev/null
+++ b/autoparallel/graph_passes/fuse_allgather.py
@@ -0,0 +1,156 @@
+# 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_permute_transpose(node: torch.fx.Node) -> bool:
+    return (
+        node.op == "call_function"
+        and node.target == torch.ops.aten.permute.default
+        and isinstance(node.args[1], (list, tuple))
+        and list(node.args[1]) == [1, 0]
+    )
+
+
+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 fuse_dp_tp_allgather(
+    graph: torch.fx.Graph,
+    full_group_size: int,
+    full_group_name: str,
+) -> int:
+    """Fuse consecutive dp + tp allgather chains into a single full-mesh allgather.
+
+    Detects the pattern::
+
+        dp_ag    = all_gather(x, dp_size, dp_pg)
+        dp_wait  = wait_tensor(dp_ag)
+        p1       = permute(dp_wait, [1, 0])
+        p2       = permute(p1, [1, 0])
+        tp_ag    = all_gather(p2, tp_size, tp_pg)
+        tp_wait  = wait_tensor(tp_ag)
+
+    and replaces it with::
+
+        full_ag   = all_gather(x, dp_size * tp_size, full_pg)
+        full_wait = wait_tensor(full_ag)
+
+    The two permutes must cancel (both ``[1, 0]`` transposes), the dp wait
+    must have a single user (the first permute), and the dp and tp group
+    sizes must multiply to ``full_group_size``.
+
+    Returns the number of fusions performed.
+    """
+    fusions = 0
+
+    # Snapshot the node list — we mutate the graph during iteration.
+    all_nodes = list(graph.nodes)
+
+    for tp_ag in all_nodes:
+        if not _is_all_gather(tp_ag):
+            continue
+
+        # --- trace backwards through the expected chain ---
+
+        p2 = tp_ag.args[0]
+        if not isinstance(p2, torch.fx.Node) or not _is_permute_transpose(p2):
+            continue
+        if len(p2.users) != 1:
+            continue
+
+        p1 = p2.args[0]
+        if not isinstance(p1, torch.fx.Node) or not _is_permute_transpose(p1):
+            continue
+        if len(p1.users) != 1:
+            continue
+
+        dp_wait = p1.args[0]
+        if not isinstance(dp_wait, torch.fx.Node) or not _is_wait_tensor(dp_wait):
+            continue
+        if len(dp_wait.users) != 1:
+            continue
+
+        dp_ag = dp_wait.args[0]
+        if not isinstance(dp_ag, torch.fx.Node) or not _is_all_gather(dp_ag):
+            continue
+        if len(dp_ag.users) != 1:
+            continue
+
+        # --- validate group sizes ---
+
+        dp_group_size = dp_ag.args[1]
+        tp_group_size = tp_ag.args[1]
+        if dp_group_size * tp_group_size != full_group_size:
+            continue
+
+        # --- validate matching dtype ---
+
+        dp_val = dp_ag.meta.get("val")
+        tp_val = tp_ag.meta.get("val")
+        if dp_val is not None and tp_val is not None and dp_val.dtype != tp_val.dtype:
+            continue
+
+        # --- find the tp wait (sole user of tp_ag) ---
+
+        tp_wait = None
+        for user in tp_ag.users:
+            if _is_wait_tensor(user):
+                tp_wait = user
+                break
+        if tp_wait is None:
+            continue
+
+        # --- build the fused allgather ---
+
+        original_input = dp_ag.args[0]
+
+        with graph.inserting_before(tp_ag):
+            full_ag = graph.call_function(
+                torch.ops._c10d_functional.all_gather_into_tensor.default,
+                args=(original_input, full_group_size, full_group_name),
+            )
+            full_ag.meta.update(tp_ag.meta)
+
+            full_wait = graph.call_function(
+                torch.ops._c10d_functional.wait_tensor.default,
+                args=(full_ag,),
+            )
+            full_wait.meta.update(tp_wait.meta)
+
+        tp_wait.replace_all_uses_with(full_wait)
+        fusions += 1
+
+        logger.debug(
+            "Fused dp_ag(%s, gs=%d) + tp_ag(gs=%d) -> full_ag(gs=%d)",
+            original_input,
+            dp_group_size,
+            tp_group_size,
+            full_group_size,
+        )
+
+    if fusions > 0:
+        graph.eliminate_dead_code()
+        logger.info(
+            "Fused %d dp+tp allgather chains into full-mesh allgathers", fusions
+        )
+
+    return fusions
diff --git a/tests/test_activation_checkpointing.py b/tests/test_activation_checkpointing.py
index abb82f4e..89c0f8c1 100644
--- a/tests/test_activation_checkpointing.py
+++ b/tests/test_activation_checkpointing.py
@@ -918,3 +918,38 @@ def input_fn():
         "inside_local_map": 2,
         "outside_checkpoint": 0,
     }
+
+
+# ---------------------------------------------------------------------------
+# 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
+    from autoparallel.graph_passes.activation_checkpointing import (
+        is_all_gather_into_tensor,
+    )
+
+    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()

From 28153aa57f7bb4280e3f6dd5bc2b3429f1612217 Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Wed, 20 May 2026 11:29:17 +0000
Subject: [PATCH 02/12] Simplify saved-tensor tagging: remove ap_save_getitems,
 tag parent instead
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Simplifies how the first compilation marks saved-for-backward tensors for
the second compilation to reproduce. No functional change — allgather counts
and peak memory are identical before and after.

For multi-output ops like SDPA, getitem metadata doesn't survive
`preserve_node_meta` (Python builtin, not dispatched). The previous approach
tracked specific getitem indices via `custom: {ap_save_getitems: [0, 1, 6, 7]}`
on the parent. This replaces it with `custom: {ap_must_save: True}` on the
parent — the second partitioner saves all getitem children, and
`_extract_fwd_bwd_modules`'s DCE removes the ones backward doesn't need.

Also removes dead code: `tag_backward` parameter and `ap_graph_part` tagging
were only used by the `must_be_in_backward` approach that was removed.

Authored with Claude.
---
 autoparallel/api.py     | 37 +++++++++++++-----------------------
 autoparallel/compile.py | 42 +++++++++++++----------------------------
 2 files changed, 26 insertions(+), 53 deletions(-)

diff --git a/autoparallel/api.py b/autoparallel/api.py
index fe378241..1b90218f 100644
--- a/autoparallel/api.py
+++ b/autoparallel/api.py
@@ -60,45 +60,35 @@
 
 
 def _boxed_nop_preserve_node_meta(
-    fx_g, example_inputs, pre_pass=None, tag_forward=False, tag_backward=False
+    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 all forward call_function nodes so the second compilation
-        # knows they belong to the original forward graph.
-        for node in fx_g.graph.nodes:
-            if node.op == "call_function":
-                node.meta.setdefault("custom", {})
-                node.meta["custom"]["ap_graph_part"] = "forward"
-        # Additionally 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.
+        # 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 nodes don't get their custom field preserved by
-                # preserve_node_meta (operator.getitem is a Python builtin,
-                # not dispatched through TorchDispatchMode). Tag the parent
-                # multi-output op with the specific getitem indices to save.
+                # 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"].setdefault("ap_save_getitems", [])
-                    parent.meta["custom"]["ap_save_getitems"].append(out.args[1])
+                    parent.meta["custom"]["ap_must_save"] = True
             else:
                 out.meta.setdefault("custom", {})
                 out.meta["custom"]["ap_must_save"] = True
 
-    if tag_backward:
-        for node in fx_g.graph.nodes:
-            if node.op == "call_function":
-                node.meta.setdefault("custom", {})
-                node.meta["custom"]["ap_graph_part"] = "backward"
-
     def run(args):
         with torch.fx.traceback.preserve_node_meta():
             return torch.fx.Interpreter(fx_g).boxed_run(args)
@@ -546,12 +536,11 @@ def apply_placement(self, sharding_placement):
             from functools import partial
 
         fw_compiler_fn = partial(compiler_fn, tag_forward=True)
-        bw_compiler_fn = partial(compiler_fn, tag_backward=True)
 
         self.parallel_model_fn = parallel_model_fn = aot_compile_joint_with_descriptors(
             self.joint_with_descriptors,
             fw_compiler=fw_compiler_fn,
-            bw_compiler=bw_compiler_fn,
+            bw_compiler=compiler_fn,
         )
 
         # Build a forward-only graph for inference (no backward, no
diff --git a/autoparallel/compile.py b/autoparallel/compile.py
index 981838da..a40bbedc 100644
--- a/autoparallel/compile.py
+++ b/autoparallel/compile.py
@@ -91,16 +91,6 @@ def __call__(
         force_save_effectful_ops(gm)
         force_save_bw_mutation_src(gm)
 
-        # Nodes from the first compilation's backward and untagged nodes
-        # (fresh autograd backward ops) should not appear in the forward
-        # graph. Mark them must_be_in_backward so classify_nodes excludes
-        # them from the forward.
-        for node in gm.graph.nodes:
-            if node.op != "call_function":
-                continue
-            if node.meta.get("custom", {}).get("ap_graph_part") != "forward":
-                node.meta["partitioner_tag"] = "must_be_in_backward"
-
         if static_lifetime_input_indices is None:
             static_lifetime_input_indices = []
         node_info = classify_nodes(gm, static_lifetime_input_indices, num_fwd_outputs)
@@ -130,34 +120,28 @@ def _maybe_save(node: torch.fx.Node) -> None:
                     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
-            # Only save nodes that the first compilation's min_cut decided
-            # to save (tagged ap_must_save via the "custom" meta field).
-            # These are exactly the forward graph's output tensors from the
-            # first partitioner — reproducing its save/recompute decisions.
-            # Placeholders (primals) don't carry the tag but must always be
-            # saved when needed.
-            #
-            # For getitem nodes (outputs of multi-output ops like SDPA),
-            # preserve_node_meta doesn't copy the custom field (Python
-            # builtin, not dispatched). Instead, the parent op carries
-            # ap_save_getitems with the specific indices to save.
-            custom = node.meta.get("custom", {})
+            # 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 custom.get("ap_must_save", False):
+            elif node.meta.get("custom", {}).get("ap_must_save"):
                 saved_values.append(node)
-            elif node.target == operator.getitem and isinstance(
-                node.args[0], torch.fx.Node
-            ):
-                parent_custom = node.args[0].meta.get("custom", {})
-                if node.args[1] in parent_custom.get("ap_save_getitems", []):
-                    saved_values.append(node)
 
         for node in node_info.required_fw_nodes:
             _maybe_save(node)

From a03aba256e6756c41422a3687c9af3863a7848df Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Wed, 20 May 2026 13:52:18 +0000
Subject: [PATCH 03/12] Fix lint

---
 tests/test_activation_checkpointing.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tests/test_activation_checkpointing.py b/tests/test_activation_checkpointing.py
index 89c0f8c1..fd8826c4 100644
--- a/tests/test_activation_checkpointing.py
+++ b/tests/test_activation_checkpointing.py
@@ -931,9 +931,6 @@ def test_autoparallel_backend_compile(device_mesh_2d):
     import torch._dynamo
 
     from autoparallel.compile import autoparallel_backend
-    from autoparallel.graph_passes.activation_checkpointing import (
-        is_all_gather_into_tensor,
-    )
 
     parallel_mod, bs, seq_len, dim = _build_parallel_module(
         AttentionBlockNoAC, device_mesh_2d

From 6e7727e9070311e89db609814f5f89e9df249b49 Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Wed, 20 May 2026 13:58:44 +0000
Subject: [PATCH 04/12] Fix test

---
 tests/test_activation_checkpointing.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/tests/test_activation_checkpointing.py b/tests/test_activation_checkpointing.py
index fd8826c4..054e79f1 100644
--- a/tests/test_activation_checkpointing.py
+++ b/tests/test_activation_checkpointing.py
@@ -908,16 +908,13 @@ def input_fn():
         f"{[n.name for n in sdpa_nodes]}"
     )
     fwd_sdpa, bwd_sdpa = sdpa_nodes
-    assert fwd_sdpa.meta["custom"] == {
-        "inside_checkpoint": 0,
-        "inside_local_map": 2,
-        "outside_checkpoint": 0,
-    }
-    assert bwd_sdpa.meta["custom"] == {
+    expected_custom = {
         "inside_checkpoint": 0,
         "inside_local_map": 2,
         "outside_checkpoint": 0,
     }
+    assert expected_custom.items() <= fwd_sdpa.meta["custom"].items()
+    assert expected_custom.items() <= bwd_sdpa.meta["custom"].items()
 
 
 # ---------------------------------------------------------------------------

From df79c65d9def50c95cbe5ef05122c242c06323ec Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Thu, 21 May 2026 06:01:50 +0000
Subject: [PATCH 05/12] Add graph pass to fuse chained allgathers on different
 subgroups into single full-mesh allgather
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When AutoParallel places weights as `S(0)S(0)` on a multi-dim mesh,
the `apply_sharding` pass decomposes the `S(0)S(0) → RR` redistribution
into per-dim allgathers: first a dp-dim allgather, then a tp-dim allgather.
In the backward graph, these appear as recomputed chains with cancelling
permute pairs between them:

```
dp_ag → wait → permute([1,0]) → permute([1,0]) → tp_ag → wait
```

Each pair produces two separate NCCL kernel launches when a single
full-mesh allgather would suffice.

This PR adds `fuse_chained_allgathers`, a graph pass that detects these
chains and replaces them with a single allgather on the flattened mesh
process group. The pass validates:
- Both allgathers are on known mesh subgroups (in descending dim order)
- Their group sizes multiply to the full mesh size
- The intermediate view ops compose to the identity (verified via
  FakeTensor shape/stride metadata, requiring at least one non-trivial
  dimension reorder like permute or transpose)
- No intermediate value has other consumers

The pass runs as a `pre_pass` on the partitioned forward and backward
graphs during both the first compilation (inside `AutoParallel`) and
the inference path. On LLaMA-3 8B (dp=16, tp=8, 32 layers), this fuses
49 allgather pairs in the backward, eliminating 49 standalone NCCL
kernel launches per iteration.

Authored with Claude.
---
 autoparallel/api.py                         |  21 +-
 autoparallel/graph_passes/fuse_allgather.py | 247 ++++++++++----
 tests/test_fuse_allgather.py                | 342 ++++++++++++++++++++
 3 files changed, 538 insertions(+), 72 deletions(-)
 create mode 100644 tests/test_fuse_allgather.py

diff --git a/autoparallel/api.py b/autoparallel/api.py
index 1b90218f..e2419963 100644
--- a/autoparallel/api.py
+++ b/autoparallel/api.py
@@ -27,7 +27,7 @@
 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_dp_tp_allgather
+from .graph_passes.fuse_allgather import fuse_chained_allgathers
 from .graph_passes.graph_utils import (
     _add_alias,
     _replace_view_mm_view_with_einsum,
@@ -510,8 +510,18 @@ def _make_fuse_allgather_pass(self):
         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)
+        }
+
         def pre_pass(graph):
-            fuse_dp_tp_allgather(graph, full_group_size, full_group_name)
+            fuse_chained_allgathers(
+                graph,
+                full_group_size,
+                full_group_name,
+                subgroup_order=subgroup_order,
+            )
 
         return pre_pass
 
@@ -555,6 +565,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/graph_passes/fuse_allgather.py b/autoparallel/graph_passes/fuse_allgather.py
index c4580c03..ebadba53 100644
--- a/autoparallel/graph_passes/fuse_allgather.py
+++ b/autoparallel/graph_passes/fuse_allgather.py
@@ -10,15 +10,6 @@
 logger: logging.Logger = logging.getLogger(__name__)
 
 
-def _is_permute_transpose(node: torch.fx.Node) -> bool:
-    return (
-        node.op == "call_function"
-        and node.target == torch.ops.aten.permute.default
-        and isinstance(node.args[1], (list, tuple))
-        and list(node.args[1]) == [1, 0]
-    )
-
-
 def _is_all_gather(node: torch.fx.Node) -> bool:
     return (
         node.op == "call_function"
@@ -33,124 +24,240 @@ def _is_wait_tensor(node: torch.fx.Node) -> bool:
     )
 
 
-def fuse_dp_tp_allgather(
+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,
 ) -> int:
-    """Fuse consecutive dp + tp allgather chains into a single full-mesh allgather.
+    """Fuse consecutive allgather chains on different subgroups into a single allgather.
 
-    Detects the pattern::
+    Detects chains of two allgathers on different process groups connected
+    through single-user view ops that compose to the identity::
 
-        dp_ag    = all_gather(x, dp_size, dp_pg)
-        dp_wait  = wait_tensor(dp_ag)
-        p1       = permute(dp_wait, [1, 0])
-        p2       = permute(p1, [1, 0])
-        tp_ag    = all_gather(p2, tp_size, tp_pg)
-        tp_wait  = wait_tensor(tp_ag)
+        ag1      = all_gather(x, size1, pg1)
+        wait1    = wait_tensor(ag1)
+        ...      = identity_view_ops(wait1)
+        ag2      = all_gather(..., size2, pg2)
+        wait2    = wait_tensor(ag2)
 
-    and replaces it with::
+    and replaces them with::
 
-        full_ag   = all_gather(x, dp_size * tp_size, full_pg)
+        full_ag   = all_gather(x, size1 * size2, full_pg)
         full_wait = wait_tensor(full_ag)
 
-    The two permutes must cancel (both ``[1, 0]`` transposes), the dp wait
-    must have a single user (the first permute), and the dp and tp group
-    sizes must multiply to ``full_group_size``.
+    Requirements:
+    - The two group sizes must multiply to ``full_group_size``.
+    - Every node between the two allgathers must have exactly one user.
+    - The view ops between them must compose to the identity (verified
+      via FakeTensor shape and stride metadata).
+    - Both allgathers must have the same dtype.
+    - When ``subgroup_order`` is provided, both process groups must be in
+      that mapping and appear in descending mesh-dim order.
 
     Returns the number of fusions performed.
     """
     fusions = 0
-
-    # Snapshot the node list — we mutate the graph during iteration.
     all_nodes = list(graph.nodes)
 
-    for tp_ag in all_nodes:
-        if not _is_all_gather(tp_ag):
+    for ag2 in all_nodes:
+        if not _is_all_gather(ag2):
             continue
 
-        # --- trace backwards through the expected chain ---
-
-        p2 = tp_ag.args[0]
-        if not isinstance(p2, torch.fx.Node) or not _is_permute_transpose(p2):
-            continue
-        if len(p2.users) != 1:
+        # Walk ag2's input backward through single-user nodes to find wait1.
+        node = ag2.args[0]
+        if not isinstance(node, torch.fx.Node):
             continue
 
-        p1 = p2.args[0]
-        if not isinstance(p1, torch.fx.Node) or not _is_permute_transpose(p1):
-            continue
-        if len(p1.users) != 1:
-            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
 
-        dp_wait = p1.args[0]
-        if not isinstance(dp_wait, torch.fx.Node) or not _is_wait_tensor(dp_wait):
+        if not _is_wait_tensor(wait1):
             continue
-        if len(dp_wait.users) != 1:
+        if len(wait1.users) != 1:
             continue
 
-        dp_ag = dp_wait.args[0]
-        if not isinstance(dp_ag, torch.fx.Node) or not _is_all_gather(dp_ag):
+        ag1 = wait1.args[0]
+        if not isinstance(ag1, torch.fx.Node) or not _is_all_gather(ag1):
             continue
-        if len(dp_ag.users) != 1:
+        if len(ag1.users) != 1:
             continue
 
-        # --- validate group sizes ---
-
-        dp_group_size = dp_ag.args[1]
-        tp_group_size = tp_ag.args[1]
-        if dp_group_size * tp_group_size != full_group_size:
+        # Validate that the view chain between wait1 and ag2 is identity.
+        if not _is_identity_view_chain(wait1, ag2):
             continue
 
-        # --- validate matching dtype ---
-
-        dp_val = dp_ag.meta.get("val")
-        tp_val = tp_ag.meta.get("val")
-        if dp_val is not None and tp_val is not None and dp_val.dtype != tp_val.dtype:
+        # 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
 
-        # --- find the tp wait (sole user of tp_ag) ---
+        # 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
+        if subgroup_order is not None:
+            if ag1_group not in subgroup_order or ag2_group not in subgroup_order:
+                continue
+            if subgroup_order[ag1_group] <= subgroup_order[ag2_group]:
+                continue
+
+        # 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
 
-        tp_wait = None
-        for user in tp_ag.users:
+        # Find wait2.
+        wait2 = None
+        for user in ag2.users:
             if _is_wait_tensor(user):
-                tp_wait = user
+                wait2 = user
                 break
-        if tp_wait is None:
+        if wait2 is None:
             continue
 
-        # --- build the fused allgather ---
-
-        original_input = dp_ag.args[0]
+        # Build the fused allgather.
+        original_input = ag1.args[0]
 
-        with graph.inserting_before(tp_ag):
+        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, full_group_name),
             )
-            full_ag.meta.update(tp_ag.meta)
+            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(tp_wait.meta)
+            full_wait.meta.update(wait2.meta)
 
-        tp_wait.replace_all_uses_with(full_wait)
+        wait2.replace_all_uses_with(full_wait)
         fusions += 1
 
         logger.debug(
-            "Fused dp_ag(%s, gs=%d) + tp_ag(gs=%d) -> full_ag(gs=%d)",
+            "Fused ag(%s, gs=%d, pg=%s) + ag(gs=%d, pg=%s) -> ag(gs=%d, pg=%s)",
             original_input,
-            dp_group_size,
-            tp_group_size,
+            ag1_group_size,
+            ag1_group,
+            ag2_group_size,
+            ag2_group,
             full_group_size,
+            full_group_name,
         )
 
     if fusions > 0:
         graph.eliminate_dead_code()
         logger.info(
-            "Fused %d dp+tp allgather chains into full-mesh allgathers", fusions
+            "Fused %d chained allgather pairs into full-mesh allgathers", fusions
         )
 
     return fusions
diff --git a/tests/test_fuse_allgather.py b/tests/test_fuse_allgather.py
new file mode 100644
index 00000000..90a26c0d
--- /dev/null
+++ b/tests/test_fuse_allgather.py
@@ -0,0 +1,342 @@
+# 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
+
+    # Correct subgroup order — should fuse
+    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
+
+
+def test_reversed_subgroup_order_does_not_fuse():
+    """Only descending mesh-dim allgather order is fuseable."""
+    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 == 0
+    assert _count_ops(graph, AG) == 2
+
+
+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 intervening views to reconcile the rank ordering between the
+    two subgroup allgathers and the flattened group, this 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)
+    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_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

From 6f9aa299e846e7d2752e68d4284861f2a856a101 Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Wed, 27 May 2026 13:53:59 +0000
Subject: [PATCH 06/12] Bypass aliases for consumers that would redistribute to
 the producer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When the ILP picks a placement for an `aten.alias.default` node that differs
from its producer's placement, any consumer whose `input_spec` matches the
producer's placement would redistribute the alias back to the producer's
placement at runtime. The original tensor must then stay alive longer than
necessary just to feed the redistribution.

The LLaMA-3 backward graph hits this on every transformer block: the
gradient at each residual add (`grad_h2`, S(0)S(1)) feeds an alias that the
optimizer assigns to S(0)R for two einsum consumers, but a third consumer
(the skip-add) wants S(0)S(1) — forcing a redistribution from R back to
S(1).

`eliminate_alias_round_trips` runs after `get_solution()` and rewires each
such consumer directly to the alias's producer. The alias keeps serving any
consumer that genuinely needs its placement; if no users remain, the alias
is erased from both the graph and the solution dict.

Unit tests in `tests/test_graph_utils.py` cover rewiring, alias erasure,
the no-op case (alias placement matches producer), intermediate
redistributions (consumer wants a third placement), and repeated inputs
(`x + x`-style consumers). On the LLaMA-3 8B example (32 layers, 128 GPUs),
the pass eliminates 32 round-trips. Authored with Claude.

Authored with Claude
---
 autoparallel/api.py                      |   5 +
 autoparallel/graph_passes/graph_utils.py |  53 ++++++
 tests/test_graph_utils.py                | 204 +++++++++++++++++++++++
 3 files changed, 262 insertions(+)

diff --git a/autoparallel/api.py b/autoparallel/api.py
index e2419963..b69706f0 100644
--- a/autoparallel/api.py
+++ b/autoparallel/api.py
@@ -33,6 +33,7 @@
     _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,
@@ -392,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))
 
diff --git a/autoparallel/graph_passes/graph_utils.py b/autoparallel/graph_passes/graph_utils.py
index b01afba1..bbe859ca 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_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

From fe18484892d7523505da407ce401cffb25b7b7ac Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Thu, 28 May 2026 07:55:22 +0000
Subject: [PATCH 07/12] Add module-level activation checkpointing from
 TorchTitan to LLaMA3 model

Port TorchTitan's apply_ac activation checkpointing implementation into AutoParallel's LLaMA3 model definition, replacing the graph-level AC pass with module-level checkpoint wrapping applied before tracing.

apply_ac(model, mode, selective_ac_option) supports three strategies matching TorchTitan: full AC (every layer), layer-selective (every nth layer), and per-op selective (saves compute-heavy ops like matmuls and SDPA but recomputes every 2nd matmul). The example uses op-level selective AC and disables the autoparallel backend's built-in AC pass to avoid double-checkpointing.

Authored with Claude.
---
 autoparallel/_testing/models/llama3.py | 111 ++++++++++++++++++++++++-
 examples/example_llama3.py             |  12 ++-
 2 files changed, 120 insertions(+), 3 deletions(-)

diff --git a/autoparallel/_testing/models/llama3.py b/autoparallel/_testing/models/llama3.py
index 9d349e1a..a7cf29f3 100644
--- a/autoparallel/_testing/models/llama3.py
+++ b/autoparallel/_testing/models/llama3.py
@@ -3,8 +3,9 @@
 # This source code is licensed under the BSD license found in the
 # LICENSE file in the root directory of this source tree.
 
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import ClassVar, Optional
+from typing import ClassVar, Literal, Optional
 
 import torch
 import torch.nn.functional as F
@@ -544,3 +545,111 @@ def forward(self, tokens: torch.Tensor, input_batch: Optional[torch.Tensor] = No
         h = self.norm(h) if self.norm else h
         output = self.output(h) if self.output else h
         return output
+
+
+# Ops whose outputs are saved (not recomputed) during selective AC.
+# Matches TorchTitan's _op_sac_save_list.
+_OP_SAC_SAVE_LIST: set = {
+    torch.ops.aten.mm.default,
+    torch.ops.aten._scaled_dot_product_efficient_attention.default,
+    torch.ops.aten._scaled_dot_product_flash_attention.default,
+    torch.ops.aten._scaled_dot_product_cudnn_attention.default,
+    torch.ops.aten._scaled_dot_product_attention_math.default,
+    torch.ops.aten._scaled_dot_product_fused_attention_overrideable.default,
+    torch.ops._c10d_functional.reduce_scatter_tensor.default,
+    torch.ops.aten.max.default,
+    torch._higher_order_ops.flex_attention,
+    torch._higher_order_ops.inductor_compiled_code,
+}
+
+# torch_attn._varlen_attn may not exist in all PyTorch builds
+if hasattr(torch.ops, "torch_attn") and hasattr(torch.ops.torch_attn, "_varlen_attn"):
+    _OP_SAC_SAVE_LIST.add(torch.ops.torch_attn._varlen_attn.default)
+
+
+def _apply_full_ac(module: nn.Module) -> nn.Module:
+    from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+        checkpoint_wrapper,
+    )
+
+    return checkpoint_wrapper(module)
+
+
+def _apply_layer_sac(module: nn.Module, layer_idx: int, ac_freq: int) -> nn.Module:
+    if not ac_freq or layer_idx % ac_freq == 0:
+        return _apply_full_ac(module)
+    return module
+
+
+def _apply_op_sac(module: nn.Module) -> nn.Module:
+    from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+        checkpoint_wrapper,
+    )
+    from torch.utils.checkpoint import (
+        CheckpointPolicy,
+        create_selective_checkpoint_contexts,
+    )
+
+    def _get_custom_policy(meta):
+        def _custom_policy(ctx, func, *args, **kwargs):
+            mode = "recompute" if ctx.is_recompute else "forward"
+            mm_count_key = f"{mode}_mm_count"
+            if func == torch.ops.aten.mm.default:
+                meta[mm_count_key] += 1
+            # Save output of all ops in save list, except every 2nd matmul
+            to_save = func in _OP_SAC_SAVE_LIST and not (
+                func == torch.ops.aten.mm.default and meta[mm_count_key] % 2 == 0
+            )
+            return (
+                CheckpointPolicy.MUST_SAVE
+                if to_save
+                else CheckpointPolicy.PREFER_RECOMPUTE
+            )
+
+        return _custom_policy
+
+    def selective_checkpointing_context_fn():
+        meta = defaultdict(int)
+        return create_selective_checkpoint_contexts(_get_custom_policy(meta))
+
+    return checkpoint_wrapper(module, context_fn=selective_checkpointing_context_fn)
+
+
+def apply_ac(
+    model: Transformer,
+    mode: Literal["full", "selective"] = "full",
+    selective_ac_option: str = "2",
+) -> None:
+    """Apply activation checkpointing to each TransformerBlock in the model.
+
+    Mirrors TorchTitan's activation checkpointing. Should be called on the
+    model before passing it to AutoParallel.
+
+    Args:
+        model: A Transformer model with a ``layers`` ModuleDict.
+        mode: ``"full"`` checkpoints every layer; ``"selective"`` uses
+            ``selective_ac_option`` to pick a strategy.
+        selective_ac_option: When mode is ``"selective"``, either a digit
+            string like ``"2"`` (checkpoint every nth layer) or ``"op"``
+            (per-op policy that recomputes every 2nd matmul).
+    """
+    for layer_id, transformer_block in model.layers.named_children():
+        if mode == "full":
+            transformer_block = _apply_full_ac(transformer_block)
+        elif mode == "selective":
+            if selective_ac_option == "op":
+                transformer_block = _apply_op_sac(transformer_block)
+            elif selective_ac_option.isdigit():
+                transformer_block = _apply_layer_sac(
+                    transformer_block, int(layer_id), int(selective_ac_option)
+                )
+            else:
+                raise ValueError(
+                    f"Invalid selective_ac_option: {selective_ac_option!r}. "
+                    "Use 'op' or a positive integer string like '2'."
+                )
+        else:
+            raise ValueError(
+                f"Invalid AC mode: {mode!r}. Valid modes: 'full', 'selective'."
+            )
+        model.layers.register_module(layer_id, transformer_block)
diff --git a/examples/example_llama3.py b/examples/example_llama3.py
index 3903ba8d..00025f5c 100644
--- a/examples/example_llama3.py
+++ b/examples/example_llama3.py
@@ -12,7 +12,11 @@
 from torch.distributed.tensor.placement_types import Partial, Replicate, Shard
 from torch.testing._internal.distributed.fake_pg import FakeStore
 
-from autoparallel._testing.models.llama3 import Transformer, TransformerModelArgs
+from autoparallel._testing.models.llama3 import (
+    Transformer,
+    TransformerModelArgs,
+    apply_ac,
+)
 from autoparallel.api import AutoParallel
 from autoparallel.compile import autoparallel_backend
 from autoparallel.cost_models.collective_runtime_estimation import set_nccl_topo_config
@@ -129,6 +133,8 @@ def post_grad_pass(graph):
 with torch.device("meta"):
     model = model_fn()
 
+apply_ac(model, mode="selective", selective_ac_option="op")
+
 mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
 
 
@@ -262,7 +268,9 @@ def _pass(graph):
 parallel_mod.init_weights()
 
 # Compile with AutoParallel-optimized Inductor passes
-parallel_mod = torch.compile(parallel_mod, backend=autoparallel_backend())
+parallel_mod = torch.compile(
+    parallel_mod, backend=autoparallel_backend(enable_ac=False)
+)
 
 # now let's run it
 x = (

From 428e9d2faf42cf6bdd92509219190d6e1b7b0883 Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Thu, 28 May 2026 12:45:53 +0000
Subject: [PATCH 08/12] Fix overlap scheduling config expiration and add
 compute batch capping
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three fixes for overlap scheduling and memory:

1. **Fix config.patch expiration for backward compilation** (`compile.py`): Overlap scheduling configs (`enable_overlap_scheduling`, `collective_bucketing`, etc.) were set via `config.patch` context manager inside `autoparallel_backend`. Since backward compilation is lazy (triggered on first `.backward()` after `compile_fx` returns), the context manager has already exited and the configs revert to defaults. Fix: set these configs globally instead. The `custom_partitioner_fn` stays in the context manager since it only affects forward partitioning.

2. **Add compute batch size capping** (`auto_bucketing.py`): After the overlap scheduler runs, AG bucketing can rewire dependencies such that `stable_topological_sort` batches hundreds of compute ops together (e.g., MM×40). This causes massive memory spikes. `_cap_compute_batch_size` adds ordering dependencies between compute nodes and their nearest reduce-scatter ops, limiting consecutive compute to `max_consecutive=8` ops. Also patches FSDP bucketing to use only the primary group and allow non-adjacent collectives, and increases `max_compute_pre_fetch` from 5 to 50.

3. **Save first residual add in SAC policy** (`llama3.py`): The selective AC policy now saves the first `aten.add.Tensor` output per transformer block (the attention residual `h = x + attn(...)`). Without this, the backward recomputes more ops between layers, giving the scheduler more latitude to prefetch and inflating peak memory. This reduced the post-scheduling memory gap vs reference from +3.30 GB to +0.39 GB.

The `example_llama3.py` changes are incidental (world_size, collective_bucketing flag, debug JSON dump).
---
 autoparallel/_testing/models/llama3.py      |   8 +
 autoparallel/compile.py                     |  20 +-
 autoparallel/graph_passes/auto_bucketing.py | 268 +++++++++++++++++++-
 3 files changed, 291 insertions(+), 5 deletions(-)

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/compile.py b/autoparallel/compile.py
index a40bbedc..eb603504 100644
--- a/autoparallel/compile.py
+++ b/autoparallel/compile.py
@@ -238,16 +238,28 @@ def autoparallel_backend(
         overlap_scheduling: Enable comm/compute overlap scheduling.
     """
     functorch_patches: dict[str, Any] = {}
-    inductor_patches: dict[str, Any] = (
-        dict(_INDUCTOR_OVERLAP_PATCHES) if overlap_scheduling else {}
-    )
 
     if enable_ac:
         functorch_patches["joint_custom_pass"] = _make_ac_joint_pass(
             ac_stage_size_in_GiB
         )
 
-    inductor_patches["custom_partitioner_fn"] = _SaveAllPartitioner()
+    # 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 (
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:

From e642341ce81a90378230864e78ecb78ae8e791c5 Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Thu, 28 May 2026 14:42:44 +0000
Subject: [PATCH 09/12] =?UTF-8?q?Fix=20fuse=5Fallgather=20pass=20to=20fire?=
 =?UTF-8?q?=20for=20dp=E2=86=92tp=20chains=20from=20reverse=20shard=20orde?=
 =?UTF-8?q?r?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 autoparallel/api.py                         | 10 +++
 autoparallel/graph_passes/fuse_allgather.py | 52 ++++++++----
 tests/test_fuse_allgather.py                | 94 +++++++++++++++++++--
 3 files changed, 134 insertions(+), 22 deletions(-)

diff --git a/autoparallel/api.py b/autoparallel/api.py
index b69706f0..54831148 100644
--- a/autoparallel/api.py
+++ b/autoparallel/api.py
@@ -520,12 +520,22 @@ def _make_fuse_allgather_pass(self):
             for dim in range(self.mesh.ndim)
         }
 
+        # Column-major flat mesh for ascending (dp→tp) chains from reverse
+        # shard order.  mesh["tp","dp"] transposes the dims so _flatten()
+        # produces ranks in column-major order.
+        reversed_full_group_name = None
+        if self.mesh.ndim == 2:
+            reversed_dim_names = tuple(reversed(self.mesh.mesh_dim_names))
+            reversed_flat_mesh = self.mesh[reversed_dim_names]._flatten()
+            reversed_full_group_name = reversed_flat_mesh.get_group().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
diff --git a/autoparallel/graph_passes/fuse_allgather.py b/autoparallel/graph_passes/fuse_allgather.py
index ebadba53..c47f9f3b 100644
--- a/autoparallel/graph_passes/fuse_allgather.py
+++ b/autoparallel/graph_passes/fuse_allgather.py
@@ -119,15 +119,17 @@ def fuse_chained_allgathers(
     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
-    through single-user view ops that compose to the identity::
+    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)
-        ...      = identity_view_ops(wait1)
+        ...      = [optional identity_view_ops(wait1)]
         ag2      = all_gather(..., size2, pg2)
         wait2    = wait_tensor(ag2)
 
@@ -136,14 +138,15 @@ def fuse_chained_allgathers(
         full_ag   = all_gather(x, size1 * size2, full_pg)
         full_wait = wait_tensor(full_ag)
 
-    Requirements:
-    - The two group sizes must multiply to ``full_group_size``.
-    - Every node between the two allgathers must have exactly one user.
-    - The view ops between them must compose to the identity (verified
-      via FakeTensor shape and stride metadata).
-    - Both allgathers must have the same dtype.
-    - When ``subgroup_order`` is provided, both process groups must be in
-      that mapping and appear in descending mesh-dim order.
+    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.
     """
@@ -182,8 +185,10 @@ def fuse_chained_allgathers(
         if len(ag1.users) != 1:
             continue
 
-        # Validate that the view chain between wait1 and ag2 is identity.
-        if not _is_identity_view_chain(wait1, ag2):
+        # 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.
@@ -199,11 +204,28 @@ def fuse_chained_allgathers(
         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
-            if subgroup_order[ag1_group] <= subgroup_order[ag2_group]:
+            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")
@@ -230,7 +252,7 @@ def fuse_chained_allgathers(
         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, full_group_name),
+                args=(original_input, full_group_size, target_group_name),
             )
             full_ag.meta.update(ag2.meta)
 
@@ -251,7 +273,7 @@ def fuse_chained_allgathers(
             ag2_group_size,
             ag2_group,
             full_group_size,
-            full_group_name,
+            target_group_name,
         )
 
     if fusions > 0:
diff --git a/tests/test_fuse_allgather.py b/tests/test_fuse_allgather.py
index 90a26c0d..93affae0 100644
--- a/tests/test_fuse_allgather.py
+++ b/tests/test_fuse_allgather.py
@@ -199,19 +199,24 @@ def test_subgroup_order_validation():
     assert fusions == 0
     assert _count_ops(graph, AG) == 2
 
-    # Correct subgroup order — should fuse
+    # 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_reversed_subgroup_order_does_not_fuse():
-    """Only descending mesh-dim allgather order is fuseable."""
+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")
@@ -224,8 +229,11 @@ def test_reversed_subgroup_order_does_not_fuse():
         subgroup_order={"dp": 0, "tp": 1},
     )
 
-    assert fusions == 0
-    assert _count_ops(graph, AG) == 2
+    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():
@@ -275,8 +283,7 @@ def test_standalone_allgather_untouched():
 def test_direct_chain_no_views():
     """Two allgathers directly chained (ag1 -> wait -> ag2) with no view ops.
 
-    Without intervening views to reconcile the rank ordering between the
-    two subgroup allgathers and the flattened group, this is NOT fuseable.
+    Without subgroup_order, this is NOT fuseable (can't validate direction).
     """
     graph = torch.fx.Graph()
     x = _add_placeholder(graph, "x", (8, 4096))
@@ -295,6 +302,79 @@ def test_direct_chain_no_views():
     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.
 

From ff446e283a400cd3ac1a1379b120d5b40edd53a2 Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Thu, 28 May 2026 15:47:54 +0000
Subject: [PATCH 10/12] =?UTF-8?q?Fix=20fuse=5Fallgather=20pass=20to=20fire?=
 =?UTF-8?q?=20for=20dp=E2=86=92tp=20chains=20from=20reverse=20shard=20orde?=
 =?UTF-8?q?r?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The `fuse_chained_allgathers` pass never fired for the unconstrained LLaMA-3
8B model despite 64 fusible dp→tp allgather chains in the backward graph
(2 per layer × 32 layers, from `S(0)S(0) → RS(0) → RR` weight unsharding).

Three independent issues prevented fusion:

1. The row-major flat mesh (`mesh._flatten()`) has the wrong rank ordering for
   dp→tp chains produced by reverse shard order. The fix creates a column-major
   process group via `dist.new_group(col_major_ranks, sort_ranks=False)` and
   uses it for ascending (dp→tp) chains, while the row-major mesh continues to
   serve descending (tp→dp) chains.

2. AOT autograd eliminates the canceling permutes between the two allgathers,
   leaving a direct chain (`ag1 → wait → ag2`) that `_is_identity_view_chain`
   rejected. Direct chains are now accepted when `subgroup_order` validates the
   direction and the matching process group is available.

3. The subgroup_order direction check (commit bf4c912 changed `<=` to `>=`)
   accepted ascending chains on the row-major flat mesh — the wrong
   combination. The check is replaced by explicit direction-to-group routing.

Authored with Claude.
---
 autoparallel/api.py          |  18 ++++--
 tests/test_fuse_allgather.py | 106 +++++++++++++++++++++++++++++++++++
 2 files changed, 118 insertions(+), 6 deletions(-)

diff --git a/autoparallel/api.py b/autoparallel/api.py
index 54831148..706159dc 100644
--- a/autoparallel/api.py
+++ b/autoparallel/api.py
@@ -520,14 +520,20 @@ def _make_fuse_allgather_pass(self):
             for dim in range(self.mesh.ndim)
         }
 
-        # Column-major flat mesh for ascending (dp→tp) chains from reverse
-        # shard order.  mesh["tp","dp"] transposes the dims so _flatten()
-        # produces ranks in column-major order.
+        # 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:
-            reversed_dim_names = tuple(reversed(self.mesh.mesh_dim_names))
-            reversed_flat_mesh = self.mesh[reversed_dim_names]._flatten()
-            reversed_full_group_name = reversed_flat_mesh.get_group().group_name
+            if not hasattr(self, "_reversed_flat_pg"):
+                import torch.distributed as dist
+
+                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(
diff --git a/tests/test_fuse_allgather.py b/tests/test_fuse_allgather.py
index 93affae0..ff3e73a7 100644
--- a/tests/test_fuse_allgather.py
+++ b/tests/test_fuse_allgather.py
@@ -420,3 +420,109 @@ def test_unsqueeze_squeeze_chain():
 
     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)

From f53aff5f58c8eb47f2d8c4dd5631edede94e060e Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Thu, 28 May 2026 16:00:25 +0000
Subject: [PATCH 11/12] Bugfix

---
 autoparallel/api.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/autoparallel/api.py b/autoparallel/api.py
index 706159dc..b46cba1b 100644
--- a/autoparallel/api.py
+++ b/autoparallel/api.py
@@ -528,11 +528,13 @@ def _make_fuse_allgather_pass(self):
         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
 
                 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
-                )
+                with unset_fake_temporarily():
+                    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):

From a22f87c767b1bb0f3f65fe7cc10039cd8969b467 Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Thu, 28 May 2026 16:12:31 +0000
Subject: [PATCH 12/12] Bugfix

---
 autoparallel/api.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/autoparallel/api.py b/autoparallel/api.py
index b46cba1b..cda994e9 100644
--- a/autoparallel/api.py
+++ b/autoparallel/api.py
@@ -530,8 +530,8 @@ def _make_fuse_allgather_pass(self):
                 import torch.distributed as dist
                 from torch._subclasses.fake_tensor import unset_fake_temporarily
 
-                col_major_ranks = self.mesh.mesh.T.contiguous().view(-1).tolist()
                 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
                     )