From 194b647cbdbca81f1d122e710f086826703a81ca Mon Sep 17 00:00:00 2001 From: Francisco Massa Date: Wed, 27 May 2026 13:53:59 +0000 Subject: [PATCH 1/2] 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 1670d509..3a938456 100644 --- a/autoparallel/api.py +++ b/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)) 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 c6c9a158d7b45a25f5920e243c7694e6b339b1d9 Mon Sep 17 00:00:00 2001 From: Francisco Massa Date: Wed, 27 May 2026 15:30:12 +0000 Subject: [PATCH 2/2] Fixes --- autoparallel/graph_passes/graph_utils.py | 10 +++-- tests/test_graph_utils.py | 47 +++++++++++++++++++++++- 2 files changed, 53 insertions(+), 4 deletions(-) diff --git a/autoparallel/graph_passes/graph_utils.py b/autoparallel/graph_passes/graph_utils.py index bbe859ca..400d686f 100644 --- a/autoparallel/graph_passes/graph_utils.py +++ b/autoparallel/graph_passes/graph_utils.py @@ -212,9 +212,12 @@ def eliminate_alias_round_trips(gm: torch.fx.GraphModule, solution: dict) -> int 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 - ] + # 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] @@ -232,6 +235,7 @@ def eliminate_alias_round_trips(gm: torch.fx.GraphModule, solution: dict) -> int changed = True if changed: + gm.graph.lint() gm.recompile() return eliminated diff --git a/tests/test_graph_utils.py b/tests/test_graph_utils.py index 7541919d..0f9968d1 100644 --- a/tests/test_graph_utils.py +++ b/tests/test_graph_utils.py @@ -346,7 +346,52 @@ def test_extract_forward_deepcopy_with_tensor_constants(): assert result._tensor_constant0.device.type == "cuda" -# --- eliminate_alias_round_trips tests --- +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):