Skip to content

Speed up sharding optimizer build#513

Open
AlbedoWang wants to merge 1 commit into
mainfrom
kaijian/fast-build-safe
Open

Speed up sharding optimizer build#513
AlbedoWang wants to merge 1 commit into
mainfrom
kaijian/fast-build-safe

Conversation

@AlbedoWang

Copy link
Copy Markdown

Summary

This PR reduces AutoParallel sharding optimizer build time while preserving the solver-visible optimization problem and selected placements.

The main changes are:

  • Skip enumeration-time DTensor redistribute_cost work during sharding metadata construction. Those costs are not consumed by the solver directly; AutoParallel recomputes and stores its own edge costs before building the ILP objective.
  • Avoid creating PuLP variables for infinite-cost strategy edges. These edges were previously represented as variables and then constrained to zero; now they are omitted before variable creation while constraints treat missing variables as zero.
  • Store repeated-subgraph cluster links at node granularity instead of per (node, arg, output, input) option. Identical repeated nodes resolve to their root node using matching option indices, reducing cluster bookkeeping size.
  • Add deterministic parallel edge-cost construction controlled by AP_PARALLEL_BUILD, with serial fallback when CUDA has already been initialized.
  • Add AP_FAST_BUILD=0 as an A/B switch for disabling the enumeration-time shortcut.
  • Update optimizer serialization to preserve the new node-level cluster links and valid decision-variable key set.
  • Add user-facing and design docs for fast optimizer build.
  • Update tests to cover public end-to-end behavior: fast/reference placement equivalence, JSON export behavior, serialization round trips, and dtype-cast placement behavior.

This reduces optimizer construction overhead while preserving the solver-visible problem. The build now skips discarded enumeration-time redistribute costs, prunes invalid edges before PuLP variable creation, stores repeated-subgraph cluster links at node granularity, and can construct edge costs in deterministic worker order.

Authored with Claude.
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant