Bypass aliases for consumers that would redistribute to the producer

autoparallel/api.py

@@ -31,6 +31,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,
@@ -361,6 +362,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))
 

autoparallel/graph_passes/graph_utils.py

@@ -183,6 +183,63 @@ 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
+            # input_specs is indexed by position within the in-solution-filtered
+            # input list (matching apply_sharding._get_input_nodes and
+            # ShardingOptimizer._all_input_nodes); get_attr / shape-computation
+            # args are filtered out.
+            filtered_inputs = [n for n in all_input_nodes(consumer) if n in solution]
+            positions = [i for i, n in enumerate(filtered_inputs) 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.graph.lint()
+        gm.recompile()
+    return eliminated
+
+
 def is_collective(node: torch.fx.Node) -> bool:
     return (
         node.op == "call_function"

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,248 @@ 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"
+
+
+def test_eliminate_alias_round_trips_handles_mixed_arg_consumer(device_mesh_2d):
+    """Consumer with a non-sharded arg (e.g. get_attr) before the alias must
+    still index input_specs by the filtered (in-solution) position, matching
+    apply_sharding's _get_input_nodes convention.
+    """
+    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")
+    # get_attr is a non-sharded node placed BEFORE the alias in the consumer's
+    # raw arg list. apply_sharding filters it out, so input_specs only has one
+    # entry (for the alias).
+    gm._mod = torch.nn.Identity()
+    submod = g.get_attr("_mod")
+    consumer = g.call_function(
+        torch.ops.higher_order.invoke_subgraph, args=(submod, alias)
+    )
+    consumer.meta["val"] = torch.empty(4, device="meta")
+    g.output((consumer,))
+    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),),
+        ),
+        consumer: OpSpec(
+            _make_spec(mesh, x_pl, shape),
+            # only one entry — for the alias (the get_attr is filtered out)
+            input_specs=(_make_spec(mesh, x_pl, shape),),
+        ),
+    }
+
+    eliminated = eliminate_alias_round_trips(gm, sol)
+
+    assert eliminated == 1
+    # alias arg position in raw all_input_nodes is 1; in filtered it is 0.
+    assert list(consumer.all_input_nodes) == [submod, x]
+
+
+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