Repository navigation
[AutoSP] Rebuild query-sharded attention masks before SDPA - #8761
Merged
Merged
Conversation
get_node_shape_meta() used `meta["val"] or meta["example_value"]`, which calls bool() on the stored tensor. Once shape propagation fills meta["val"] with a real tensor, this raises for multi-element tensors and silently falls back to example_value for a zero-valued scalar. Compare with None instead. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
AutoSP shards inputs along the sequence dimension before Ulysses all-to-all. After the all-to-all, SDPA gets the full query sequence, but an attention mask built from sharded position_ids still has only the local query length. That fails with a shape mismatch. Add an autosp::all_gather_sequence op and use it to gather the mask query dimension (-2) over the SP group before the mask reaches SDPA. Masks are gathered once even when several layers share them. A mask whose key dimension is also sharded cannot be rebuilt from local pieces, so raise a clear error for that case. Depends on the get_node_shape_meta fix, because the comparison reads shapes after FakeTensor propagation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Differentiable masks lack autograd support, and valid singleton key dimensions can be incorrectly rejected.
Review effort: Balanced
Findings: 2
Open (2)
What changed in this PR
Adds AutoSP support for rebuilding query-sharded SDPA attention masks after Ulysses redistribution.
Changes:
- Adds sequence all-gather custom operation and mask-rewriting pass.
- Refreshes shape metadata before attention transformation.
- Adds documentation and distributed coverage for positional and keyword masks.
| File | Description |
|---|---|
deepspeed/compile/custom_ops/all_to_all.py |
Adds sequence all-gather operation and fake implementation. |
deepspeed/compile/passes/sp_compile.py |
Detects and gathers query-sharded SDPA masks. |
deepspeed/compile/fx.py |
Includes the stacked shape-metadata fix. |
docs/code-docs/source/training.rst |
Documents attention-mask limitations. |
tests/unit/v1/compile/test_compile_autosp.py |
Adds mask integration and metadata tests. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Address review on the attention mask gather: - Register autograd for autosp::all_gather_sequence. Every rank uses the whole gathered mask, so the backward sums the gradient over the SP group and keeps the local shard. A new autosp::reduce_scatter_sequence op does this, so the backward graph does not capture process groups. - Do not treat a broadcast key dimension of 1 as sharded. - Redistribute masks with a full head dimension the same way as the query: all-to-all them when their query rows are sharded, and otherwise keep only this rank's block of heads. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
SDPA reads a 3D mask as [heads, queries, keys]; add the batch dimension before redistributing its heads. With sequence_parallel_size 1 the mask is already complete, and its key dimension must not be mistaken for a sharded one. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
tohtana
approved these changes
Oct 11, 2026
tohtana
left a comment
tohtana
left a comment
Collaborator
There was a problem hiding this comment.
Thank you for your contributions, @jinyouzhi! This is a significant improvement, and I don't see a clear issue. Let's merge it.
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.
Split out of #8457. This PR is stacked on #8760. It contains that PR's commit, so please review only the last commit.
Problem
AutoSP shards
input_ids,labelsandposition_idsalong the sequence dimension. A model can build its causal or attention mask from the shardedposition_ids, which gives the mask the shape[B, 1, S/sp, S]. After the Ulysses all-to-all, SDPA sees allSqueries, but the mask still covers onlyS/spof them. The run then fails with a shape mismatch:size of tensor a (s) must match ... (s//2).Fix
autosp::all_gather_sequence(input, dim)with a fake (meta) implementation.pass_insert_attention_all_to_all, check whether the mask's query dimension equals the local sequence length. If it does, all-gather the mask along dim -2 over the SP group. The mask can be passed as a kwarg or positionally. When several layers share one mask node, it is gathered only once.pass_propagate_shapesonce before the attention pass so that the mask shapes are current. This is why the PR depends on theget_node_shape_metafix.all_gather_sequencehas a backward pass, anautosp::reduce_scatter_sequenceop, so masks that require gradients (for example a learnable position bias) also work.Tests
Added
TestAutoSPAttentionMask(world_size=2):Without the fix, all 3 tests fail. With the fix but without the
fx.pychange, all 3 also fail.Hardware: 8× Intel Arc Pro B70 (XPU), torch 2.14.0.dev+xpu, transformers 5.17. 42 AutoSP tests pass.
TestAutoSPCompileis excluded because of a pre-existing environment failure that also happens on master.