Repository navigation
Conversation
_validate_holders compared every piece with every other piece, and _sources compared every source region with every other region, so a map with 1024 pieces already exceeded the default max_pair_checks and planning time grew quadratically. When no rank repeats a box and distinct boxes do not overlap, a rank covers a box exactly when it holds that very box. The holders of each box can then be read from a table. Add boxes_disjoint to affine_geometry: it sweeps along the most varied axis and compares a box only with the boxes still open where it starts. Maps that fail the fast checks take the original pairwise walk, which names the offending piece, so error messages are unchanged. A randomized test checks that plans and errors match the pairwise walk, including replicas split at different points, overlapping pieces and wrong holder lists. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
A rank whose pieces sweep as disjoint in its shard and fill it passes _shards_tile with boxes_disjoint; any other rank takes the pairwise walk, so error messages are unchanged. Maps that pass no longer hit the per-rank len(held)**2 budget. Suggested-by: 0z5a Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
A disjoint layout can keep every box open for the whole sweep, so boxes_disjoint was still quadratic without a budget of its own. Give it an optional max_checks: once exceeded it reports the boxes as not shown disjoint, and since the sweep compares at most n*(n-1)/2 pairs, the caller's pairwise fallback under the same budget refuses at once. The planner passes max_pair_checks at every call. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
A ZeRO-3 partition cut at a tensor-parallel column edge holds boxes narrower than the parameter, stored at the parameter's own row-major strides. Their address intervals interleave, so plan_transfer refused them, and the only workaround was one piece per row, which needs a larger budget. Accept them when a flat shard keeps one shift for all its pieces: the address of logical row-major index i is then i + shift, which is injective, so disjoint logical boxes, the shard bounds and the volume sum prove the shard is tiled. Other mixes stay refused. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a6feceda4c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| hold a multi-axis box as a dense run, which is how ZeRO stores a tensor-parallel shard, or at the | ||
| parameter's own row-major strides, which is how a ZeRO partition holds a box cut out of the rows it | ||
| stores. ``_shards_tile`` proves the second kind by its shift, not by its address interval. |
There was a problem hiding this comment.
When callers build these parameter-stride boxes, the checked-in contract still tells them the opposite: deepspeed/checkpoint/affine_ir_spec.md:671-688 says every box in a 1-D shard is dense, is treated as [dest_offset, dest_offset + numel), and that a narrower block must remain one piece per row. Following that guidance recreates the max_pieces_per_map failure this change is intended to remove, so update the storage-map contract and derivation guidance alongside this new behavior.
AGENTS.md reference: AGENTS.md:L71-L71
Useful? React with 👍 / 👎.
| monkeypatch.setattr(transfer_module, '_validate_holders', transfer_module._validate_holders_pairwise) | ||
| assert fast == _outcome(target, source) |
There was a problem hiding this comment.
This test replaces the production _validate_holders binding with another private helper and treats the exact plan or error text from that implementation as its oracle. A correct refactor that removes or renames either helper, or changes diagnostics without changing plan_transfer behavior, will fail this test, while a defect shared by both paths can still pass; exercise the public planner against independently derived valid/invalid maps or observable transfer results instead.
AGENTS.md reference: AGENTS.md:L76-L80
Useful? React with 👍 / 👎.
gpu-test-labeler
Bot
added
the
skip-gpu-test
Stacked on #9011: only the last commit (a6feced) is new here. I'll rebase onto master once #9011 merges. Follow-up to the column case in #9026.
A ZeRO-3 partition cut at a tensor-parallel column edge holds boxes narrower than the parameter, and it stores them at the parameter's own row-major strides (
(cols, 1)), not as dense runs.plan_transferrefuses those today (dest axis 0 strides …). The only workaround is one piece per row, which runs intomax_pieces_per_mapfor real weights.Change. In
_frame_of, a flat shard may now hold a piece at the parameter's own strides, alongside the dense run it already accepts. Each such piece records its shift: its shard offset minus the row-major index of its origin. The tiling check for these pieces is a separate branch,_in_order_shard_tiles:isits at addressi + shift, which is injective.max_pair_checks, the same as in [AffineIR] Check affine holders and source regions without pairwise comparisons #9011.Shards whose pieces are all dense runs take the existing path unchanged.
Tests (
tests/unit/checkpoint/test_affine_transfer_plan.py, all single-process):Without the planner change, the 5 acceptance cases are refused. The 5 refusal cases refuse either way; the tests pin the new reasons. On CPU, every affine test file under
tests/unit/checkpoint/passes (716 tests), as doestests/unit/v1/moe/test_autoep_checkpoint.py(45 passed, 4 skipped).Outside the unit tests (CPU, against this branch). These use the two-layout map from #9026: ZeRO-3 partitions plus the surviving TP shards, cut at each other's edges.
execute_transferson 4 gloo processes, in both directions, at staging budgets of 8 MiB and 256 B. Nothing is wrong: 0 of 200 target shards reading from the cut map, and 0 of 483 writing into it.