From ed3da72af40467919fd367162acfbbc09c9a764e Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Sat, 7 Mar 2026 17:50:52 +0000
Subject: [PATCH 1/5] Model Comms / Compute Overlap in the ILP

---
 autoparallel/api.py               |   5 +
 autoparallel/optimize_sharding.py | 162 +++++++++++++++++++++++++++++-
 tests/test_optimize_placement.py  |  48 +++++++++
 3 files changed, 213 insertions(+), 2 deletions(-)

diff --git a/autoparallel/api.py b/autoparallel/api.py
index 4230d55c..0e410311 100644
--- a/autoparallel/api.py
+++ b/autoparallel/api.py
@@ -294,6 +294,7 @@ def __init__(
         reshard_after_forward: bool = True,
         dynamic: bool = False,
         numerics_logger: NumericsLogger | None = None,
+        enable_prefetch_overlap: bool = False,
         **kwargs,
     ):
         self.stack = ExitStack()
@@ -339,6 +340,7 @@ def __init__(
         self.enable_ac = enable_ac
         self.ac_stage_size_in_GiB = ac_stage_size_in_GiB
         self.reshard_after_forward = reshard_after_forward
+        self.enable_prefetch_overlap = enable_prefetch_overlap
 
         if dynamic:
             self.fake_mode.shape_env = ShapeEnv()
@@ -370,6 +372,7 @@ def __enter__(self):
             self.mesh,
             rescale_grad_comm_cost_for_mp,
             repeated_subgraphs=self.kwargs.get("repeated_subgraphs", False),
+            enable_prefetch_overlap=self.enable_prefetch_overlap,
         )
 
         self.sharding_optimizer = sharding_optimizer
@@ -853,6 +856,7 @@ def auto_parallel(
     mp_policy: Optional[MixedPrecisionPolicy] = None,
     compile: bool = True,
     parameter_memory_budget: Optional[tuple[Optional[float], Optional[float]]] = None,
+    enable_prefetch_overlap: bool = False,
 ) -> torch.nn.Module:
     """
     Parallelize a model with automatic sharding optimization.
@@ -932,6 +936,7 @@ def auto_parallel(
         compile=compile,
         # enable_ac=True,
         enable_ac=False,
+        enable_prefetch_overlap=enable_prefetch_overlap,
     ) as autop:
         # Add constraints
         autop.add_input_constraints(input_placements)
diff --git a/autoparallel/optimize_sharding.py b/autoparallel/optimize_sharding.py
index 3fa6e229..7afd4c01 100644
--- a/autoparallel/optimize_sharding.py
+++ b/autoparallel/optimize_sharding.py
@@ -130,7 +130,12 @@ def _assert_has_tensor_meta(spec_or_specs, node, label):
 
 class ShardingOptimizer:
     def __init__(
-        self, gm, mesh, rescale_grad_comm_cost_for_mp=1.0, repeated_subgraphs=False
+        self,
+        gm,
+        mesh,
+        rescale_grad_comm_cost_for_mp=1.0,
+        repeated_subgraphs=False,
+        enable_prefetch_overlap=False,
     ):
         self.gm = gm
         self.graph = gm.graph
@@ -151,7 +156,10 @@ def __init__(
         self.decision_vars = self._build_decision_vars()
         self.validate()
         self.prob = pulp.LpProblem("AutoParallel", pulp.LpMinimize)
+        self.savings_vars: list[pulp.LpVariable] = []
         self.add_default_constraints()
+        if enable_prefetch_overlap:
+            self.add_prefetch_overlap_constraints()
 
     def _get_next_name(self, prefix):
         idx = self._name_counters.setdefault(prefix, 0)
@@ -525,6 +533,153 @@ def add_default_constraints(self):
         self.add_inf_cost_constraint()
         self.add_forward_backward_consistency_constraints()
 
+    # ---- Prefetch overlap ----
+
+    def _build_param_derived_set(self):
+        """Compute the set of nodes whose inputs are ALL parameter-derived.
+
+        A node is parameter-derived if every one of its inputs is either a
+        parameter placeholder or itself parameter-derived. This propagates
+        through dtype_cast, views, aliases, etc.
+        """
+        param_derived = set(get_param_nodes(self.graph))
+        for node in self.graph.nodes:
+            all_inputs = self._all_input_nodes(node)
+            if all_inputs and all(inp in param_derived for inp in all_inputs):
+                param_derived.add(node)
+        return param_derived
+
+    def _build_terminal_derived_set(self):
+        """Compute the set of nodes where all downstream paths lead to output.
+
+        A node is terminal-derived if every user is either the output node or
+        itself terminal-derived. This propagates backward through alias chains,
+        dtype_cast_bwd nodes, etc.
+        """
+        terminal_derived: set[torch.fx.Node] = set()
+        for node in reversed(list(self.graph.nodes)):
+            if node.op == "output":
+                continue
+            if node.users and all(
+                u.op == "output" or u in terminal_derived for u in node.users
+            ):
+                terminal_derived.add(node)
+        return terminal_derived
+
+    def _comm_cost_expr_for_edge(self, node, argi):
+        """Build an LP expression for the selected comm cost on a given edge.
+
+        Returns Σ_{o,j} dv[node, argi, o, j].comm_cost * x[node, argi, o, j].
+        Skips infinite-cost entries (already forced to x=0 by add_inf_cost_constraint).
+        """
+        node_idx = self.node_map[node]
+        terms = []
+        for _, out_idx, inp_idx in self.walk_over_options(node, constrain_arg=argi):
+            key = (node_idx, argi, out_idx, inp_idx)
+            dv = self.decision_vars[key]
+            if dv.comm_cost != 0 and math.isfinite(dv.comm_cost):
+                terms.append(dv.comm_cost * dv.var)
+        return pulp.lpSum(terms)
+
+    def _compute_cost_expr_for_node(self, node):
+        """Build an LP expression for the selected full compute cost of a node.
+
+        Uses arg 0 per_arg_compute * num_args to recover the full compute cost,
+        since all args agree on the output strategy.
+        """
+        node_idx = self.node_map[node]
+        if node.op == "output" or node not in self.strats:
+            return pulp.lpSum([])
+        num_args = len(self.strats[node].strategies[0].input_specs)
+        terms = []
+        for _, out_idx, inp_idx in self.walk_over_options(node, constrain_arg=0):
+            key = (node_idx, 0, out_idx, inp_idx)
+            dv = self.decision_vars[key]
+            if dv.compute_cost != 0:
+                terms.append(num_args * dv.compute_cost * dv.var)
+        return pulp.lpSum(terms)
+
+    def add_prefetch_overlap_constraints(self):
+        """Model communication-computation overlap within the ILP.
+
+        Creates continuous savings variables that the solver maximizes (by
+        subtracting them from the objective). Each savings variable is bounded
+        by both the communication cost it overlaps and the available compute
+        budget from neighboring nodes, preventing double-counting.
+        """
+        param_derived = self._build_param_derived_set()
+        terminal_derived = self._build_terminal_derived_set()
+
+        # Maps each compute node to the list of savings variables that use its
+        # compute budget (from either scan direction).
+        compute_partners: dict[torch.fx.Node, list[pulp.LpVariable]] = defaultdict(list)
+
+        # --- Forward scan: prefetch overlap for parameter-derived inputs ---
+        current_group: list[torch.fx.Node] = []
+        for node in self.graph.nodes:
+            if node.op == "output":
+                continue
+
+            all_inputs = self._all_input_nodes(node)
+            param_derived_args = [
+                argi for argi, inp in enumerate(all_inputs) if inp in param_derived
+            ]
+
+            if param_derived_args and current_group:
+                for argi in param_derived_args:
+                    comm_expr = self._comm_cost_expr_for_edge(node, argi)
+                    savings = pulp.LpVariable(
+                        self._get_next_name("fwd_savings"),
+                        lowBound=0,
+                        cat=pulp.LpContinuous,
+                    )
+                    self.prob += (
+                        savings <= comm_expr,
+                        self._get_next_name("fwd_savings_le_comm"),
+                    )
+                    for partner in current_group:
+                        compute_partners[partner].append(savings)
+                    self.savings_vars.append(savings)
+                current_group = []
+
+            if node in self.strats and node.op != "output":
+                current_group.append(node)
+
+        # --- Reverse scan: post-compute overlap for terminal-derived nodes ---
+        current_group = []
+        for node in reversed(list(self.graph.nodes)):
+            if node.op == "output":
+                continue
+
+            if node in terminal_derived and current_group:
+                all_inputs = self._all_input_nodes(node)
+                for argi in range(len(all_inputs)):
+                    comm_expr = self._comm_cost_expr_for_edge(node, argi)
+                    savings = pulp.LpVariable(
+                        self._get_next_name("bwd_savings"),
+                        lowBound=0,
+                        cat=pulp.LpContinuous,
+                    )
+                    self.prob += (
+                        savings <= comm_expr,
+                        self._get_next_name("bwd_savings_le_comm"),
+                    )
+                    for partner in current_group:
+                        compute_partners[partner].append(savings)
+                    self.savings_vars.append(savings)
+                current_group = []
+
+            if node in self.strats and node.op != "output":
+                current_group.append(node)
+
+        # --- Compute budget constraints: prevent double-counting ---
+        for compute_node, savings_list in compute_partners.items():
+            compute_expr = self._compute_cost_expr_for_node(compute_node)
+            self.prob += (
+                pulp.lpSum(savings_list) <= compute_expr,
+                self._get_next_name("compute_budget"),
+            )
+
     # ---- Solution ----
 
     def _set_objective(self):
@@ -533,7 +688,10 @@ def _set_objective(self):
         cost_per_var = defaultdict(int)
         for dv in self.decision_vars.values():
             cost_per_var[dv.var] += dv.cost
-        self.prob += pulp.lpSum([var * cost for var, cost in cost_per_var.items()])
+        savings_sum = pulp.lpSum(self.savings_vars) if self.savings_vars else 0
+        self.prob += (
+            pulp.lpSum([var * cost for var, cost in cost_per_var.items()]) - savings_sum
+        )
 
     def _solve(self, verbose=False):
         solver = pulp.PULP_CBC_CMD(msg=verbose)
diff --git a/tests/test_optimize_placement.py b/tests/test_optimize_placement.py
index e185ad23..fd21898f 100644
--- a/tests/test_optimize_placement.py
+++ b/tests/test_optimize_placement.py
@@ -591,3 +591,51 @@ def input_fn():
     param_nodes = get_param_nodes(autop.gm.graph)
     for node in param_nodes:
         assert sharding_placement[node].output_specs.placements == (Replicate(),)
+
+
+def _run_ffn_optimization(mesh, enable_prefetch_overlap):
+    """Helper to run FFN optimization with or without prefetch overlap."""
+    model_fn, input_fn = _make_model_and_input_fn(mesh)
+    with torch.device("meta"):
+        model = model_fn()
+    with AutoParallel(
+        model, input_fn, mesh, enable_prefetch_overlap=enable_prefetch_overlap
+    ) as autop:
+        placement = (Shard(0), Replicate())
+        autop.add_input_constraints([placement] * 2)
+        autop.add_output_constraints([placement] * 3)
+        autop.add_parameter_memory_constraint(low=0, high=None)
+        sharding_placement = autop.optimize_placement(verbose=False)
+    return autop, sharding_placement
+
+
+@patch("torch.cuda.device_count", lambda: 8)
+@patch("torch.cuda.get_device_name", lambda device: "H100")
+def test_prefetch_overlap_reduces_cost(device_mesh_2d):
+    autop, _ = _run_ffn_optimization(device_mesh_2d, enable_prefetch_overlap=True)
+    opt = autop.sharding_optimizer
+
+    # The overlap-enabled objective subtracts savings from the base cost.
+    # Verify that savings are non-negative (the solver found overlap to exploit).
+    total_savings = sum(v.value() for v in opt.savings_vars)
+    assert total_savings >= -1e-9, f"Total savings should be >= 0, got {total_savings}"
+
+
+@patch("torch.cuda.device_count", lambda: 8)
+@patch("torch.cuda.get_device_name", lambda device: "H100")
+def test_prefetch_overlap_savings_structure(device_mesh_2d):
+    autop, _ = _run_ffn_optimization(device_mesh_2d, enable_prefetch_overlap=True)
+    opt = autop.sharding_optimizer
+
+    assert len(opt.savings_vars) > 0, "Expected savings variables to be created"
+
+    # Check that savings variables have names indicating both forward and
+    # backward scan directions
+    fwd_savings = [v for v in opt.savings_vars if v.name.startswith("fwd_savings")]
+    bwd_savings = [v for v in opt.savings_vars if v.name.startswith("bwd_savings")]
+    assert len(fwd_savings) > 0, "Expected forward-scan savings variables"
+    assert len(bwd_savings) > 0, "Expected backward-scan savings variables"
+
+    # All savings values should be >= 0
+    for v in opt.savings_vars:
+        assert v.value() >= -1e-9, f"Savings variable {v.name} has negative value"

From b08a67ce55d7f582aa5c653ef999dc624eb8991a Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Sun, 8 Mar 2026 12:18:59 +0000
Subject: [PATCH 2/5] Allow savings to draw from combined compute budget of
 multiple nodes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Previously each savings variable was added directly to every compute partner's budget constraint, which meant savings <= min(compute_cost(A), compute_cost(B)) — bounding it
by the smallest node in the group. This is too conservative: a 12-unit comm overlapping with two compute nodes (5 and 10) should allow up to 12 units of savings, not 5.

Fix by splitting each savings into per-node contribution variables (savings = contrib_A + contrib_B), where each contribution is non-negative and participates in its node's
budget constraint. The solver can now allocate e.g. 5 from A and 7 from B to fully hide 12 units of comm.

Authored with Claude.
---
 autoparallel/optimize_sharding.py | 68 +++++++++++++++++++------------
 1 file changed, 42 insertions(+), 26 deletions(-)

diff --git a/autoparallel/optimize_sharding.py b/autoparallel/optimize_sharding.py
index 7afd4c01..539aafae 100644
--- a/autoparallel/optimize_sharding.py
+++ b/autoparallel/optimize_sharding.py
@@ -599,6 +599,42 @@ def _compute_cost_expr_for_node(self, node):
                 terms.append(num_args * dv.compute_cost * dv.var)
         return pulp.lpSum(terms)
 
+    def _add_savings_var(self, node, argi, compute_group, compute_partners, prefix):
+        """Create a savings variable for overlapping comm on (node, argi) with
+        the compute budget of nodes in compute_group.
+
+        The savings is split into per-node contributions so that the total
+        savings can draw from the combined compute of the entire group, while
+        the per-node budget constraints still prevent double-counting.
+        """
+        comm_expr = self._comm_cost_expr_for_edge(node, argi)
+        savings = pulp.LpVariable(
+            self._get_next_name(f"{prefix}_savings"),
+            lowBound=0,
+            cat=pulp.LpContinuous,
+        )
+        self.prob += (
+            savings <= comm_expr,
+            self._get_next_name(f"{prefix}_savings_le_comm"),
+        )
+        # Split savings into per-node contributions that sum to savings.
+        # Each contribution is bounded by its node's compute budget via
+        # the aggregate constraint added later in add_prefetch_overlap_constraints.
+        contribs = []
+        for partner in compute_group:
+            contrib = pulp.LpVariable(
+                self._get_next_name(f"{prefix}_contrib"),
+                lowBound=0,
+                cat=pulp.LpContinuous,
+            )
+            contribs.append(contrib)
+            compute_partners[partner].append(contrib)
+        self.prob += (
+            savings == pulp.lpSum(contribs),
+            self._get_next_name(f"{prefix}_savings_split"),
+        )
+        self.savings_vars.append(savings)
+
     def add_prefetch_overlap_constraints(self):
         """Model communication-computation overlap within the ILP.
 
@@ -610,8 +646,8 @@ def add_prefetch_overlap_constraints(self):
         param_derived = self._build_param_derived_set()
         terminal_derived = self._build_terminal_derived_set()
 
-        # Maps each compute node to the list of savings variables that use its
-        # compute budget (from either scan direction).
+        # Maps each compute node to the list of contribution variables that
+        # draw from its compute budget (from either scan direction).
         compute_partners: dict[torch.fx.Node, list[pulp.LpVariable]] = defaultdict(list)
 
         # --- Forward scan: prefetch overlap for parameter-derived inputs ---
@@ -627,19 +663,9 @@ def add_prefetch_overlap_constraints(self):
 
             if param_derived_args and current_group:
                 for argi in param_derived_args:
-                    comm_expr = self._comm_cost_expr_for_edge(node, argi)
-                    savings = pulp.LpVariable(
-                        self._get_next_name("fwd_savings"),
-                        lowBound=0,
-                        cat=pulp.LpContinuous,
-                    )
-                    self.prob += (
-                        savings <= comm_expr,
-                        self._get_next_name("fwd_savings_le_comm"),
+                    self._add_savings_var(
+                        node, argi, current_group, compute_partners, "fwd"
                     )
-                    for partner in current_group:
-                        compute_partners[partner].append(savings)
-                    self.savings_vars.append(savings)
                 current_group = []
 
             if node in self.strats and node.op != "output":
@@ -654,19 +680,9 @@ def add_prefetch_overlap_constraints(self):
             if node in terminal_derived and current_group:
                 all_inputs = self._all_input_nodes(node)
                 for argi in range(len(all_inputs)):
-                    comm_expr = self._comm_cost_expr_for_edge(node, argi)
-                    savings = pulp.LpVariable(
-                        self._get_next_name("bwd_savings"),
-                        lowBound=0,
-                        cat=pulp.LpContinuous,
-                    )
-                    self.prob += (
-                        savings <= comm_expr,
-                        self._get_next_name("bwd_savings_le_comm"),
+                    self._add_savings_var(
+                        node, argi, current_group, compute_partners, "bwd"
                     )
-                    for partner in current_group:
-                        compute_partners[partner].append(savings)
-                    self.savings_vars.append(savings)
                 current_group = []
 
             if node in self.strats and node.op != "output":

From 0977c89fd572e78a695539b58ae9f0b60f290ff7 Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Sun, 8 Mar 2026 12:56:53 +0000
Subject: [PATCH 3/5] Refine overlap grouping, fix constraint validation
 tolerance, and add savings logging
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three improvements to the prefetch overlap model:

The forward scan now creates savings variables for all param-derived input edges (not just boundary edges), since the ILP may place the all-gather anywhere in the param
chain (e.g. param → dtype_cast). The compute group is only reset at boundary edges (non-param-derived consumer), so intermediate param-derived edges share the same compute
window and don't fragment the budget. The reverse scan applies symmetric boundary logic.

The violated-constraints logger now uses a 1e-6 tolerance, fixing false positives from floating-point residuals in the continuous savings/contribution equality constraints.

The cost summary now reports overlap_savings and effective_cost when prefetch overlap is enabled.

Authored with Claude.
---
 autoparallel/log_formatting.py    |  7 +++++++
 autoparallel/optimize_sharding.py | 35 ++++++++++++++++++++++++-------
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/autoparallel/log_formatting.py b/autoparallel/log_formatting.py
index 2bbea74a..8612b20d 100644
--- a/autoparallel/log_formatting.py
+++ b/autoparallel/log_formatting.py
@@ -22,6 +22,7 @@ def format_sharding_log(
     colored: bool = False,
     verbose: bool = False,
     violated_constraints_log: str = "",
+    savings_vars: list[Any] | None = None,
 ) -> str:
     """
     Format the sharding optimization results as annotated Python code.
@@ -38,6 +39,8 @@ def format_sharding_log(
         colored: Whether to use ANSI color codes in the output.
         verbose: Whether to include verbose information (shapes, stack traces).
         violated_constraints_log: Optional string with violated constraints info.
+        savings_vars: Optional list of PuLP savings variables from prefetch
+            overlap modeling.
 
     Returns:
         A string containing the annotated Python code representation of the graph.
@@ -205,6 +208,10 @@ def is_node_assignment_line(line: str, node_repr: str) -> bool:
     code += f"\n  comm_cost: {total_comm_cost:.2f}"
     code += f"\n  compute_cost: {total_compute_cost:.2f}"
     code += f"\n  transition_cost: {total_transition_cost:.2f}"
+    if savings_vars:
+        total_savings = sum(v.value() for v in savings_vars)
+        code += f"\n  overlap_savings: {total_savings:.2f}"
+        code += f"\n  effective_cost: {total_cost - total_savings:.2f}"
     if violated_constraints_log:
         code += "\n" + violated_constraints_log
     return code
diff --git a/autoparallel/optimize_sharding.py b/autoparallel/optimize_sharding.py
index 539aafae..25841e9c 100644
--- a/autoparallel/optimize_sharding.py
+++ b/autoparallel/optimize_sharding.py
@@ -651,6 +651,12 @@ def add_prefetch_overlap_constraints(self):
         compute_partners: dict[torch.fx.Node, list[pulp.LpVariable]] = defaultdict(list)
 
         # --- Forward scan: prefetch overlap for parameter-derived inputs ---
+        # Create savings for every edge with a param-derived input (the ILP
+        # may place the all-gather on any edge in the param chain). Only reset
+        # the compute group at boundary edges (non-param-derived node consuming
+        # a param-derived input), which marks the end of one prefetch window
+        # and the start of the next. Intermediate param-derived edges share
+        # the same compute group so the budget isn't fragmented.
         current_group: list[torch.fx.Node] = []
         for node in self.graph.nodes:
             if node.op == "output":
@@ -666,12 +672,18 @@ def add_prefetch_overlap_constraints(self):
                     self._add_savings_var(
                         node, argi, current_group, compute_partners, "fwd"
                     )
-                current_group = []
+                # Only reset at boundary edges — the prefetch window for the
+                # next layer starts after a non-param-derived consumer.
+                if node not in param_derived:
+                    current_group = []
 
             if node in self.strats and node.op != "output":
                 current_group.append(node)
 
         # --- Reverse scan: post-compute overlap for terminal-derived nodes ---
+        # Symmetric logic: create savings for every edge into a terminal-
+        # derived node from a non-terminal-derived input (where the reduce-
+        # scatter happens). Only reset at boundary edges.
         current_group = []
         for node in reversed(list(self.graph.nodes)):
             if node.op == "output":
@@ -679,11 +691,17 @@ def add_prefetch_overlap_constraints(self):
 
             if node in terminal_derived and current_group:
                 all_inputs = self._all_input_nodes(node)
-                for argi in range(len(all_inputs)):
-                    self._add_savings_var(
-                        node, argi, current_group, compute_partners, "bwd"
-                    )
-                current_group = []
+                boundary_args = [
+                    argi
+                    for argi, inp in enumerate(all_inputs)
+                    if inp not in terminal_derived
+                ]
+                if boundary_args:
+                    for argi in boundary_args:
+                        self._add_savings_var(
+                            node, argi, current_group, compute_partners, "bwd"
+                        )
+                    current_group = []
 
             if node in self.strats and node.op != "output":
                 current_group.append(node)
@@ -755,9 +773,9 @@ def get_solution(self, verbose=False):
 
     # ---- Logging ----
 
-    def get_violated_constraints_log(self):
+    def get_violated_constraints_log(self, eps=1e-6):
         violated_constraints = [
-            (k, c) for k, c in self.prob.constraints.items() if not c.valid()
+            (k, c) for k, c in self.prob.constraints.items() if not c.valid(eps)
         ]
         log_str = f"Violated constraints: {[x[0] for x in violated_constraints]}"
         for cname, c in violated_constraints:
@@ -780,6 +798,7 @@ def get_log(self, colored=False, verbose=False):
             colored=colored,
             verbose=verbose,
             violated_constraints_log=self.get_violated_constraints_log(),
+            savings_vars=self.savings_vars,
         )
 
     def print_costs_for_node(self, node, arg=0, **kwargs):

From 85e17c3bf65309af3c9f7e9a5cc30e85b3166d09 Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Sun, 8 Mar 2026 18:02:48 +0000
Subject: [PATCH 4/5] Replace group-based overlap model with cumulative budget
 chain
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The group-based overlap model splits each savings variable into per-node
contributions bounded by per-node compute budgets. When a savings variable
is created at a boundary edge, the compute group resets, so leftover
compute budget can't carry forward across windows.

This replaces it with a cumulative budget chain: a continuous LP variable
that grows with compute and shrinks with savings as we scan the graph.

```
B_0 = 0
B_i = B_{i-1} + compute(i) - savings(i)
B_i >= 0
```

The non-negativity of B_i implicitly prevents savings from exceeding
accumulated compute, and leftover budget naturally carries forward. This
is strictly more expressive — a comm preceded by compute windows of 3, 4,
5 can now draw from the full cumulative budget of 12 rather than only its
immediate group.

Also fixes a regression in the backward scan where a boundary_args
filter was too restrictive: in graphs where all nodes are terminal-derived,
it produced zero backward savings.

Authored with Claude.
---
 autoparallel/optimize_sharding.py | 165 ++++++++++++------------------
 1 file changed, 68 insertions(+), 97 deletions(-)

diff --git a/autoparallel/optimize_sharding.py b/autoparallel/optimize_sharding.py
index 25841e9c..dfee418e 100644
--- a/autoparallel/optimize_sharding.py
+++ b/autoparallel/optimize_sharding.py
@@ -599,120 +599,91 @@ def _compute_cost_expr_for_node(self, node):
                 terms.append(num_args * dv.compute_cost * dv.var)
         return pulp.lpSum(terms)
 
-    def _add_savings_var(self, node, argi, compute_group, compute_partners, prefix):
-        """Create a savings variable for overlapping comm on (node, argi) with
-        the compute budget of nodes in compute_group.
+    def _run_budget_chain(self, nodes, should_create_savings, prefix):
+        """Run a cumulative budget chain over nodes, creating savings variables.
 
-        The savings is split into per-node contributions so that the total
-        savings can draw from the combined compute of the entire group, while
-        the per-node budget constraints still prevent double-counting.
+        The budget is a continuous LP variable that grows with compute and
+        shrinks with savings as we scan:
+            B_0 = 0
+            B_i = B_{i-1} + compute(i) - savings(i)
+            B_i >= 0
+
+        The non-negativity constraint on B_i implicitly enforces that total
+        savings never exceed total accumulated compute.
+
+        Args:
+            nodes: Ordered sequence of nodes to scan.
+            should_create_savings: Function(node) -> list of arg indices that
+                need savings variables, or empty list if none.
+            prefix: Name prefix for LP variables ("fwd" or "bwd").
         """
-        comm_expr = self._comm_cost_expr_for_edge(node, argi)
-        savings = pulp.LpVariable(
-            self._get_next_name(f"{prefix}_savings"),
-            lowBound=0,
-            cat=pulp.LpContinuous,
-        )
-        self.prob += (
-            savings <= comm_expr,
-            self._get_next_name(f"{prefix}_savings_le_comm"),
-        )
-        # Split savings into per-node contributions that sum to savings.
-        # Each contribution is bounded by its node's compute budget via
-        # the aggregate constraint added later in add_prefetch_overlap_constraints.
-        contribs = []
-        for partner in compute_group:
-            contrib = pulp.LpVariable(
-                self._get_next_name(f"{prefix}_contrib"),
+        budget_prev = None
+        for node in nodes:
+            if node.op == "output":
+                continue
+
+            compute_expr = self._compute_cost_expr_for_node(node)
+            savings_args = should_create_savings(node)
+
+            node_savings = []
+            for argi in savings_args:
+                comm_expr = self._comm_cost_expr_for_edge(node, argi)
+                savings = pulp.LpVariable(
+                    self._get_next_name(f"{prefix}_savings"),
+                    lowBound=0,
+                    cat=pulp.LpContinuous,
+                )
+                self.prob += (
+                    savings <= comm_expr,
+                    self._get_next_name(f"{prefix}_savings_le_comm"),
+                )
+                node_savings.append(savings)
+                self.savings_vars.append(savings)
+
+            has_compute = node in self.strats and node.op != "output"
+            if not has_compute and not node_savings:
+                continue
+
+            budget = pulp.LpVariable(
+                self._get_next_name(f"{prefix}_budget"),
                 lowBound=0,
                 cat=pulp.LpContinuous,
             )
-            contribs.append(contrib)
-            compute_partners[partner].append(contrib)
-        self.prob += (
-            savings == pulp.lpSum(contribs),
-            self._get_next_name(f"{prefix}_savings_split"),
-        )
-        self.savings_vars.append(savings)
+            rhs = compute_expr - pulp.lpSum(node_savings)
+            if budget_prev is not None:
+                rhs += budget_prev
+            self.prob += (
+                budget == rhs,
+                self._get_next_name(f"{prefix}_budget_eq"),
+            )
+            budget_prev = budget
 
     def add_prefetch_overlap_constraints(self):
         """Model communication-computation overlap within the ILP.
 
-        Creates continuous savings variables that the solver maximizes (by
-        subtracting them from the objective). Each savings variable is bounded
-        by both the communication cost it overlaps and the available compute
-        budget from neighboring nodes, preventing double-counting.
+        Uses cumulative budget chains: a continuous LP variable that grows
+        with compute and shrinks with savings as we scan the graph. The
+        non-negativity of the budget variable naturally prevents savings
+        from exceeding available compute, and leftover compute carries
+        forward across multiple windows.
         """
         param_derived = self._build_param_derived_set()
         terminal_derived = self._build_terminal_derived_set()
 
-        # Maps each compute node to the list of contribution variables that
-        # draw from its compute budget (from either scan direction).
-        compute_partners: dict[torch.fx.Node, list[pulp.LpVariable]] = defaultdict(list)
-
         # --- Forward scan: prefetch overlap for parameter-derived inputs ---
-        # Create savings for every edge with a param-derived input (the ILP
-        # may place the all-gather on any edge in the param chain). Only reset
-        # the compute group at boundary edges (non-param-derived node consuming
-        # a param-derived input), which marks the end of one prefetch window
-        # and the start of the next. Intermediate param-derived edges share
-        # the same compute group so the budget isn't fragmented.
-        current_group: list[torch.fx.Node] = []
-        for node in self.graph.nodes:
-            if node.op == "output":
-                continue
-
+        def fwd_savings(node):
             all_inputs = self._all_input_nodes(node)
-            param_derived_args = [
-                argi for argi, inp in enumerate(all_inputs) if inp in param_derived
-            ]
+            return [argi for argi, inp in enumerate(all_inputs) if inp in param_derived]
 
-            if param_derived_args and current_group:
-                for argi in param_derived_args:
-                    self._add_savings_var(
-                        node, argi, current_group, compute_partners, "fwd"
-                    )
-                # Only reset at boundary edges — the prefetch window for the
-                # next layer starts after a non-param-derived consumer.
-                if node not in param_derived:
-                    current_group = []
-
-            if node in self.strats and node.op != "output":
-                current_group.append(node)
-
-        # --- Reverse scan: post-compute overlap for terminal-derived nodes ---
-        # Symmetric logic: create savings for every edge into a terminal-
-        # derived node from a non-terminal-derived input (where the reduce-
-        # scatter happens). Only reset at boundary edges.
-        current_group = []
-        for node in reversed(list(self.graph.nodes)):
-            if node.op == "output":
-                continue
+        self._run_budget_chain(self.graph.nodes, fwd_savings, "fwd")
 
-            if node in terminal_derived and current_group:
-                all_inputs = self._all_input_nodes(node)
-                boundary_args = [
-                    argi
-                    for argi, inp in enumerate(all_inputs)
-                    if inp not in terminal_derived
-                ]
-                if boundary_args:
-                    for argi in boundary_args:
-                        self._add_savings_var(
-                            node, argi, current_group, compute_partners, "bwd"
-                        )
-                    current_group = []
-
-            if node in self.strats and node.op != "output":
-                current_group.append(node)
+        # --- Backward scan: post-compute overlap for terminal-derived nodes ---
+        def bwd_savings(node):
+            if node not in terminal_derived:
+                return []
+            return list(range(len(self._all_input_nodes(node))))
 
-        # --- Compute budget constraints: prevent double-counting ---
-        for compute_node, savings_list in compute_partners.items():
-            compute_expr = self._compute_cost_expr_for_node(compute_node)
-            self.prob += (
-                pulp.lpSum(savings_list) <= compute_expr,
-                self._get_next_name("compute_budget"),
-            )
+        self._run_budget_chain(reversed(list(self.graph.nodes)), bwd_savings, "bwd")
 
     # ---- Solution ----
 
@@ -773,7 +744,7 @@ def get_solution(self, verbose=False):
 
     # ---- Logging ----
 
-    def get_violated_constraints_log(self, eps=1e-6):
+    def get_violated_constraints_log(self, eps=1e-4):
         violated_constraints = [
             (k, c) for k, c in self.prob.constraints.items() if not c.valid(eps)
         ]

From 1efb6d8bbe856bd80e88da5bf9190a342250fc44 Mon Sep 17 00:00:00 2001
From: Francisco Massa <fvsmassa@gmail.com>
Date: Tue, 24 Mar 2026 16:17:51 +0000
Subject: [PATCH 5/5] Split compute budget between forward and backward overlap
 chains

The two cumulative budget chains (forward for all-gather prefetch,
backward for reduce-scatter) were both counting every node's compute
independently, allowing the same compute to be "spent" twice. This meant
overlap savings could exceed total compute cost.

Fix by introducing a per-node continuous LP variable that splits each
node's compute between the two chains. The forward chain gets
compute_share, the backward gets compute_expr - compute_share, and
the solver decides the optimal allocation. This guarantees total savings
never exceed total compute while still allowing full flexibility in how
compute is distributed between the two overlap directions.

Authored with Claude.
---
 autoparallel/optimize_sharding.py | 46 ++++++++++++++++++++++++++++---
 examples/example_autoparallel.py  |  4 ++-
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/autoparallel/optimize_sharding.py b/autoparallel/optimize_sharding.py
index dfee418e..9777255e 100644
--- a/autoparallel/optimize_sharding.py
+++ b/autoparallel/optimize_sharding.py
@@ -599,7 +599,7 @@ def _compute_cost_expr_for_node(self, node):
                 terms.append(num_args * dv.compute_cost * dv.var)
         return pulp.lpSum(terms)
 
-    def _run_budget_chain(self, nodes, should_create_savings, prefix):
+    def _run_budget_chain(self, nodes, should_create_savings, prefix, compute_fn=None):
         """Run a cumulative budget chain over nodes, creating savings variables.
 
         The budget is a continuous LP variable that grows with compute and
@@ -616,13 +616,18 @@ def _run_budget_chain(self, nodes, should_create_savings, prefix):
             should_create_savings: Function(node) -> list of arg indices that
                 need savings variables, or empty list if none.
             prefix: Name prefix for LP variables ("fwd" or "bwd").
+            compute_fn: Optional function(node) -> LP expression for the
+                compute contribution. Defaults to _compute_cost_expr_for_node.
         """
+        if compute_fn is None:
+            compute_fn = self._compute_cost_expr_for_node
+
         budget_prev = None
         for node in nodes:
             if node.op == "output":
                 continue
 
-            compute_expr = self._compute_cost_expr_for_node(node)
+            compute_expr = compute_fn(node)
             savings_args = should_create_savings(node)
 
             node_savings = []
@@ -666,16 +671,42 @@ def add_prefetch_overlap_constraints(self):
         non-negativity of the budget variable naturally prevents savings
         from exceeding available compute, and leftover compute carries
         forward across multiple windows.
+
+        Each node's compute is split between forward and backward chains
+        via a continuous LP variable, so the solver decides the optimal
+        allocation and no compute is double-counted.
         """
         param_derived = self._build_param_derived_set()
         terminal_derived = self._build_terminal_derived_set()
 
+        # Split each node's compute between forward and backward chains.
+        # fwd_share is the portion allocated to the forward chain; the
+        # remainder goes to the backward chain.
+        fwd_share: dict[torch.fx.Node, pulp.LpVariable] = {}
+        for node in self.graph.nodes:
+            if node.op == "output" or node not in self.strats:
+                continue
+            compute_expr = self._compute_cost_expr_for_node(node)
+            share = pulp.LpVariable(
+                self._get_next_name("compute_share"),
+                lowBound=0,
+                cat=pulp.LpContinuous,
+            )
+            self.prob += (
+                share <= compute_expr,
+                self._get_next_name("compute_share_ub"),
+            )
+            fwd_share[node] = share
+
         # --- Forward scan: prefetch overlap for parameter-derived inputs ---
         def fwd_savings(node):
             all_inputs = self._all_input_nodes(node)
             return [argi for argi, inp in enumerate(all_inputs) if inp in param_derived]
 
-        self._run_budget_chain(self.graph.nodes, fwd_savings, "fwd")
+        def fwd_compute(node):
+            return fwd_share.get(node, pulp.lpSum([]))
+
+        self._run_budget_chain(self.graph.nodes, fwd_savings, "fwd", fwd_compute)
 
         # --- Backward scan: post-compute overlap for terminal-derived nodes ---
         def bwd_savings(node):
@@ -683,7 +714,14 @@ def bwd_savings(node):
                 return []
             return list(range(len(self._all_input_nodes(node))))
 
-        self._run_budget_chain(reversed(list(self.graph.nodes)), bwd_savings, "bwd")
+        def bwd_compute(node):
+            if node in fwd_share:
+                return self._compute_cost_expr_for_node(node) - fwd_share[node]
+            return self._compute_cost_expr_for_node(node)
+
+        self._run_budget_chain(
+            reversed(list(self.graph.nodes)), bwd_savings, "bwd", bwd_compute
+        )
 
     # ---- Solution ----
 
diff --git a/examples/example_autoparallel.py b/examples/example_autoparallel.py
index b2c7404f..50f2df39 100644
--- a/examples/example_autoparallel.py
+++ b/examples/example_autoparallel.py
@@ -119,7 +119,9 @@ def input_fn():
 mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
 # mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16)
 
-with AutoParallel(model, input_fn, mesh, mp_policy, compile=True) as autop:
+with AutoParallel(
+    model, input_fn, mesh, mp_policy, compile=True, enable_prefetch_overlap=True
+) as autop:
     assert any(n.meta.get("nn_module_stack") for n in autop.gm.graph.nodes)
     assert any(n.meta.get("fwd_nn_module_stack") for n in autop.gm.graph.nodes)
     autop.add_parameter_memory_constraint(low=None, high=None)