diff --git a/autoparallel/api.py b/autoparallel/api.py index 498c248f..8a9f2d01 100644 --- a/autoparallel/api.py +++ b/autoparallel/api.py @@ -169,6 +169,7 @@ def __init__( reshard_after_forward: bool = True, dynamic: bool = False, numerics_logger: NumericsLogger | None = None, + enable_prefetch_overlap: bool = False, cost_model: Any = None, repeated_subgraphs: bool = False, ): @@ -204,6 +205,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() @@ -258,6 +260,7 @@ def __enter__(self): self.mesh, rescale_grad_comm_cost_for_mp, repeated_subgraphs=self.repeated_subgraphs, + enable_prefetch_overlap=self.enable_prefetch_overlap, ) self.sharding_optimizer = sharding_optimizer @@ -500,6 +503,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. @@ -579,6 +583,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/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 edd4caca..f8c5cefb 100644 --- a/autoparallel/optimize_sharding.py +++ b/autoparallel/optimize_sharding.py @@ -133,7 +133,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 @@ -162,7 +167,10 @@ def __init__( self.validate() t2 = time.perf_counter() 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() t3 = time.perf_counter() n_unique_vars = len(set(id(v) for v in self.pulp_variables.values())) n_constraints = len(self.prob.constraints) @@ -632,6 +640,196 @@ 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 _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 + 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"). + 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 = compute_fn(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, + ) + 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. + + 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. + + 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] + + 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): + if node not in terminal_derived: + return [] + return list(range(len(self._all_input_nodes(node)))) + + 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 ---- def _set_objective(self): @@ -640,7 +838,8 @@ def _set_objective(self): for key, dv in self.decision_vars.items(): multiplier = 1 + len(self._root_to_linked.get(key, [])) terms.append(dv.var * dv.cost * multiplier) - self.prob += pulp.lpSum(terms) + savings_sum = pulp.lpSum(self.savings_vars) if self.savings_vars else 0 + self.prob += pulp.lpSum(terms) - savings_sum def _solve(self, verbose=False): solver = pulp.PULP_CBC_CMD(msg=verbose) @@ -697,9 +896,9 @@ def get_solution(self, verbose=False): # ---- Logging ---- - def get_violated_constraints_log(self): + 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() + (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: @@ -724,6 +923,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): diff --git a/examples/example_autoparallel.py b/examples/example_autoparallel.py index a4b4f50e..c4712715 100644 --- a/examples/example_autoparallel.py +++ b/examples/example_autoparallel.py @@ -122,7 +122,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: autop.add_parameter_memory_constraint(low=None, high=None) x_sharding = (Shard(0),) + (Replicate(),) * (mesh.ndim - 1) diff --git a/tests/test_optimize_placement.py b/tests/test_optimize_placement.py index cbc853b4..b2ec4b28 100644 --- a/tests/test_optimize_placement.py +++ b/tests/test_optimize_placement.py @@ -633,3 +633,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"