Aggregate: optimizer build, TRW-S, lazy costs, and split-dim seed

[AlbedoWang]

Aggregate only. Do not review or merge this combined diff.

Review and merge the stack in order: #520 -> #521 -> #522 -> #523.

Review stack

1. Speed up sharding optimizer construction #520: algorithm-preserving optimizer construction, compact cluster links, invalid-edge pruning and serialization.
2. Add LP relaxation and eager TRW-S solvers #521: LP relaxation and eager full-cost TRW-S.
3. Build TRW-S costs lazily without PuLP #522: lazy/no-PuLP cost construction for TRW-S.
4. Add corrected split-dimension seed search #523: corrected split-dimension seed search and per-dimension PR514 local_map projection.

This Draft keeps the complete top-of-stack diff and profiling context in one place. Its head mirrors #523. The four stacked PRs are the merge targets.

Final-tree policy

The stack was rebuilt from this PR's final implementation rather than replaying PR484/PR511/PR513 wholesale. This preserves the selected main+PR511 propagation behavior and excludes abandoned PR484 annotation/DP/min-sum paths.

Compared with the previous aggregate head, implementation files are unchanged except for the direct-LP objective fix. Additional differences remove machine-local configuration and unrelated artifacts, update tests for node-level cluster links, and add a fast-build equivalence regression.

PR514 remains the source of truth for 3+D MoE local_map, flattened EP groups, efsdp boundary placements, and backward-gradient duality.

Validation

• Changed-file pre-commit passes whitespace, AST, conflict, large-file, EOF, Black, flake8, isort and mypy.
• Full repository suite: 528 passed, 1 xfailed, 4 xpassed.
• Foundation related suites: 33 passed.
• LP/eager-TRW-S suites: 7 passed.
• Lazy/no-PuLP suites: 8 passed.
• Split/local-map suites: 3 passed.

The five new failures from the previous aggregate CI are resolved: LP constraint feasibility, export JSON cluster links, both repeated-subgraph dtype constraints, and direct split-seed LP objective initialization.

Corrected MoE profiles

Profiles use fake tensors and a fake process group; they measure search latency and host RSS, not distributed execution. Each result is one deterministic run; variance is unavailable.

• 2D (8,8) uses direct unseeded lazy TRW-S, not split-dim: 33.05s versus 37.02s full-cost TRW-S, with identical objective 45680.5744 and placement.
• 3D (2,4,8) corrected split + lazy TRW-S: 111.15s, objective 44260.3021, +1.698% versus unrestricted full/lazy 43521.3868. The unrestricted solution lies inside radius 2; the delta is compute 0, communication +675.915, transition +63.
• 4D (8,2,2,2) corrected split + lazy TRW-S: 201.41s, objective 47952.1351, peak RSS 3.44 GiB. Unrestricted eager and lazy searches exceeded 20 minutes, so no objective gap or speedup is claimed. No 4D sampling experiment was run.

Authored with Claude.

Inherited base checks

PR514 currently has the same unrelated repo-wide mypy failure in autoparallel/tools/overlap_simulator/run.py:404 and the same TorchTitan upstream config-registry failure. The stack does not include unrelated fixes for either path.