Speed up sharding optimizer build

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,29 +265,31 @@ 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
     opt.mesh = _MeshPlaceholder(save_dict["mesh_shape"], save_dict["mesh_dim_names"])
 
     # 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):

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

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)

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.

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.

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 ----