Add LP relaxation and eager TRW-S solvers

[AlbedoWang]

Summary

Second PR in the PR519 review stack. It adds LP relaxation and eager full-cost TRW-S over the same optimizer problem; ILP remains the default. Lazy/no-PuLP construction and split-dimension seed search are intentionally left to PR522 and PR523.

Current head 5fa1e9df8e39eb0a272b903ff4fbe4ba353d8f5d also replaces the colliding invalid-cost sentinel and makes approximate solving and serialization honor removal of active memory constraints without changing the search space.

Interface

The supported entry point is AutoParallel:

autop = AutoParallel(
    model,
    input_fn,
    mesh,
    solver="ilp",  # "ilp" (default), "lp", or "approx"
)

with autop:
    solution = autop.optimize_placement(
        solver=None,             # optionally override the constructor choice
        approximate_options=None,
        optimality_check=False,
    )

• ilp: exact PuLP/CBC solve.
• lp: solve the continuous relaxation and extract it only when integral; otherwise raise and direct the caller to ilp.
• approx: eager pairwise-factor TRW-S followed by constrained local-search polish. approximate_options is forwarded to this solver.
• optimality_check=True: solve the LP relaxation as a lower bound and report the certified gap. It requires an ILP/LP-backed optimizer build.

ShardingOptimizer.solve_lp_relaxation() remains the lower-level diagnostic interface. It restores binary variable categories after the relaxation so a later ILP solve is unaffected.

Reproducible tests

The harness and its tests are checked in as tests/search_profile.py and tests/test_search_profile.py; there is no separate profiling directory.

Unit and fake-process-group coverage:

PYTHONPATH=. python -m pytest -q \
  tests/test_search_profile.py \
  tests/test_approximate_sharding.py \
  tests/test_lp_relaxation.py \
  tests/test_optimize_placement.py::test_invalid_strategies_are_pruned \
  -k "not ilp_and_approx_match"

Real four-GPU end-to-end behavior, also run by the Test CUDA multi-GPU job:

PYTHONPATH=. python -m pytest \
  tests/test_search_profile.py::TestRealDsv3SolverE2E::test_ilp_and_approx_match \
  -v -s

This test enters through AutoParallel for both ILP and approximate solvers, adds the same memory/input/output constraints, solves, applies the placement, initializes identical weights, runs forward, reconstructs the full output, computes the same scalar loss, and runs backward. It compares the objective, full output, and every parameter gradient and emits per-rank phase/memory breakdown as AUTOPARALLEL_E2E_BREAKDOWN.

Full suite:

PYTHONPATH=. python -m pytest -q tests

Search-only reproduction

The following matrix uses meta models, a fake process group, and explicit H100 properties. It measures placement search only, not distributed model execution. Each JSON records the expanded config, environment, Git commit/status, constraints, solver status, objective, placement hash, node/variable counts, phase timings, and process peak RSS.

mkdir -p results

for solver in ilp lp approx; do
  PYTHONPATH=. PYTHONHASHSEED=0 timeout --signal=TERM --kill-after=30s 20m \
    python tests/search_profile.py \
      --model llama1b --mesh 8,8 --solver "$solver" \
      --revision-label "$(git rev-parse --short HEAD)" \
      --output "results/llama1b_${solver}.json"

  PYTHONPATH=. PYTHONHASHSEED=0 timeout --signal=TERM --kill-after=30s 20m \
    python tests/search_profile.py \
      --model llama8b --mesh 8,8 --solver "$solver" \
      --revision-label "$(git rev-parse --short HEAD)" \
      --output "results/llama8b_${solver}.json"

  PYTHONPATH=. PYTHONHASHSEED=0 timeout --signal=TERM --kill-after=30s 20m \
    python tests/search_profile.py \
      --model dsv3 --moe-layout 2d --solver "$solver" \
      --revision-label "$(git rev-parse --short HEAD)" \
      --output "results/dsv3_${solver}.json"
done

Add --detailed-solution to emit per-node placement and component cost contributions. Reproduction should use a clean worktree and verify that the recorded git.status is empty.

Search settings:

• Python 3.12.13, PyTorch 2.14.0.dev20260629+cu130, CUDA 13.0, PuLP 3.3.2; AMD EPYC 9654 host.
• PYTHONHASHSEED=0, torch.manual_seed(0), repeated subgraphs, default NCCL cost model, parameter bfloat16, reduction float32, and the same parameter memory/input/output constraints. Dynamic tracing is enabled only for DSv3.
• Fake H100: capability 9.0, 80 GiB, 132 SMs, 50 MiB L2, eight reported local devices; the fake process group world size is 64.
• LLaMA1B (8,8): dim 2048, 16 layers, 32 heads, 8 KV heads, FFN multiplier 1.5, FFN multiple 256, RoPE theta 500000, vocab 128256, sequence 2048, global batch 16. Mesh axes are (dp,tp); input is (Shard(0), Replicate()) and output is (Shard(0), Shard(2)).
• LLaMA8B (8,8): dim 4096, 32 layers, 32 heads, 8 KV heads, FFN multiplier 1.3, FFN multiple 1024; remaining settings and placements as LLaMA1B.
• DSv3 (8,8): dim 256, 6 layers (1 dense), 16 heads, 64 experts, vocab 2048, sequence 2048, global batch 512; degrees (8,8,1,1,8) for (dp_replicate,dp_shard,cp,tp,ep). Input and output are Shard(0) on both mesh axes.

Search-only results

The matrix was run once per row at 1b92f5fb7d81315b9a9d417d25eb4c213498b9a4. The final PR head moves the harness to tests/, extracts the same validation checks for unit coverage, and adds the real E2E/CI invocation; it does not change the solver implementation. Timings use time.perf_counter; RSS uses process ru_maxrss. factor build and solver core are contained in solve call.

Workload	Solver	Trace	Optimizer init	Solve call	Factor build	Solver core	Unaccounted	Search total	RSS GiB	Objective
LLaMA1B	ILP	7.748s	12.780s	13.964s	-	11.816s	0.191s	34.685s	1.834	71178.582224
LLaMA1B	LP	6.280s	12.669s	15.457s	-	12.714s	0.161s	34.567s	1.813	71178.582224
LLaMA1B	eager TRW-S	9.829s	14.215s	8.590s	1.744s	4.940s	0.157s	32.792s	1.731	71178.582224
LLaMA8B	ILP	13.421s	15.529s	17.521s	-	15.851s	0.354s	46.826s	2.756	257769.291686
LLaMA8B	LP	13.996s	17.358s	14.475s	-	12.055s	0.353s	46.182s	2.770	257769.291686
LLaMA8B	eager TRW-S	13.810s	19.676s	10.167s	1.268s	5.548s	0.402s	44.066s	2.712	257769.291686
DSv3	ILP	16.979s	20.716s	21.545s	-	17.744s	0.304s	59.545s	2.144	45680.574430
DSv3	LP	16.465s	18.228s	15.996s	-	11.730s	0.271s	50.962s	2.073	45680.574430
DSv3	eager TRW-S	18.813s	18.940s	12.780s	0.835s	8.826s	0.296s	50.829s	1.765	45680.574430

All nine jobs returned success, finite objectives, nonempty placements, and zero constraint violations. ILP and LP report Optimal; approximate reports Heuristic and Solution Found. ILP, LP, and approximate objectives match for all three workloads. LLaMA8B placements are identical. LLaMA1B differs at 16 of 4,299 equal-cost view placements and DSv3 differs at 10 of 2,288 equal-cost permute placements; detailed compute, communication, transition, and total cost contributions match at recorded precision.

These are single observations with no run-to-run variance. They are diagnostic breakdowns, not stable latency, speedup, throughput, or memory-improvement claims.

Real 2x2 GPU E2E result

Run at final head 9e6181f5bae8ff1fb02f75416b672f1c84e531c5 with Python 3.12.13, PyTorch 2.14.0.dev20260629+cu130, CUDA 13.0, NCCL 2.30.7+cuda13.3, and four NVIDIA H100 GPUs.

Configuration: DeepSeekV3 debug model, dim 256, 6 layers (1 dense), 16 heads, 8 experts, vocab 2048, sequence 2048, global/local batch 32/8, compute bfloat16, reduction float32; mesh (2,2) with dp_replicate and dp_shard_in_ep, degrees (2,2,1,1,2). Inputs and outputs are Shard(0) on both mesh axes. Runtime collectives are real NCCL operations. The non-canonical four-GPU topology uses the repository's documented default cost-model fallback.

Behavior result:

• Both solvers: objective 34091.15472077961, 2,288 solution nodes, zero constraint violations.
• Full output: match at BF16-derived rtol=atol=0.015625; observed max/mean absolute difference 0.03125 / 0.0004442967.
• All parameter gradients match using BF16-derived tolerance for BF16 tensors and PyTorch defaults for other dtypes; observed max/mean absolute difference is 4.6566e-10 / 1.1185e-12.

Breakdown method: time.perf_counter per rank, with CUDA synchronization around materialization, forward, full-output construction, and backward. Values are mean [min, max] over four ranks from one ILP run followed by one approximate run.

Phase	ILP seconds	Approx seconds
Model setup	0.009 [0.008, 0.009]	0.007 [0.007, 0.008]
Graph trace	12.950 [12.477, 13.212]	10.864 [10.559, 11.154]
Optimizer init	19.403 [18.342, 21.626]	16.127 [14.417, 18.362]
User constraints	0.001 [0.001, 0.001]	0.001 [0.001, 0.001]
Solve call	15.023 [13.906, 16.092]	13.492 [12.857, 14.611]
Factor build (inside solve)	n/a	1.110 [0.972, 1.207]
Solver core (inside solve)	11.770 [10.871, 12.347]	9.418 [9.267, 9.754]
Apply placement	8.539 [8.229, 9.201]	8.861 [8.642, 9.103]
Materialize/init	4.001 [0.833, 6.531]	4.015 [0.064, 6.765]
Forward	1.116 [1.114, 1.117]	0.204 [0.203, 0.205]
Full-output construction	0.004 [0.003, 0.004]	0.004 [0.002, 0.005]
Backward	1.078 [1.072, 1.083]	0.192 [0.189, 0.196]
E2E through backward	62.916 [62.854, 63.014]	54.395 [53.935, 54.583]

Mean [min, max] CUDA peak allocated memory is 4.829 [4.821, 4.839] GiB for ILP and 4.904 [4.896, 4.914] GiB for approximate. Peak reserved memory is 4.992 [4.992, 4.992] GiB and 5.062 [5.061, 5.063] GiB, respectively. This is one ordered run, so later execution can benefit from process, compiler, allocator, and filesystem caches; no performance improvement is claimed.

Published evidence

• Complete reproducibility and validation report (Markdown, SHA-256 41aac1783a5c67acccdf570f64d16b51ea0188b2455beb34b62b79a0815e6e4b).

Validation status

• Targeted unit/search/solver command above: 22 passed, 1 real-GPU case deselected, 240.89s.
• Post-restack focused invalid-cost/memory-state regressions: 3 passed at current head.
• Pre-restack implementation-head related coverage: 24 passed, 203.08s; range-diff confirms the PR521 production diff is unchanged.
• Final clean-commit real E2E: 1 passed, 181.34s.
• Full local suite including the real-GPU E2E: 543 passed, 1 xfailed, 4 xpassed, 913.29s.
• Pre-restack GitHub Test CUDA single GPU, examples, and multi GPU (including this E2E): passed. Current-head CI was restarted by the restack.
• CodeQL and CLA: passed.
• Repository-wide lint remains red on the same two pre-existing mypy errors in autoparallel/tools/overlap_simulator/run.py as the PR base. TorchTitan integration remains red on the same upstream config_registry import error as the prior PR head.

Authored with Claude.

Stack

• 1/4: Speed up sharding optimizer construction #520
• 2/4: this PR
• 3/4: Build TRW-S costs lazily without PuLP #522
• 4/4: Add corrected split-dimension seed search #523
• Aggregate index only: Aggregate: optimizer build, TRW-S, lazy costs, and split-dim seed #519

[AlbedoWang]

CI note for reviewers: the repo-wide lint job reaches and passes isort, Black, and flake8, then remains red only for mypy errors in autoparallel/tools/overlap_simulator/run.py at lines 255 and 402. PR514/base is already red in that unrelated file (see #514), and this stack does not modify it. The PR514 TorchTitan integration check is also already red upstream; layer validation and the saved full-suite evidence are linked from #519.

[Copilot]

Pull request overview

This PR extends AutoParallel’s sharding optimizer with additional solver backends over the same eagerly-constructed optimization problem: an LP-relaxation path (used directly when integral) and an approximate TRW-S-based solver, while keeping ILP/CBC as the default. This fits into the core “build ILP once, solve with selectable backend” workflow of autoparallel/api.py + autoparallel/optimize_sharding.py.

Changes:

• Add an eager TRW-S approximate solver (ApproximateShardingSolver) that reuses the existing PuLP-built optimizer and writes back assignments for exact scoring/validation.
• Add LP-relaxation solving and lower-bound certification utilities in the sharding optimizer, plus API plumbing to select "ilp" | "lp" | "approx".
• Add targeted tests for LP integrality/lower-bound certification and approximate-solver faithfulness + constraint adherence.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
autoparallel/api.py	Adds solver selection to AutoParallel.optimize_placement() and optional LP lower-bound optimality checking/logging.
autoparallel/optimize_sharding.py	Adds LP relaxation/lower-bound routines, prunes infinite-cost edges, and updates constraint/objective handling to support alternate solvers.
autoparallel/approximate_sharding.py	Introduces TRW-S + local search approximate solver that operates on the existing optimizer’s decision vars/constraints.
tests/test_approximate_sharding.py	New tests covering approx-vs-ILP gap bounds, feasibility/constraint validity, LP==ILP on the fixture, and optimality-check logging.
tests/test_lp_relaxation.py	New test validating LP relaxation provides a certified lower bound and restores optimizer state afterward.
tests/conftest.py	Adds an autouse fixture to reset placement-options caching between tests.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.