Repository navigation
[AutoSP] Fix tensor truthiness in get_node_shape_meta - #8760
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>
This was referenced Oct 6, 2026
tohtana
approved these changes
Oct 11, 2026
tohtana
left a comment
tohtana
left a comment
Collaborator
There was a problem hiding this comment.
@jinyouzhi Thank you for the improvement! Looks good to me.
tohtana
enabled auto-merge
October 11, 2026 23:02
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Oct 12, 2026
…ai#8761) Split out of deepspeedai#8457. **This PR is stacked on deepspeedai#8760.** It contains that PR's commit, so please review only the last commit. ## Problem AutoSP shards `input_ids`, `labels` and `position_ids` along the sequence dimension. A model can build its causal or attention mask from the sharded `position_ids`, which gives the mask the shape `[B, 1, S/sp, S]`. After the Ulysses all-to-all, SDPA sees all `S` queries, but the mask still covers only `S/sp` of them. The run then fails with a shape mismatch: `size of tensor a (s) must match ... (s//2)`. ## Fix - Add a custom op `autosp::all_gather_sequence(input, dim)` with a fake (meta) implementation. - In `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. - If the mask's key dimension is also sharded, it cannot be rebuilt correctly from the local pieces, so raise a clear error. - Run `pass_propagate_shapes` once before the attention pass so that the mask shapes are current. This is why the PR depends on the `get_node_shape_meta` fix. - `all_gather_sequence` has a backward pass, an `autosp::reduce_scatter_sequence` op, so masks that require gradients (for example a learnable position bias) also work. - A mask with a full head dimension is redistributed like the query: if its query rows are sharded, it goes through all-to-all; otherwise each rank keeps only its own block of heads. A broadcast key dimension of 1 is never treated as sharded. - Document the behavior under AutoSP Limitations. ## Tests Added `TestAutoSPAttentionMask` (world_size=2): - With a learnable mask bias (shared, per-head with sharded queries, or per-head with full queries), the loss and all parameter gradients match a single-device reference. - With a query-sharded mask, passed as a kwarg and positionally, the forward loss matches a single-device full-sequence reference. - A mask whose key dimension is sharded is rejected with a clear error. Without the fix, all 3 tests fail. With the fix but without the `fx.py` change, all 3 also fail. Hardware: 8× Intel Arc Pro B70 (XPU), torch 2.14.0.dev+xpu, transformers 5.17. 42 AutoSP tests pass. `TestAutoSPCompile` is excluded because of a pre-existing environment failure that also happens on master. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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
2 participants
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.
Problem
get_node_shape_metaindeepspeed/compile/fx.pypicks betweenmeta["val"]andmeta["example_value"]withor. Whenmeta["val"]is a tensor, thatorevaluates the tensor's truthiness. As a result:RuntimeError: Boolean value of Tensor with more than one value is ambiguous.example_value.This only shows up after
pass_propagate_shapes, which runs FakeTensorProp and fills inmeta["val"]. Any AutoSP pass that reads shapes after that step can hit it.Fix
Check
is not Noneexplicitly.Tests
Added
TestNodeShapeMetatotests/unit/v1/compile/test_compile_autosp.py. It covers:example_valueBefore the fix, 2 of these tests fail. After the fix, all 3 pass.
Hardware: 8× Intel Arc Pro B70 (XPU), torch 2.14.0.dev+xpu, transformers 5.17. The rest of the AutoSP unit tests pass.
TestAutoSPCompileis not counted: it fails the same way on master in this environment, in the Dynamo Ulysses baseline.