Skip to content

Fix int32 overflow in embedding_dense_backward#165095

Closed
yinghai wants to merge 8 commits intopytorch:mainfrom
yinghai:yinghai/fix_emb_bwd
Closed

Fix int32 overflow in embedding_dense_backward#165095
yinghai wants to merge 8 commits intopytorch:mainfrom
yinghai:yinghai/fix_emb_bwd

Conversation

@yinghai
Copy link
Contributor

@yinghai yinghai commented Oct 9, 2025

If max_partial_segment is large we can overflow gid and cause a bunch of IMA.

@pytorch-bot pytorch-bot bot added the release notes: cuda release notes category label Oct 9, 2025
@pytorch-bot
Copy link

pytorch-bot bot commented Oct 9, 2025

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/165095

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit b7fb34d with merge base a2f29bc (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@yinghai yinghai changed the title FIx int32 overflow in embedding_dense_backward Fix int32 overflow in embedding_dense_backward Oct 9, 2025
@ngimel
Copy link
Collaborator

ngimel commented Oct 9, 2025

Test?

)

@onlyCUDA
@dtypes(torch.bfloat16,)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

decorate is with @largeTensorTest to avoid ooms

@ngimel ngimel added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 9, 2025
@pytorch-bot pytorch-bot bot removed the ciflow/trunk Trigger trunk jobs on your pull request label Oct 9, 2025
@Aidyn-A Aidyn-A added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 9, 2025
@pytorch-bot pytorch-bot bot removed the ciflow/trunk Trigger trunk jobs on your pull request label Oct 9, 2025
const int gid = blockIdx.x * blockDim.x + threadIdx.x;
const int id = gid / stride_warped;
const int startFeature = gid % stride_warped;
const int64_t gid = blockIdx.x * blockDim.x + threadIdx.x;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One of these values needs to be cast to 64-bit, otherwise the math is still done in 32bit
see e.g., #142010 #164705

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks updated.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ugh interesting that the test isn't catching it

@yinghai yinghai requested a review from eqy October 10, 2025 03:12
@ngimel ngimel added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 10, 2025
@pytorch-bot pytorch-bot bot removed the ciflow/trunk Trigger trunk jobs on your pull request label Oct 10, 2025
@Aidyn-A Aidyn-A added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 10, 2025
@ngimel
Copy link
Collaborator

ngimel commented Oct 10, 2025

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@yinghai yinghai deleted the yinghai/fix_emb_bwd branch October 10, 2025 22:10
Chao1Han pushed a commit to Chao1Han/pytorch that referenced this pull request Oct 21, 2025
If `max_partial_segment` is large we can overflow `gid` and cause a bunch of IMA.
Pull Request resolved: pytorch#165095
Approved by: https://github.com/ngimel, https://github.com/eqy
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/trunk Trigger trunk jobs on your pull request Merged open source release notes: cuda release notes category

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants