Add corrected split-dimension seed search - #523
Open
AlbedoWang wants to merge 5 commits into
Open
Conversation
This was referenced Jul 27, 2026
meta-cla
Bot
added
the
CLA Signed
This was referenced Jul 27, 2026
AlbedoWang
force-pushed
the
kaijian/final-opt-split
branch
from
July 27, 2026 03:12
3f897d0 to
0d9ee3d
Compare
Author
|
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. |
AlbedoWang
force-pushed
the
kaijian/final-opt-split
branch
from
July 28, 2026 02:31
0d9ee3d to
e2b379f
Compare
AlbedoWang
force-pushed
the
kaijian/final-opt-split
branch
from
July 28, 2026 02:33
e2b379f to
bfbe4fb
Compare
There was a problem hiding this comment.
Pull request overview
Adds a split-dimension seed mechanism to accelerate approximate sharding search by running per-dimension 1D solves (with per-dim local_map projection and optional fabric-aware NCCL topology) and then restricting subsequent strategy enumeration to a Hamming ball around the per-node seed placements.
Changes:
- Introduces
autoparallel.mesh_search.build_split_dim_seed()plus helpers to projectlocal_mapplacements per dimension and cache 1D solve results. - Threads
strategy_seed/strategy_radiusthroughAutoParallel→ShardingOptimizer, and applies seed-based filtering during placement-option generation. - Adds unit tests and test fixtures to validate local-map projection, cache keys, and feasibility of seeded ILP/LP solves.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_split_dim_seed.py | New tests covering local_map per-dim projection and seeded solve feasibility. |
| tests/conftest.py | Adds CUDA device-property patching and an autouse fixture to clear the global placement-options cache. |
| autoparallel/shardings/propagation_rules.py | Adds global seed/radius state and filters _create_all_options() by seed ball. |
| autoparallel/shardings/placement_options.py | Incorporates seed into placement-options caching and filters strategies by seed ball. |
| autoparallel/optimize_sharding.py | Adds seed/radius plumbing; sets/clears seed state around strategy enumeration. |
| autoparallel/mesh_search.py | New module implementing split-dimension seed construction and cache keying. |
| autoparallel/cost_models/nccl_cost_model.py | Adds mesh_dim_topo_override support for per-dimension NCCL topology derivation. |
| autoparallel/api.py | Adds strategy_radius/seed_input_placements API and integrates seed building into optimizer construction. |
| autoparallel/init.py | Exports build_split_dim_seed at package top level. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| "name": "H100", | ||
| "total_memory": 80 * 1024**3, | ||
| "multi_processor_count": 132, | ||
| "L2_cache_size": 50 * 1024**2, |
Comment on lines
+287
to
+298
| def _seed_cache_key(): | ||
| seed = get_strategy_seed() | ||
| if seed is None: | ||
| return None | ||
| node_name = get_current_seed_node() | ||
| placements = seed.get(node_name) if node_name is not None else None | ||
| placement_key = None if placements is None else tuple(str(p) for p in placements) | ||
| # Key on the seed PLACEMENT (not node_name): the seed-filtered result depends only | ||
| # on (op, input specs, seed placement, radius), so repeated ops across layers with | ||
| # the same seed placement share the placement-options cache instead of missing | ||
| # per node (which re-runs get_op_strategy for every node). | ||
| return placement_key, get_strategy_radius() |
| if placements is None or within_strategy_seed_ball(placements): | ||
| kept.append(strategy) | ||
|
|
||
| return out_strat if not kept else OpStrategy(kept) |
Comment on lines
+350
to
+353
| assert self.seed_input_placements is not None, ( | ||
| "strategy_radius requires seed_input_placements " | ||
| "(one input placement per mesh dim)" | ||
| ) |
Comment on lines
+222
to
+232
| per_dim: list[dict[str, Placement]] = [] | ||
| for dim_idx, size in enumerate(mesh_shape): | ||
| input_placement = input_placements[dim_idx] | ||
| key = _split_dim_seed_cache_key( | ||
| int(size), | ||
| input_placement, | ||
| seed_cost_model, | ||
| mesh_shape, | ||
| dim_idx, | ||
| fabric_aware=fabric_aware, | ||
| ) |
AlbedoWang
force-pushed
the
kaijian/final-opt-split
branch
from
July 28, 2026 03:23
bfbe4fb to
b374f7f
Compare
AlbedoWang
force-pushed
the
kaijian/final-opt-split
branch
from
July 28, 2026 20:54
b374f7f to
63a2dde
Compare
AlbedoWang
marked this pull request as draft
July 30, 2026 00:09
AlbedoWang
force-pushed
the
kaijian/final-opt-split
branch
from
August 1, 2026 04:50
cfb8be1 to
42dbe53
Compare
Solve each mesh dimension independently, preserve its fabric topology, and project local_map contracts onto the active dimension before restricting the final search. Authored with Claude.
Document and validate the public split-search options, retain feasible operator domains explicitly, broaden approximate candidates, and expose solver profile data for reproducible validation.\n\nAuthored with Claude.
Check in the search-only LLaMA and DSv3 profiler used for heuristic and LP result validation, with full configuration, environment, solution, and timing output.\n\nAuthored with Claude.
Keep radius as the only user-facing control and derive per-dimension seeds from the existing input and output constraints. Include reviewer-runnable profiling commands alongside the harness.\n\nAuthored with Claude.
AlbedoWang
force-pushed
the
kaijian/final-opt-split
branch
from
August 1, 2026 05:26
42dbe53 to
adb8cf7
Compare
Authored with Claude.
AlbedoWang
marked this pull request as ready for review
August 2, 2026 03:42
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fourth PR in the PR519 review stack. It adds split-dimension seed search on top of lazy TRW-S. Each active 1D seed solve projects the existing user input/output constraints,
local_mapplacements, and fabric topology to that dimension; final N-D strategy enumeration uses a Hamming ball around the composed seed.The composed seed now retains the complete selected
OpSpecwitness, including input and output specs. If upstream filtering prevents an input-conditioned rule from enumerating that exact legal witness, final enumeration validates and restores it without requiring the producer to expose the consumer placement. The output-placement radius and empty-ball feasibility fallback are unchanged.The only new user-facing control is
AutoParallel(strategy_radius=2). Forsolver="approx"on meshes with more than two dimensions, radius 1 or 2 enables split search; 0 orNoneretains unrestricted lazy TRW-S. Values greater than 2 warn, and a radius at least as large as the mesh dimensionality is rejected. ILP/LP and 1D/2D meshes retain their original behavior. Seed construction is internal.Reproduction
The existing
tests/search_profile.pycovers 2D/3D LLaMA and 2D/3D/4D DSv3.tests/profiling/run_search_profiles.shcontains the exact row matrix, requires an explicit result directory, records every row's command/status, and returns nonzero if any row fails. Use a clean worktree and place results outside the source checkout:The harness uses repository example models, meta tensors, 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 state, solver status, objective, placement hash, phase timings, and process peak RSS. Lazy approximate rows have no PuLP constraint list; feasibility is represented by
Heuristic/Solution Found, not a post-hoc PuLP constraint scan.Current performance
The following single-run observations were collected at
946f0643bea1e40165a63696b3e18982aea07838. Environment: Python 3.12.13, PyTorch2.14.0.dev20260629+cu130, CUDA 13.0, PuLP 3.3.2, AMD EPYC 9654 host. Each row used a deterministic fresh process withPYTHONHASHSEED=0,torch.manual_seed(0), parameter bf16, reduction fp32, repeated subgraphs, lazy approximate costs, candidate limit 128, and radius 2 for split rows. Peak RSS is GNUtime -vmaximum resident set size. Heuristic timeout was 20 minutes. Variance is unavailable.Seed/otheris search time not attributed to graph trace, final optimizer initialization, factor build, solver, or user constraint calls; split rows include the 1D seed solves.Split vs full searchreports full/split speedup and the split row's reduction in total search time.Objective vs fullreports(split / full - 1) * 100; lower is better. These are single observations, not stable performance claims.These compare two heuristic searches, not certified optimality gaps. DSv3 3D full and split now have identical objectives and placement hashes. The 4D row is a standalone feasibility result because no unrestricted 4D result was recorded.
LP probe status
The serial no-timeout LP matrix remains available as:
LP probes were not rerun because they do not use split seeds and their search space is unaffected by this correction. Historical probes produced no accepted 3D LP objective: LLaMA1B reached CBC output after 67 minutes but PuLP could not parse the solution file; LLaMA8B returned after 64 minutes but its rounded assignment violated constraints; no accepted DSv3 3D result was recorded. This PR makes no global-optimality or certified suboptimality-gap claim.
Current validation
The real E2E used the repository DSv3 model with a
(2, 2, 2)physical mesh, applied the selected placement, and compared distributed forward/backward results with a single-GPU reference. All outputs and gradients were finite; the complete seed-witness scan had zero missing nodes and included the intendedPPP -> SSSredistribution fromsum_2todtype_cast_73. Limits were 0.005 for output relative error and 0.02 for gradient relative error.The full repository suite was not run locally. Repository-wide mypy remains red on the same two pre-existing
autoparallel/tools/overlap_simulator/run.pyerrors as the PR base; changed-file mypy passes.Authored with Claude.
Stack