Build TRW-S costs lazily without PuLP

[AlbedoWang]

Summary

Third PR in the PR519 review stack. For solver="approx", it avoids eager per-edge cost and DecisionVar materialization, derives the constraint topology directly, and memoizes complete factor costs on demand. Factor ranking still includes compute, communication, and transition costs. ILP and LP construction remain eager and unchanged.

Interface

autop = AutoParallel(
    model,
    input_fn,
    mesh,
    solver="approx",
    lazy_costs=None,
)

• solver="approx", lazy_costs=None or True: lazy approximate build with no PuLP problem and no eager decision variables.
• solver="approx", lazy_costs=False: eager approximate cost construction for A/B correctness checks; it still does not construct a PuLP problem.
• solver="ilp" or "lp": the existing eager PuLP-backed build.
• lazy_costs=True with ILP/LP is rejected at construction. Full optimizer serialization also requires an ILP/LP build; approximate-only placements can still be serialized.

The approximate solver reports Heuristic / Solution Found; it does not misreport global optimality. Removing active memory constraints clears their state in both eager and no-PuLP builds, while the default memory budget remains unchanged when the caller does not remove it.

Reproduction

The checked-in tests/search_profile.py supports the 3D LLaMA configuration used below. It uses the repository LLaMA example, meta tensors, a fake process group, and explicit H100 properties; it measures placement search, not distributed model execution.

mkdir -p results

# Eager approximate reference
PYTHONPATH=. PYTHONHASHSEED=0 /usr/bin/time -v -o results/eager.time.txt \
  timeout --signal=TERM --kill-after=30s 20m \
  python tests/search_profile.py --model llama1b --mesh 2,4,8 \
  --solver approx --lazy-costs false --detailed-solution \
  --revision-label "$(git rev-parse --short HEAD)" \
  --output results/eager.json

# Default lazy approximate path
PYTHONPATH=. PYTHONHASHSEED=0 /usr/bin/time -v -o results/lazy.time.txt \
  timeout --signal=TERM --kill-after=30s 20m \
  python tests/search_profile.py --model llama1b --mesh 2,4,8 \
  --solver approx --lazy-costs true --detailed-solution \
  --revision-label "$(git rev-parse --short HEAD)" \
  --output results/lazy.json

Each JSON records the expanded model/mesh config, environment, Git revision and status, objective, placement hash, optimizer counts, phase timings, and process peak RSS. Lazy mode has no PuLP constraint list; feasibility is represented by the approximate solver status rather than a post-hoc PuLP constraint scan.

Historical performance

The following is the previously collected single-run result at 829b11720bed03673aefce33161f8912908a3863. It was not rerun after the stack rebase. The current harness exposes the same model, mesh, and eager/lazy modes, but these numbers should be read as pinned historical evidence, not a current-head benchmark.

Environment: Python 3.12.13, PyTorch 2.14.0.dev20260629+cu130, CUDA 13.0, PuLP 3.3.2, AMD EPYC 9654 host. The model is LLaMA1B with dim 2048, 16 layers, 32 heads, 8 KV heads, FFN multiplier 1.5, vocab 128256, sequence length 2048, batch 16, parameter bfloat16, reduction float32, and mesh (2,4,8) named (dp,cp,tp). Peak RSS is GNU time -v maximum resident set size. Timeout was 20 minutes. There was one deterministic run per mode, so variance is unavailable.

Mode	Optimizer init	Factor build	TRW-S solve	Search total	Peak RSS
eager	169.932s	10.805s	17.242s	205.475s	6.166 GiB
lazy	19.715s	112.190s	17.667s	156.058s	4.579 GiB

Observed component deltas were -150.217s optimizer initialization, +101.385s factor construction, and +0.425s TRW-S solve. Search total was 49.417s lower (1.32x), and peak RSS was 1.587 GiB lower (25.7%). Both modes are approximate/no-PuLP builds; the observed difference comes from avoiding eager edge-cost and DecisionVar materialization, with required cost work moving into lazy factor construction. No throughput, training latency, or distributed runtime claim is made.

Both historical modes recorded objective 49282.70137059267, placement SHA-256 b23e62433fe08ccba75fce1103cd47302987c5a75c0dc5e3ba6bf5391bffb742, and identical placements for all 4,299 nodes. The 3D LP baseline exceeded the 20-minute cap, so this establishes equality to eager TRW-S, not a certified global optimality gap.

Correctness validation

At the rebased implementation head, focused coverage compares lazy and eager approximate builds on the existing transformer example, requires equal objectives and selected keys, verifies that valid cost 10000.0 is not treated as forbidden, verifies removal of active memory constraints, checks placement serialization, and rejects lazy ILP/LP construction.

# Current-head eager/lazy build-state regressions.
PYTHONPATH=. PYTHONHASHSEED=0 python -m pytest -q \
  tests/test_optimize_placement.py::test_approx_lazy_build_matches_eager_example \
  tests/test_optimize_placement.py::test_lazy_costs_requires_approx_solver

# Broad approximate, serialization, profile, and placement coverage.
PYTHONPATH=. PYTHONHASHSEED=0 python -m pytest -q \
  tests/test_approximate_sharding.py tests/test_serialization.py \
  tests/test_search_profile.py tests/test_optimize_placement.py \
  -k "not ilp_and_approx_match"

# Profile harness and its 3D argument subset.
PYTHONPATH=. PYTHONHASHSEED=0 python -m pytest -q \
  tests/test_search_profile.py -k "not ilp_and_approx_match"
PYTHONPATH=. PYTHONHASHSEED=0 python -m pytest -q \
  tests/test_search_profile.py -k validate_search_profile_args

Current restacked head: focused build-state regressions 3 passed. At the pre-restack implementation head e403a490158e8c519758b11ab70285216950ec7b, whose PR522 range-diff is unchanged: broad coverage 63 passed, 1 deselected, profile-harness coverage 17 passed, 1 deselected, and 3D argument coverage 9 passed, 9 deselected. Changed-file Black, isort, flake8, mypy, and git diff --check passed.

The existing four-GPU DeepSeekV3 ILP-vs-approximate E2E now reads no-PuLP approximate results correctly. It was not executable on the local host and is left to the GitHub multi-GPU job. Repository-wide mypy remains red on the same two pre-existing errors in autoparallel/tools/overlap_simulator/run.py as the PR base; changed-file mypy passes.

Authored with Claude.

Stack

• 1/4: Speed up sharding optimizer construction #520
• 2/4: Add LP relaxation and eager TRW-S solvers #521
• 3/4: this PR
• 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 (3/4 in the PR519 stack) refactors approximate (TRW-S) solver construction to avoid PuLP variable/constraint creation and to compute TRW-S costs lazily on demand, aiming to reduce optimizer initialization time and host RSS while preserving equivalence with the eager TRW-S path.

Changes:

• Add build_pulp / build_costs switches to ShardingOptimizer to support a “lite” build mode (no PuLP, optionally no eager per-edge cost materialization).
• Teach ApproximateShardingSolver to derive constraint topology directly from the FX graph + cluster_links + _constraint_log when PuLP is absent, and to memoize compute/edge costs as needed.
• Add a regression test ensuring lite-build approximate results match the full-build approximate results byte-identically on the fixture.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
autoparallel/optimize_sharding.py	Introduces lite build controls and makes several constraint/solution helpers tolerant of prob=None.
autoparallel/approximate_sharding.py	Adds direct-topology extraction and lazy cost providers for TRW-S when PuLP/costs aren’t prebuilt.
autoparallel/api.py	Exposes lazy_costs and wires solver choice to optimizer build mode; adds runtime guardrails for incompatible solves.
tests/test_approximate_sharding.py	Adds an equivalence regression for full vs lite optimizer build under TRW-S.

---

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