From 7d06b891f469387448d37ee0f32457100f44f496 Mon Sep 17 00:00:00 2001 From: Kaijian Wang Date: Sun, 5 Jul 2026 19:45:58 -0700 Subject: [PATCH] Speed up sharding optimizer build This reduces optimizer construction overhead while preserving the solver-visible problem. The build now skips discarded enumeration-time redistribute costs, prunes invalid edges before PuLP variable creation, stores repeated-subgraph cluster links at node granularity, and can construct edge costs in deterministic worker order. Authored with Claude. --- autoparallel/optimize_sharding.py | 424 ++++++++++++------ autoparallel/serialization.py | 41 +- .../tools/overlap_simulator/colls32_8.table | 14 +- docs/README.md | 2 + docs/fast_build.md | 40 ++ docs/fast_build_design.md | 43 ++ tests/test_export_json.py | 81 +--- tests/test_grad_reduce_dtype.py | 233 +++------- tests/test_optimize_placement.py | 58 +++ 9 files changed, 523 insertions(+), 413 deletions(-) create mode 100644 docs/fast_build.md create mode 100644 docs/fast_build_design.md diff --git a/autoparallel/optimize_sharding.py b/autoparallel/optimize_sharding.py index 6b72878b..f67e943a 100644 --- a/autoparallel/optimize_sharding.py +++ b/autoparallel/optimize_sharding.py @@ -41,10 +41,10 @@ → Implemented in: add_output_input_consistent_constraint() -4. COST CONSTRAINTS: Variables with infinite cost (invalid configurations) are - forced to zero. Efficiency penalties for inefficient collective operations - (e.g., non-batch-dim shard-to-replicate) are embedded in the cost coefficients - computed by the cost model (see cost_models/collective_runtime_estimation.py). +4. COST CONSTRAINTS: Invalid configurations are excluded. Efficiency penalties + for inefficient collective operations (e.g., non-batch-dim + shard-to-replicate) are embedded in the cost coefficients computed by the + cost model (see cost_models/collective_runtime_estimation.py). ∀i,a,o,j: c_{i,a,o,j} = ∞ ⟹ x_{i,a,o,j} = 0 @@ -69,9 +69,11 @@ runtime cost while satisfying all constraints. """ +import contextlib import logging import math import operator +import os import tempfile import time from collections import defaultdict @@ -108,6 +110,87 @@ logger = logging.getLogger(__name__) +def _fast_build_enabled(): + return os.environ.get("AP_FAST_BUILD", "1") != "0" + + +@contextlib.contextmanager +def _skip_enumeration_redistribute_cost(): + if not _fast_build_enabled(): + yield + return + + import torch.distributed.tensor._ops.utils as dt_utils + + orig = dt_utils.redistribute_cost + dt_utils.redistribute_cost = lambda *args, **kwargs: 0.0 + try: + yield + finally: + dt_utils.redistribute_cost = orig + + +def _parallel_build_workers(): + raw = os.environ.get("AP_PARALLEL_BUILD") + if raw is None: + workers = min(32, os.cpu_count() or 1) + else: + try: + workers = int(raw) + except ValueError as exc: + raise ValueError("AP_PARALLEL_BUILD must be an integer") from exc + if workers < 1: + raise ValueError("AP_PARALLEL_BUILD must be >= 1") + + if workers > 1 and torch.cuda.is_initialized(): + logger.info( + "AP_PARALLEL_BUILD requested %s workers, but CUDA is already " + "initialized; using serial optimizer build.", + workers, + ) + return 1 + return workers + + +_FORK_OPT: "ShardingOptimizer | None" = None + + +def _par_node_edge_costs(node_idx): + opt = _FORK_OPT + assert opt is not None + node = opt.nodes[node_idx] + op_strategy = opt.strats[node] + num_args = len(op_strategy.strategies[0].input_specs) + all_input_nodes = opt._all_input_nodes(node) + producer_strategies = [opt.strats[n] for n in all_input_nodes] + + out_data = [] + for output_strategy in op_strategy.strategies: + per_arg_compute = ( + estimate_strategy_runtime_cost(node, output_strategy) / num_args + ) + arg_rows = [] + for argi, redist_costs in enumerate(output_strategy.redistribute_cost): + producer_strategy = ( + producer_strategies[argi] if argi < len(producer_strategies) else None + ) + arg_rows.append( + [ + opt._compute_edge_costs( + node, + output_strategy, + argi, + inp_idx, + default_comm_cost, + producer_strategy, + ) + for inp_idx, default_comm_cost in enumerate(redist_costs) + ] + ) + out_data.append((per_arg_compute, arg_rows)) + return node_idx, out_data + + def concretize_symint(val): """Concretize a SymInt to a plain int, pass through other values. @@ -188,7 +271,7 @@ def concretize_gm(gm): return concrete_gm, orig_to_concrete, concrete_to_orig -@dataclass +@dataclass(slots=True) class DecisionVar: """A decision variable in the ILP, representing one (node, arg, output_placement, input_placement) choice with its associated costs and strategy metadata.""" @@ -246,6 +329,7 @@ def __init__( # remove_constraints can keep this in sync. self._node_constraint_names: dict[str, str] = {} self._name_counters: dict[str, int] = {} + self._valid_keys: set[tuple] | None = None t0 = time.perf_counter() self.strats = self.build_sharding_metadata() # nodes/node_map are derived from strats (not graph.nodes) so that @@ -258,7 +342,8 @@ def __init__( get_placement_options_timer().report() - self.cluster_links: dict[tuple, tuple] = {} + self.cluster_links: dict[int, int] = {} + self._root_to_copies: dict[int, list[int]] = defaultdict(list) if repeated_subgraphs: t = time.time() clusters = get_identical_regions(self.gm.graph, self.strats) @@ -305,58 +390,59 @@ def _normalize_node(self, node): def build_sharding_metadata(self): strats = {} - for node in self.graph.nodes: - if node.op in ("placeholder", "get_attr"): - val = node.meta.get("val") - if isinstance(val, torch.Tensor): - strats[node] = _create_all_options(self.mesh, val.shape, tensor=val) - elif node.op == "placeholder": - # Non-tensor placeholders (e.g. baked-in booleans/strings): - # keep them in strats with empty-shape replicate options - # so the constraint system can reference them. - strats[node] = _create_all_options(self.mesh, ()) + with _skip_enumeration_redistribute_cost(): + for node in self.graph.nodes: + if node.op in ("placeholder", "get_attr"): + val = node.meta.get("val") + if isinstance(val, torch.Tensor): + strats[node] = _create_all_options( + self.mesh, val.shape, tensor=val + ) + elif node.op == "placeholder": + # Non-tensor placeholders (e.g. baked-in booleans/strings): + # keep them in strats with empty-shape replicate options + # so the constraint system can reference them. + strats[node] = _create_all_options(self.mesh, ()) + else: + # Non-tensor get_attr: GraphModule submodules used by + # HOPs — not added to strats, invisible to the ILP. + # _all_input_nodes filters them. + assert node.op == "get_attr" + assert any( + isinstance(u.target, torch._ops.HigherOrderOperator) + or "local_map" in u.name + for u in node.users + ), f"Non-tensor get_attr {node} is not used by a HOP" + elif node.op == "call_function": + if not _produces_tensor(node.meta.get("val")): + # Shape-computation nodes (sym_size, operator.mul, etc.) + # produce scalars, not tensors — skip sharding. + continue + user_strats = tree_map_only( + torch.fx.Node, + lambda x: strats.get(x, x.meta.get("val")), + node.args, + ) + user_args = tree_map_only( + torch.fx.Node, lambda x: x.meta.get("val"), node.args + ) + user_kwargs = tree_map_only( + torch.fx.Node, lambda x: x.meta.get("val"), node.kwargs + ) + strats[node] = get_placement_options_for_node( + self.mesh, node, user_strats, user_args, user_kwargs + ) + elif node.op == "output": + user_strats = tree_map_only( + torch.fx.Node, lambda x: strats[x], node.args + ) + strats[node] = user_strats else: - # Non-tensor get_attr: GraphModule submodules used by - # HOPs — not added to strats, invisible to the ILP. - # _all_input_nodes filters them. - assert node.op == "get_attr" - assert any( - isinstance(u.target, torch._ops.HigherOrderOperator) - or "local_map" in u.name - for u in node.users - ), f"Non-tensor get_attr {node} is not used by a HOP" - elif node.op == "call_function": - if not _produces_tensor(node.meta.get("val")): - # Shape-computation nodes (sym_size, operator.mul, etc.) - # produce scalars, not tensors — skip sharding. - continue - user_strats = tree_map_only( - torch.fx.Node, - lambda x: strats.get(x, x.meta.get("val")), - node.args, - ) - user_args = tree_map_only( - torch.fx.Node, lambda x: x.meta.get("val"), node.args - ) - user_kwargs = tree_map_only( - torch.fx.Node, lambda x: x.meta.get("val"), node.kwargs - ) - strats[node] = get_placement_options_for_node( - self.mesh, node, user_strats, user_args, user_kwargs - ) - elif node.op == "output": - user_strats = tree_map_only( - torch.fx.Node, lambda x: strats[x], node.args - ) - strats[node] = user_strats - else: - raise ValueError(f"Unexpected node op: {node.op}") + raise ValueError(f"Unexpected node op: {node.op}") return strats def create_cluster_links(self, clusters): - """Create a mapping between identical optimization nodes to reduce the - optimization space. If cluster_links[key1] == key2, the optimization - problem uses key2's variable in place of key1.""" + """Create a node-level mapping between identical optimization nodes.""" for cluster_group in clusters: cluster0 = cluster_group[0] for cluster_i in cluster_group[1:]: @@ -369,13 +455,15 @@ def create_cluster_links(self, clusters): f"Problem with graph clustering: {n0} and {ni} don't have the same number " "of input/output placements. Please report a bug" ) - for argi, out_idx, inp_idx in options_n0: - self.cluster_links[(idx1, argi, out_idx, inp_idx)] = ( - idx0, - argi, - out_idx, - inp_idx, - ) + self.cluster_links[idx1] = idx0 + + def _cluster_root_key(self, key): + root_idx = self.cluster_links.get(key[0]) + return key if root_idx is None else (root_idx, key[1], key[2], key[3]) + + def _linked_option_keys(self, root_key): + for copy_idx in self._root_to_copies.get(root_key[0], ()): + yield (copy_idx, root_key[1], root_key[2], root_key[3]) def _all_input_nodes(self, node): """Variant of node.all_input_nodes that preserves duplicate nodes. @@ -417,7 +505,7 @@ def _create_pulp_variables(self): to their PuLP variables. Linked keys are not stored; use _get_pulp_variable() to resolve them through cluster_links. """ - cluster_linked_node_idxs = {key[0] for key in self.cluster_links} + cluster_linked_node_idxs = set(self.cluster_links) pulp_variables = {} for node, _ in self.strats.items(): @@ -428,6 +516,8 @@ def _create_pulp_variables(self): continue for argi, out_idx, inp_idx in self.walk_over_options(node): key = (node_idx, argi, out_idx, inp_idx) + if self._valid_keys is not None and key not in self._valid_keys: + continue root_node = self.nodes[node_idx] pulp_variables[key] = pulp.LpVariable( f"n={root_node},s={node_idx},arg={argi}," @@ -440,8 +530,8 @@ def _create_pulp_variables(self): def _get_pulp_variable(self, key): """Look up the PuLP variable for a key, resolving through cluster_links if the key belongs to a linked node.""" - root_key = self.cluster_links.get(key, key) - return self.pulp_variables[root_key] + root_key = self._cluster_root_key(key) + return self.pulp_variables.get(root_key) def _compute_edge_costs( self, @@ -482,16 +572,17 @@ def _build_decision_vars(self): """Build DecisionVar entries for every (node_idx, argi, out_idx, inp_idx) combination in the strategy space.""" t_pulp_start = time.perf_counter() - self.pulp_variables = self._create_pulp_variables() - t_pulp_end = time.perf_counter() + self.pulp_variables = {} + self._valid_keys = set() # Precompute which node indices are cluster-linked so we can # copy costs from the root instead of recomputing them. - self._cluster_linked_node_idxs = {key[0] for key in self.cluster_links} + self._cluster_linked_node_idxs = set(self.cluster_links) t_compute = 0.0 t_edge = 0.0 n_vars = 0 + n_pruned = 0 n_cluster_copied = 0 decision_vars = {} @@ -499,49 +590,43 @@ def _build_decision_vars(self): (self.node_map[node], node, strat) for node, strat in self.strats.items() ] - # Build DVs for root nodes only (not cluster-linked). - for node_idx, node, op_strategy in strats_items: - if node.op == "output": - continue - if node_idx in self._cluster_linked_node_idxs: - continue - - num_args = len(op_strategy.strategies[0].input_specs) - - for out_idx, output_strategy in enumerate(op_strategy.strategies): - tc0 = time.perf_counter() - compute_cost = estimate_strategy_runtime_cost(node, output_strategy) - tc1 = time.perf_counter() - t_compute += tc1 - tc0 - per_arg_compute = compute_cost / num_args + root_idxs = [ + node_idx + for node_idx, node, _ in strats_items + if node.op != "output" and node_idx not in self._cluster_linked_node_idxs + ] + tc0 = time.perf_counter() + node_results = self._compute_node_edge_costs(root_idxs) + t_edge = time.perf_counter() - tc0 + for node_idx, out_data in node_results: + node = self.nodes[node_idx] + op_strategy = self.strats[node] + for out_idx, (per_arg_compute, arg_rows) in enumerate(out_data): + output_strategy = op_strategy.strategies[out_idx] for argi, redist_costs in enumerate(output_strategy.redistribute_cost): - for inp_idx, default_comm_cost in enumerate(redist_costs): + for inp_idx, (comm_cost, transition_cost) in enumerate( + arg_rows[argi] + ): key = (node_idx, argi, out_idx, inp_idx) - - all_input_nodes = self._all_input_nodes(node) - producer_strategy = ( - self.strats[all_input_nodes[argi]] - if all_input_nodes - else None - ) - te0 = time.perf_counter() - comm_cost, transition_cost = self._compute_edge_costs( - node, - output_strategy, - argi, - inp_idx, - default_comm_cost, - producer_strategy, - ) - te1 = time.perf_counter() - t_edge += te1 - te0 - redist_costs[inp_idx] = comm_cost + cost = comm_cost + per_arg_compute + transition_cost + if not math.isfinite(cost): + n_pruned += 1 + continue + + var = pulp.LpVariable( + f"n={node},s={node_idx},arg={argi}," + f"output_p={out_idx},input_p={inp_idx}", + cat=pulp.LpBinary, + ) + self.pulp_variables[key] = var + self._valid_keys.add(key) + n_vars += 1 decision_vars[key] = DecisionVar( - var=self.pulp_variables[key], - cost=comm_cost + per_arg_compute + transition_cost, + var=var, + cost=cost, compute_cost=per_arg_compute, comm_cost=comm_cost, sharding_transition_cost=transition_cost, @@ -549,16 +634,12 @@ def _build_decision_vars(self): output_spec=output_strategy.output_specs, input_spec=output_strategy.input_specs[argi], ) - n_vars += 1 # Batch-copy redistribute_cost from root strats to linked strats. # The root pass above updated redistribute_cost in place with # edge-computed costs; linked strats need the same values for # _compute_solution_cost and other readers. - linked_node_to_root_node: dict[int, int] = {} - for linked_key, root_key in self.cluster_links.items(): - linked_node_to_root_node[linked_key[0]] = root_key[0] - for linked_node_idx, root_node_idx in linked_node_to_root_node.items(): + for linked_node_idx, root_node_idx in self.cluster_links.items(): linked_node = self.nodes[linked_node_idx] root_node = self.nodes[root_node_idx] linked_op = self.strats[linked_node] @@ -570,16 +651,17 @@ def _build_decision_vars(self): list(costs) for costs in root_spec.redistribute_cost ] n_cluster_copied = len(self.cluster_links) - n_vars += n_cluster_copied - self._root_to_linked: dict[tuple, list[tuple]] = defaultdict(list) - for linked_key, root_key in self.cluster_links.items(): - self._root_to_linked[root_key].append(linked_key) + self._root_to_copies = defaultdict(list) + for copy_idx, root_idx in self.cluster_links.items(): + self._root_to_copies[root_idx].append(copy_idx) + t_pulp_end = time.perf_counter() logger.debug( - "_build_decision_vars breakdown (%d vars, %d cluster-copied): " + "_build_decision_vars breakdown (%d vars, %d pruned-inf, %d cluster-copied): " "pulp_vars=%.3fs, compute_cost=%.3fs, edge_cost=%.3fs", n_vars, + n_pruned, n_cluster_copied, t_pulp_end - t_pulp_start, t_compute, @@ -587,12 +669,28 @@ def _build_decision_vars(self): ) return decision_vars + def _compute_node_edge_costs(self, root_idxs): + global _FORK_OPT + _FORK_OPT = self + try: + workers = _parallel_build_workers() + if workers <= 1 or len(root_idxs) < 64: + return [_par_node_edge_costs(ni) for ni in root_idxs] + + import multiprocessing as mp + + ctx = mp.get_context("fork") + with ctx.Pool(workers) as pool: + return list(pool.imap(_par_node_edge_costs, root_idxs, chunksize=4)) + finally: + _FORK_OPT = None + def _resolve_decision_var(self, key): """Return a DecisionVar for key, reconstructing on the fly for linked keys.""" dv = self.decision_vars.get(key) if dv is not None: return dv - root_key = self.cluster_links[key] + root_key = self._cluster_root_key(key) root_dv = self.decision_vars[root_key] node_idx, argi, out_idx, _ = key strategy = self.strats[self.nodes[node_idx]].strategies[out_idx] @@ -607,6 +705,15 @@ def _resolve_decision_var(self, key): input_spec=strategy.input_specs[argi], ) + def _find_decision_var(self, node_idx, argi, out_idx): + node = self.nodes[node_idx] + for _, _, inp_idx in self.walk_over_options(node, argi): + key = (node_idx, argi, out_idx, inp_idx) + root_key = self._cluster_root_key(key) + if root_key in self.decision_vars: + return self._resolve_decision_var(key) + return None + def _collect_vars(self, node, node_idx, argi, group_by, resolve_clusters=False): """Collect PuLP variables for a node's options, grouped by strategy index. @@ -619,12 +726,14 @@ def _collect_vars(self, node, node_idx, argi, group_by, resolve_clusters=False): result = {} for _, out_idx, inp_idx in self.walk_over_options(node, argi): key = (node_idx, argi, out_idx, inp_idx) - if key in self.cluster_links: + if key[0] in self.cluster_links: if not resolve_clusters: continue - var = self.pulp_variables[self.cluster_links[key]] + var = self._get_pulp_variable(key) else: - var = self.pulp_variables[key] + var = self.pulp_variables.get(key) + if var is None: + continue group_key = out_idx if group_by == "out_idx" else inp_idx result.setdefault(group_key, []).append(var) return result @@ -677,13 +786,16 @@ def add_unique_decision_constraint(self): if node_idx in self._cluster_linked_node_idxs: continue arg_vars = {} + num_args = len(self.strats[node].strategies[0].input_specs) for argi, out_idx, inp_idx in self.walk_over_options(node): key = (node_idx, argi, out_idx, inp_idx) - var = self.pulp_variables[key] + var = self.pulp_variables.get(key) + if var is None: + continue arg_vars.setdefault(argi, []).append(var) - for eqs in arg_vars.values(): + for argi in range(num_args): self.prob += ( - pulp.lpSum(eqs) == 1, + pulp.lpSum(arg_vars.get(argi, [])) == 1, self._get_next_name("unique_decision"), ) @@ -703,17 +815,24 @@ def add_same_output_across_args_constraint(self): continue if len(self._all_input_nodes(node)) <= 1: continue + num_args = len(self._all_input_nodes(node)) + num_outputs = len(self.strats[node].strategies) vars_per_output = {} for argi, out_idx, inp_idx in self.walk_over_options(node): key = (node_idx, argi, out_idx, inp_idx) - var = self.pulp_variables[key] + var = self.pulp_variables.get(key) + if var is None: + continue vars_per_output.setdefault((argi, out_idx), []).append(var) - eqs_per_arg = [[] for _ in self._all_input_nodes(node)] - for (argi, out_idx), value in vars_per_output.items(): - eqs_per_arg[argi].append(pulp.lpSum(value)) + eqs_per_arg = [ + [ + pulp.lpSum(vars_per_output.get((argi, out_idx), [])) + for out_idx in range(num_outputs) + ] + for argi in range(num_args) + ] arg0 = eqs_per_arg[0] for arg_eqs in eqs_per_arg[1:]: - assert len(arg0) == len(arg_eqs) for i in range(len(arg0)): self.prob += ( arg0[i] == arg_eqs[i], @@ -790,22 +909,15 @@ def add_output_input_consistent_constraint(self): ) continue - assert ( - vars_producer.keys() == vars_consumer.keys() - ), f"{vars_producer}, {vars_consumer}" - - for k in vars_producer: + for k in vars_producer.keys() | vars_consumer.keys(): self.prob += ( - pulp.lpSum(vars_producer[k]) == pulp.lpSum(vars_consumer[k]), + pulp.lpSum(vars_producer.get(k, [])) + == pulp.lpSum(vars_consumer.get(k, [])), self._get_next_name("output_input_consistent"), ) def add_inf_cost_constraint(self): - """COST (Category 4): Variables with infinite cost (invalid configurations) - are forced to zero. - - ∀i,a,o,j: c_{i,a,o,j} = ∞ ⟹ x_{i,a,o,j} = 0 - """ + """COST (Category 4): Invalid configurations are excluded.""" for key, dv in self.decision_vars.items(): if not math.isfinite(dv.cost): dv.cost = 10000.0 @@ -880,7 +992,7 @@ def _set_objective(self): """Add the cost minimization objective to the ILP.""" terms = [] for key, dv in self.decision_vars.items(): - multiplier = 1 + len(self._root_to_linked.get(key, [])) + multiplier = 1 + len(self._root_to_copies.get(key[0], ())) terms.append(dv.var * dv.cost * multiplier) self.prob += pulp.lpSum(terms) @@ -898,7 +1010,7 @@ def _solve(self, verbose=False): key for key, dv in self.decision_vars.items() if dv.var.value() == 1 ] for root_key in list(self.selected_keys): - self.selected_keys.extend(self._root_to_linked.get(root_key, [])) + self.selected_keys.extend(self._linked_option_keys(root_key)) if self.prob.status == -1: logger.warning(self.get_violated_constraints_log()) @@ -1073,7 +1185,9 @@ def _compute_solution_cost(self, solution): # Use pre-computed costs from decision vars instead of # estimate_strategy_runtime_cost, which needs node.meta["val"] # (absent on loaded optimizers). - dv = self._resolve_decision_var((node_idx, 0, out_idx, 0)) + dv = self._find_decision_var(node_idx, 0, out_idx) + if dv is None: + continue num_args = max(len(strategy.input_specs), 1) total_compute += dv.compute_cost * num_args @@ -1158,9 +1272,9 @@ def get_json(self): # Build node-level cluster mapping: linked_node -> root_node cluster_roots: dict[torch.fx.Node, torch.fx.Node] = {} - for linked_key, root_key in self.cluster_links.items(): - linked_node = self.nodes[linked_key[0]] - root_node = self.nodes[root_key[0]] + for linked_idx, root_idx in self.cluster_links.items(): + linked_node = self.nodes[linked_idx] + root_node = self.nodes[root_idx] cluster_roots[linked_node] = root_node _normalize_cluster_layer(cluster_roots) @@ -1404,11 +1518,13 @@ def _add_node_constraint( if constraint_name is None: constraint_name = "user_constraint" node_idx = self.node_map[node] - vars_per_arg = {} + num_args = len(self.strats[node].strategies[0].input_specs) + vars_per_arg = {argi: [] for argi in range(num_args)} for argi, out_idx, inp_idx in self.walk_over_options(node): if out_idx in output_constraint_indices: var = self._get_pulp_variable((node_idx, argi, out_idx, inp_idx)) - vars_per_arg.setdefault(argi, []).append(var) + if var is not None: + vars_per_arg[argi].append(var) names = [] for eqs in vars_per_arg.values(): name = self._get_next_name(constraint_name) @@ -1435,8 +1551,10 @@ def _add_paired_output_constraint(self, node_a, node_b, constraint_name): # This placement exists in node_a but not in node_b. # Disable it: force sum of its decision variables to 0. v_a = [ - self._get_pulp_variable((idx_a, 0, out_idx, inp_idx)) + var for inp_idx in range(num_inp_a) + if (var := self._get_pulp_variable((idx_a, 0, out_idx, inp_idx))) + is not None ] self.prob += ( pulp.lpSum(v_a) == 0, @@ -1445,12 +1563,16 @@ def _add_paired_output_constraint(self, node_a, node_b, constraint_name): continue out_idx_b = strat_b.index(sp) v_a = [ - self._get_pulp_variable((idx_a, 0, out_idx, inp_idx)) + var for inp_idx in range(num_inp_a) + if (var := self._get_pulp_variable((idx_a, 0, out_idx, inp_idx))) + is not None ] v_b = [ - self._get_pulp_variable((idx_b, 0, out_idx_b, inp_idx)) + var for inp_idx in range(num_inp_b) + if (var := self._get_pulp_variable((idx_b, 0, out_idx_b, inp_idx))) + is not None ] self.prob += ( pulp.lpSum(v_b) == pulp.lpSum(v_a), @@ -1653,7 +1775,9 @@ def _apply_memory_constraint(self): num_out_strat = len(self.strats[node].strategies) ratios: list[float] = [] for out_idx in range(num_out_strat): - dv = self._resolve_decision_var((node_idx, 0, out_idx, 0)) + dv = self._find_decision_var(node_idx, 0, out_idx) + if dv is None: + continue spec: DTensorSpec = dv.input_spec assert spec.tensor_meta is not None tensor_shape: torch.Size = spec.tensor_meta.shape @@ -1663,6 +1787,8 @@ def _apply_memory_constraint(self): ratio = new_size / old_size ratios.append(ratio) elms.append(dv.var * ratio) + if not ratios: + continue best_ratio: float = min(ratios) budget_low += max(best_ratio, memory_factor_low) budget_high += max(best_ratio, memory_factor_high) diff --git a/autoparallel/serialization.py b/autoparallel/serialization.py index 46cb3fde..3ae3d789 100644 --- a/autoparallel/serialization.py +++ b/autoparallel/serialization.py @@ -135,7 +135,7 @@ def save_optimizer(opt, path): # Re-key strats by node name, saving only root nodes (non-linked). # Linked nodes share identical strats with their root and are # reconstructed on load from cluster_links. - linked_node_names = {opt.nodes[lk[0]].name for lk in opt.cluster_links} + linked_node_names = {opt.nodes[linked_idx].name for linked_idx in opt.cluster_links} strats_by_name = { node.name: strat for node, strat in opt.strats.items() @@ -193,8 +193,8 @@ def save_optimizer(opt, path): "dv_costs_keys": dv_costs_keys, "dv_costs_vals": dv_costs_vals, "cluster_links_node_by_name": { - opt.nodes[lk[0]].name: opt.nodes[rk[0]].name - for lk, rk in opt.cluster_links.items() + opt.nodes[linked_idx].name: opt.nodes[root_idx].name + for linked_idx, root_idx in opt.cluster_links.items() }, "constraint_log": opt._constraint_log, "selected_keys_by_name": selected_keys_by_name, @@ -265,22 +265,15 @@ def load_optimizer(cls, path): opt._node_constraint_names = {} opt._name_counters = {} - # Reconstruct cluster_links by expanding the node-level mapping over - # all (argi, out_idx, inp_idx) combinations. + # Reconstruct node-level cluster links. opt.cluster_links = {} for linked_name, root_name in cluster_links_node_by_name.items(): linked_node = nodes_by_name[linked_name] root_node = nodes_by_name[root_name] linked_idx = opt.node_map[linked_node] root_idx = opt.node_map[root_node] - for argi, out_idx, inp_idx in opt.walk_over_options(linked_node): - opt.cluster_links[(linked_idx, argi, out_idx, inp_idx)] = ( - root_idx, - argi, - out_idx, - inp_idx, - ) - opt._cluster_linked_node_idxs = {key[0] for key in opt.cluster_links} + opt.cluster_links[linked_idx] = root_idx + opt._cluster_linked_node_idxs = set(opt.cluster_links) # Mesh placeholder — provides shape/dim_names for get_json() and ndim # for add_node_constraint() default placement, without needing a PG @@ -288,6 +281,15 @@ def load_optimizer(cls, path): # Rebuild PuLP variables and decision vars from saved costs. t2 = time.perf_counter() + save_node_names = save_dict["dv_costs_node_names"] + keys_t = save_dict["dv_costs_keys"].tolist() + vals_t = save_dict["dv_costs_vals"].tolist() + mapped_keys = set() + for save_node_idx, argi, out_idx, inp_idx in keys_t: + node_name = save_node_names[save_node_idx] + node = nodes_by_name[node_name] + mapped_keys.add((opt.node_map[node], argi, out_idx, inp_idx)) + opt._valid_keys = mapped_keys opt.pulp_variables = opt._create_pulp_variables() t3 = time.perf_counter() logger.debug( @@ -296,9 +298,6 @@ def load_optimizer(cls, path): len(opt.pulp_variables), ) # Reconstruct decision_vars from compact tensors. - save_node_names = save_dict["dv_costs_node_names"] - keys_t = save_dict["dv_costs_keys"].tolist() - vals_t = save_dict["dv_costs_vals"].tolist() opt.decision_vars = {} for (save_node_idx, argi, out_idx, inp_idx), ( compute_cost, @@ -329,9 +328,9 @@ def load_optimizer(cls, path): len(opt.decision_vars), ) - opt._root_to_linked = defaultdict(list) - for linked_key, root_key in opt.cluster_links.items(): - opt._root_to_linked[root_key].append(linked_key) + opt._root_to_copies = defaultdict(list) + for copy_idx, root_idx in opt.cluster_links.items(): + opt._root_to_copies[root_idx].append(copy_idx) opt.prob = pulp.LpProblem("AutoParallel", pulp.LpMinimize) opt.add_default_constraints() @@ -382,9 +381,9 @@ def _restore_solution(opt, selected_keys_by_name, nodes_by_name): if dv is not None: dv.var.varValue = 1.0 - # Expand cluster links + # Expand cluster copies. for root_key in list(opt.selected_keys): - opt.selected_keys.extend(opt._root_to_linked.get(root_key, [])) + opt.selected_keys.extend(opt._linked_option_keys(root_key)) def save_placements(opt, path): diff --git a/autoparallel/tools/overlap_simulator/colls32_8.table b/autoparallel/tools/overlap_simulator/colls32_8.table index 50426ef2..82001356 100644 --- a/autoparallel/tools/overlap_simulator/colls32_8.table +++ b/autoparallel/tools/overlap_simulator/colls32_8.table @@ -1,8 +1,8 @@ - Group Group Size Collective 1MB (ms) 2MB (ms) 4MB (ms) 8MB (ms) 16MB (ms) 32MB (ms) 64MB (ms) 128MB (ms) 256MB (ms) 512MB (ms) 1024MB (ms) 2048MB (ms) + Group Group Size Collective 1MB (ms) 2MB (ms) 4MB (ms) 8MB (ms) 16MB (ms) 32MB (ms) 64MB (ms) 128MB (ms) 256MB (ms) 512MB (ms) 1024MB (ms) 2048MB (ms) ------- ------------ -------------------------- ---------- ---------- ---------- ---------- ----------- ----------- ----------- ------------ ------------ ------------ ------------- ------------- - 1 8 all_gather_into_tensor 0.0495 0.0716 0.1138 0.1953 0.3584 0.6846 1.3371 2.642 5.2518 10.4714 20.9105 41.7888 - 1 8 reduce_scatter_tensor 0.0173 0.0238 0.0368 0.0495 0.0716 0.1138 0.1953 0.3584 0.6846 1.3371 2.642 5.2518 - 1 8 all_reduce 0.028 0.041 0.0628 0.0849 0.1292 0.2179 0.3822 0.7084 1.3609 2.6658 5.2756 10.4952 - 0 32 all_gather_into_tensor 1.0136 1.7497 3.1512 5.86 11.2777 22.113 43.7835 87.1247 173.807 347.171 693.901 1387.36 - 0 32 reduce_scatter_tensor 0.2114 0.2612 0.3608 0.4615 0.6455 1.0136 1.7497 3.1512 5.86 11.2777 22.113 43.7835 - 0 32 all_gather_into_tensor_out 1.0136 1.7497 3.1512 5.86 11.2777 22.113 43.7835 87.1247 173.807 347.171 693.901 1387.36 + 1 8 all_gather_into_tensor 0.0495 0.0716 0.1138 0.1953 0.3584 0.6846 1.3371 2.642 5.2518 10.4714 20.9105 41.7888 + 1 8 reduce_scatter_tensor 0.0173 0.0238 0.0368 0.0495 0.0716 0.1138 0.1953 0.3584 0.6846 1.3371 2.642 5.2518 + 1 8 all_reduce 0.028 0.041 0.0628 0.0849 0.1292 0.2179 0.3822 0.7084 1.3609 2.6658 5.2756 10.4952 + 0 32 all_gather_into_tensor 1.0136 1.7497 3.1512 5.86 11.2777 22.113 43.7835 87.1247 173.807 347.171 693.901 1387.36 + 0 32 reduce_scatter_tensor 0.2114 0.2612 0.3608 0.4615 0.6455 1.0136 1.7497 3.1512 5.86 11.2777 22.113 43.7835 + 0 32 all_gather_into_tensor_out 1.0136 1.7497 3.1512 5.86 11.2777 22.113 43.7835 87.1247 173.807 347.171 693.901 1387.36 diff --git a/docs/README.md b/docs/README.md index 9299286f..c685dd1d 100644 --- a/docs/README.md +++ b/docs/README.md @@ -25,3 +25,5 @@ If you're new to the project, use the reading order below. - [Using `local_map` for MoE and Custom Communication Patterns](local_map_and_moe.md) - [Saving and Loading Optimizer State](save_load.md) +- [Fast Optimizer Build](fast_build.md) +- [Fast Optimizer Build Design](fast_build_design.md) diff --git a/docs/fast_build.md b/docs/fast_build.md new file mode 100644 index 00000000..6af142e1 --- /dev/null +++ b/docs/fast_build.md @@ -0,0 +1,40 @@ +# Fast Optimizer Build + +AutoParallel builds an optimization problem before solving for placements. Fast +optimizer build reduces the time and memory needed for that construction while +preserving the same placement result. + +Fast build is enabled by default. + +## Controls + +`AP_FAST_BUILD=0` disables the enumeration-time shortcut used during optimizer +construction. This is useful for A/B validation. + +```bash +AP_FAST_BUILD=0 python examples/example_autoparallel.py +``` + +`AP_PARALLEL_BUILD` controls how many worker processes are used for optimizer +edge-cost construction. + +```bash +AP_PARALLEL_BUILD=1 python examples/example_autoparallel.py +AP_PARALLEL_BUILD=8 python examples/example_autoparallel.py +``` + +Use `AP_PARALLEL_BUILD=1` to force serial construction. AutoParallel also falls +back to serial construction when CUDA is already initialized, because forking +after CUDA initialization is not safe in real GPU training processes. + +## Validation + +To compare fast build with the reference construction path, run the same model +twice and compare the returned placements: + +```bash +AP_FAST_BUILD=0 AP_PARALLEL_BUILD=1 python your_autoparallel_script.py +AP_FAST_BUILD=1 AP_PARALLEL_BUILD=1 python your_autoparallel_script.py +``` + +The placements should match. If they do not, treat it as a correctness bug. diff --git a/docs/fast_build_design.md b/docs/fast_build_design.md new file mode 100644 index 00000000..fe92a8c3 --- /dev/null +++ b/docs/fast_build_design.md @@ -0,0 +1,43 @@ +# Fast Optimizer Build Design + +Fast optimizer build changes how AutoParallel constructs the sharding ILP. It +does not change the objective, the feasible placements, or the selected +placement result. + +## Equivalence + +Enumeration-time redistribute costs are skipped because the optimizer later +computes and stores AutoParallel's own communication costs for every surviving +strategy edge. No solver-visible cost uses the skipped values. + +Infinite-cost strategy edges are omitted before PuLP variable creation. This is +equivalent to the previous formulation that created those variables and added +constraints forcing them to zero. + +Repeated-subgraph cluster links are stored per node instead of per option. The +option indices are identical between a repeated node and its cluster root, so the +same root variable can be resolved from the node-level link. + +`DecisionVar` uses `slots=True` to reduce Python object overhead. This changes +memory layout only. + +Edge-cost computation can run in forked workers. Results are consumed in the +same node order as the serial path, and PuLP variables are still assembled in +the parent process. + +## Safety + +`AP_PARALLEL_BUILD=1` forces the serial path. + +When CUDA is already initialized, AutoParallel uses the serial path even if +`AP_PARALLEL_BUILD` requests multiple workers. This avoids fork-after-CUDA +failures in real training runs. + +`AP_FAST_BUILD=0` disables the enumeration-time shortcut for A/B validation. + +## Test Strategy + +The fast-build regression test compares public optimizer behavior between the +reference and fast paths: the optimizer must solve successfully and return the +same placements. It does not rely on private optimizer data structures such as +PuLP variables or decision-variable key sets. diff --git a/tests/test_export_json.py b/tests/test_export_json.py index 9ac9ac0a..f202daab 100644 --- a/tests/test_export_json.py +++ b/tests/test_export_json.py @@ -9,7 +9,6 @@ from torch.distributed.tensor.placement_types import Shard from autoparallel.api import AutoParallel -from autoparallel.export_json import _get_layer_index, _normalize_cluster_layer class _RepeatedLayerModel(nn.Module): @@ -37,53 +36,8 @@ def input_fn(): return AutoParallel(model, input_fn, device_mesh, repeated_subgraphs=True) -# ---- _get_layer_index tests ---- - - @apply_cuda_patches -def test_get_layer_index_from_nn_module_stack(device_mesh_1d): - dim = 64 - with torch.device("meta"): - model = _RepeatedLayerModel(dim, n_layers=3) - - autop = _setup_autop(model, dim, device_mesh_1d) - with autop: - autop.add_input_constraints([(Shard(0),)]) - autop.add_output_constraints([(Shard(0),)]) - graph = autop.sharding_optimizer.graph - - found_layers = set() - found_non_layer = False - for node in graph.nodes: - idx = _get_layer_index(node) - if idx is not None: - found_layers.add(idx) - elif node.op not in ("placeholder", "output"): - found_non_layer = True - - assert found_layers == {0, 1, 2}, f"Expected layers 0-2, got {found_layers}" - assert found_non_layer, "Expected some non-layer nodes (embed, head)" - - -def test_get_layer_index_returns_none_for_non_layer_node(): - """Nodes without nn_module_stack should return None.""" - - class FakeNode: - def __init__(self): - self.meta = {} - - assert _get_layer_index(FakeNode()) is None - - node_with_empty_stack = FakeNode() - node_with_empty_stack.meta = {"nn_module_stack": {}} - assert _get_layer_index(node_with_empty_stack) is None - - -# ---- _normalize_cluster_layer tests ---- - - -@apply_cuda_patches -def test_normalize_cluster_layer_swaps_backward_roots(device_mesh_1d): +def test_export_json_handles_repeated_layer_clusters(device_mesh_1d): dim = 64 with torch.device("meta"): model = _RepeatedLayerModel(dim, n_layers=4) @@ -92,36 +46,11 @@ def test_normalize_cluster_layer_swaps_backward_roots(device_mesh_1d): with autop: autop.add_input_constraints([(Shard(0),)]) autop.add_output_constraints([(Shard(0),)]) - opt = autop.sharding_optimizer - - # Build cluster_roots from cluster_links - cluster_roots = {} - for linked_key, root_key in opt.cluster_links.items(): - linked_node = opt.nodes[linked_key[0]] - root_node = opt.nodes[root_key[0]] - cluster_roots[linked_node] = root_node - - if not cluster_roots: - return # no clusters to test - - _normalize_cluster_layer(cluster_roots) - - # After normalization, all roots should be in the canonical (lowest) layer - root_layers = set() - for root in set(cluster_roots.values()): - idx = _get_layer_index(root) - if idx is not None: - root_layers.add(idx) - - if root_layers: - canonical = min(root_layers) - # Most roots should be in the canonical layer (some may not have - # a copy there due to boundary asymmetry) - assert canonical == min(root_layers) - + autop.sharding_optimizer.get_solution() + data = autop.sharding_optimizer.get_json() -def test_normalize_cluster_layer_empty(): - _normalize_cluster_layer({}) + assert data["nodes"] + assert any("cluster_id" in node for node in data["nodes"]) # ---- export_sharding_json tests ---- diff --git a/tests/test_grad_reduce_dtype.py b/tests/test_grad_reduce_dtype.py index 7fb9383c..dd47f8ce 100644 --- a/tests/test_grad_reduce_dtype.py +++ b/tests/test_grad_reduce_dtype.py @@ -37,7 +37,7 @@ def forward(self, x): def _run_autop(mesh, model_fn, input_fn, mp_policy, repeated_subgraphs=False): - """Run AutoParallel and return the solution and optimizer.""" + """Run AutoParallel and return the placement solution.""" with torch.device("meta"): model = model_fn() @@ -48,80 +48,71 @@ def _run_autop(mesh, model_fn, input_fn, mp_policy, repeated_subgraphs=False): autop.add_output_constraints([(Shard(0),) * mesh.ndim]) sharding_placement = autop.optimize_placement(verbose=False) - return sharding_placement, autop + return sharding_placement -def _assert_no_pre_cast_redistribution( - sharding_placement, autop, require_linked_validation=False -): - """Assert that no chosen pre-cast decision var has comm_cost > 0. - - This directly validates the solver constraint: no redistribution - should occur on unary-chain edges before the dtype_cast node. - """ - from torch._functorch._aot_autograd.fx_utils import get_param_and_grad_nodes +def _iter_dtype_casts(sharding_placement, dtype): + for node, strategy in sharding_placement.items(): + if node.target != torch.ops.autoparallel.dtype_cast.default: + continue + if node.meta["val"].dtype == dtype: + yield node, strategy + + +def _assert_unary_chain_has_no_pre_cast_redistribution(sharding_placement, cast_node): + node = cast_node + while True: + input_nodes = [n for n in node.all_input_nodes if n in sharding_placement] + if len(input_nodes) != 1: + return + + producer = input_nodes[0] + input_spec = sharding_placement[node].input_specs[0] + producer_spec = sharding_placement[producer].output_specs + if not hasattr(input_spec, "placements") or not hasattr( + producer_spec, "placements" + ): + return + + assert input_spec.placements == producer_spec.placements, ( + f"dtype_cast pre-chain edge {producer.name}->{node.name} changes " + f"placement from {producer_spec.placements} to {input_spec.placements}" + ) + node = producer - opt = autop.sharding_optimizer - validated_pre_cast_keys = 0 - validated_linked_keys = 0 - for param, grad in get_param_and_grad_nodes(opt.graph).values(): - if grad is None: +def _assert_reduce_after_cast(sharding_placement, reduce_dtype, min_casts=1): + matched = 0 + for node, strategy in _iter_dtype_casts(sharding_placement, reduce_dtype): + output_spec = strategy.output_specs + if not hasattr(output_spec, "placements") or not any( + p.is_partial() for p in output_spec.placements + ): continue + _assert_unary_chain_has_no_pre_cast_redistribution(sharding_placement, node) + matched += 1 + + assert matched >= min_casts, ( + f"Expected at least {min_casts} dtype_cast outputs with Partial placement, " + f"but found {matched}" + ) + - # Build the pre-cast node set (same logic as the constraint) - chain = [grad] - n = grad - while len(n.all_input_nodes) == 1: - parent = n.all_input_nodes[0] - if len(parent.all_input_nodes) != 1: - break - chain.append(parent) - n = parent - - cast_idx = None - for i, node in enumerate(chain): - if node.target == torch.ops.autoparallel.dtype_cast.default: - cast_idx = i - break - - if cast_idx is None: +def _assert_forward_allgather_after_cast(sharding_placement, param_dtype, min_casts=1): + matched = 0 + for node, strategy in _iter_dtype_casts(sharding_placement, param_dtype): + output_spec = strategy.output_specs + if not hasattr(output_spec, "placements") or any( + p.is_partial() for p in output_spec.placements + ): continue + _assert_unary_chain_has_no_pre_cast_redistribution(sharding_placement, node) + matched += 1 - pre_cast_node_idxs = set() - for node in chain[cast_idx:]: - if node in opt.node_map: - pre_cast_node_idxs.add(opt.node_map[node]) - - # Check that no chosen decision var on a pre-cast node has comm_cost > 0 - for key in opt.selected_keys: - node_idx, argi, out_idx, inp_idx = key - if node_idx not in pre_cast_node_idxs: - continue - dv = opt._resolve_decision_var(key) - validated_pre_cast_keys += 1 - if key in opt.cluster_links: - validated_linked_keys += 1 - assert dv.comm_cost == 0, ( - f"Pre-cast node {opt.nodes[node_idx].name} has chosen decision var " - f"with comm_cost={dv.comm_cost} > 0 (strategy {dv.strategy}). " - f"Redistribution should not happen before the dtype_cast." - ) - - # Also check dtype_cast output has Partial (indirect but readable check) - dtype_cast_node = chain[cast_idx] - spec = sharding_placement.get(dtype_cast_node) - if spec is not None: - assert any(p.is_partial() for p in spec.output_specs.placements), ( - f"dtype_cast {dtype_cast_node.name} should have Partial output " - f"(no pre-cast reduction), but got {spec.output_specs.placements}" - ) - - assert validated_pre_cast_keys > 0, "Expected to validate at least one pre-cast key" - if require_linked_validation: - assert ( - validated_linked_keys > 0 - ), "Expected to validate at least one cluster-linked pre-cast key" + assert matched >= min_casts, ( + f"Expected at least {min_casts} forward dtype_cast outputs, " + f"but found {matched}" + ) @apply_cuda_patches @@ -139,8 +130,8 @@ def input_fn(): mp_policy = MixedPrecisionPolicy( param_dtype=torch.bfloat16, reduce_dtype=torch.float32 ) - sharding_placement, autop = _run_autop(mesh, model_fn, input_fn, mp_policy) - _assert_no_pre_cast_redistribution(sharding_placement, autop) + sharding_placement = _run_autop(mesh, model_fn, input_fn, mp_policy) + _assert_reduce_after_cast(sharding_placement, torch.float32) @apply_cuda_patches @@ -162,12 +153,10 @@ def input_fn(): mp_policy = MixedPrecisionPolicy( param_dtype=torch.bfloat16, reduce_dtype=torch.float32 ) - sharding_placement, autop = _run_autop( + sharding_placement = _run_autop( mesh, model_fn, input_fn, mp_policy, repeated_subgraphs=True ) - _assert_no_pre_cast_redistribution( - sharding_placement, autop, require_linked_validation=True - ) + _assert_reduce_after_cast(sharding_placement, torch.float32, min_casts=2) @apply_cuda_patches @@ -187,7 +176,7 @@ def input_fn(): mp_policy = MixedPrecisionPolicy( param_dtype=torch.float32, reduce_dtype=torch.bfloat16 ) - sharding_placement, _ = _run_autop(mesh, model_fn, input_fn, mp_policy) + sharding_placement = _run_autop(mesh, model_fn, input_fn, mp_policy) assert sharding_placement is not None, "Optimizer should find a feasible solution" @@ -206,80 +195,13 @@ def input_fn(): mp_policy = MixedPrecisionPolicy( param_dtype=torch.float32, reduce_dtype=torch.float32 ) - sharding_placement, _ = _run_autop(mesh, model_fn, input_fn, mp_policy) + sharding_placement = _run_autop(mesh, model_fn, input_fn, mp_policy) assert sharding_placement is not None, "Optimizer should find a feasible solution" # ---- Forward dtype_cast constraint tests ---- -def _assert_no_fwd_pre_cast_redistribution( - sharding_placement, autop, require_linked_validation=False -): - """Assert that no forward dtype_cast has redistribution on its input edge. - - Validates that every forward dtype_cast node matches its producer's - placement, so allgather happens on the cast output (smaller dtype) - rather than on the raw parameter (larger storage dtype). - """ - from torch._functorch._aot_autograd.fx_utils import get_param_and_grad_nodes - - opt = autop.sharding_optimizer - validated_keys = 0 - validated_linked_keys = 0 - - for param, _grad in get_param_and_grad_nodes(opt.graph).values(): - # Walk forward from param through single-user unary nodes - n = param - while True: - if n.target == torch.ops.autoparallel.dtype_cast.default: - break - users = list(n.users.keys()) - if len(users) != 1: - break - child = users[0] - if len(child.all_input_nodes) != 1: - break - n = child - - if n.target != torch.ops.autoparallel.dtype_cast.default: - continue - - # Only check downcasts (storage > param_dtype) - storage_dtype = param.meta["val"].dtype - cast_dtype = n.meta["val"].dtype - if cast_dtype.itemsize >= storage_dtype.itemsize: - continue - - # Collect pre-cast node indices - fwd_pre_cast_node_idxs = set() - node = n - while node != param: - if node in opt.node_map: - fwd_pre_cast_node_idxs.add(opt.node_map[node]) - node = node.all_input_nodes[0] - - for key in opt.selected_keys: - node_idx, argi, out_idx, inp_idx = key - if node_idx not in fwd_pre_cast_node_idxs: - continue - dv = opt._resolve_decision_var(key) - validated_keys += 1 - if key in opt.cluster_links: - validated_linked_keys += 1 - assert dv.comm_cost == 0, ( - f"Forward pre-cast node {opt.nodes[node_idx].name} has chosen " - f"decision var with comm_cost={dv.comm_cost} > 0. " - f"Redistribution should not happen before the dtype_cast." - ) - - assert validated_keys > 0, "Expected to validate at least one forward pre-cast key" - if require_linked_validation: - assert ( - validated_linked_keys > 0 - ), "Expected to validate at least one cluster-linked forward pre-cast key" - - @apply_cuda_patches def test_fwd_allgather_in_param_dtype(device_mesh_1d): """With param_dtype=bf16 and f32 storage, forward allgather should happen @@ -296,8 +218,8 @@ def input_fn(): mp_policy = MixedPrecisionPolicy( param_dtype=torch.bfloat16, reduce_dtype=torch.float32 ) - sharding_placement, autop = _run_autop(mesh, model_fn, input_fn, mp_policy) - _assert_no_fwd_pre_cast_redistribution(sharding_placement, autop) + sharding_placement = _run_autop(mesh, model_fn, input_fn, mp_policy) + _assert_forward_allgather_after_cast(sharding_placement, torch.bfloat16) @apply_cuda_patches @@ -317,8 +239,8 @@ def input_fn(): mp_policy = MixedPrecisionPolicy( param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16 ) - sharding_placement, autop = _run_autop(mesh, model_fn, input_fn, mp_policy) - _assert_no_fwd_pre_cast_redistribution(sharding_placement, autop) + sharding_placement = _run_autop(mesh, model_fn, input_fn, mp_policy) + _assert_forward_allgather_after_cast(sharding_placement, torch.bfloat16) @apply_cuda_patches @@ -336,11 +258,11 @@ def input_fn(): mp_policy = MixedPrecisionPolicy( param_dtype=torch.bfloat16, reduce_dtype=torch.float32 ) - sharding_placement, autop = _run_autop( + sharding_placement = _run_autop( mesh, model_fn, input_fn, mp_policy, repeated_subgraphs=True ) - _assert_no_fwd_pre_cast_redistribution( - sharding_placement, autop, require_linked_validation=True + _assert_forward_allgather_after_cast( + sharding_placement, torch.bfloat16, min_casts=2 ) @@ -363,14 +285,5 @@ def input_fn(): mp_policy = MixedPrecisionPolicy( param_dtype=torch.float32, reduce_dtype=torch.bfloat16 ) - sharding_placement, autop = _run_autop(mesh, model_fn, input_fn, mp_policy) - - # Verify no fwd_param_dtype constraints were added to the ILP - opt = autop.sharding_optimizer - fwd_constraint_names = [ - name for name in opt.prob.constraints if name.startswith("fwd_param_dtype") - ] - assert len(fwd_constraint_names) == 0, ( - f"Forward dtype constraint should not fire for upcasting, " - f"but found {len(fwd_constraint_names)} constraints: {fwd_constraint_names}" - ) + sharding_placement = _run_autop(mesh, model_fn, input_fn, mp_policy) + assert sharding_placement is not None, "Optimizer should find a feasible solution" diff --git a/tests/test_optimize_placement.py b/tests/test_optimize_placement.py index 59a4cf7c..56dc43ed 100644 --- a/tests/test_optimize_placement.py +++ b/tests/test_optimize_placement.py @@ -63,6 +63,19 @@ def forward(self, x): return o +class StackedLinear(nn.Module): + def __init__(self, dim, n_layers): + super().__init__() + self.layers = nn.ModuleList( + [nn.Linear(dim, dim, bias=False) for _ in range(n_layers)] + ) + + def forward(self, x): + for layer in self.layers: + x = layer(x).relu() + return x + + def _make_model_and_input_fn( mesh, model_type="ffn_with_multiple_input_output", device="cuda" ): @@ -92,6 +105,38 @@ def input_fn(): return model_fn, input_fn +def _placement_signature(solution): + def output_signature(output_specs): + if isinstance(output_specs, (tuple, list)): + return tuple( + None if spec is None else tuple(str(p) for p in spec.placements) + for spec in output_specs + ) + return tuple(str(p) for p in output_specs.placements) + + return { + node.name: output_signature(strategy.output_specs) + for node, strategy in solution.items() + } + + +def _solve_stacked_linear(device_mesh, n_layers): + batch_size = 2 * device_mesh.shape[0] + dim = 16 + + with torch.device("meta"): + model = StackedLinear(dim, n_layers) + + def input_fn(): + return torch.randn(batch_size, dim, device="cuda") + + with AutoParallel(model, input_fn, device_mesh, repeated_subgraphs=False) as autop: + autop.add_input_constraints([(Shard(0),)]) + autop.add_output_constraints([(Shard(0),)]) + autop.add_parameter_memory_constraint(low=None, high=None) + return _placement_signature(autop.optimize_placement()) + + @apply_cuda_patches @pytest.mark.parametrize( "model_type", ["ffn_with_multiple_input_output", "transformer_block"] @@ -183,6 +228,19 @@ def test_optimization_finds_fsdp_and_ddp_1d(device_mesh_1d, high_mem, model_type assert p.input_specs[1].placements == (Replicate(),) +@apply_cuda_patches +def test_fast_build_preserves_optimizer_solution(device_mesh_1d, monkeypatch): + monkeypatch.setenv("AP_FAST_BUILD", "0") + monkeypatch.setenv("AP_PARALLEL_BUILD", "1") + baseline = _solve_stacked_linear(device_mesh_1d, n_layers=8) + + monkeypatch.setenv("AP_FAST_BUILD", "1") + monkeypatch.setenv("AP_PARALLEL_BUILD", "2") + fast = _solve_stacked_linear(device_mesh_1d, n_layers=8) + + assert fast == baseline + + _expected_param_placements_ffn = [(Shard(0), Shard(0)), (Shard(0), Shard(1))]