Fix IMA in TBE grad indices kernel for int32 indices #3877
Closed
Conversation
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
|
This pull request was exported from Phabricator. Differential Revision: D71796826 |
✅ Deploy Preview for pytorch-fbgemm-docs ready!
To edit notification comments on pull requests, go to your Netlify site configuration. |
Summary: This diff forces casting the weight tensor index to `overflow_safe_int_t` whichi is int64_t to address the 32-bit int overflow issue when the TBE `indices` is int32_t. Before this diff, `idx_j` and `D_emb` are both int32_t when using 32-bit int `indices`. When accessing the embedding `weights` tensor, we computed an index by multiplying `idx_j` and `D_emb` together. Their product could be larger than the max value of int32_t. This led to an integer overflow problem, resulting in illegal memory access. By forcing the `idx_j` and `D_emb` to be int64_t, we can prevent the 32-bit int overflow problem. Note that using int64_t is safe since its max value is much larger than the memory sizes of modern GPUs and CPUs. We also add unit tests for this issue. X-link: facebookresearch/FBGEMM#967 **Facebook:** This is the fix for S498528. The full root cause details are in https://fburl.com/gdoc/lhmvenw3. Reviewed By: brad-mengchi, spcyppt Differential Revision: D71796826
|
This pull request was exported from Phabricator. Differential Revision: D71796826 |
|
This pull request has been merged in 6a6db7c. |
q10
pushed a commit
to q10/FBGEMM
that referenced
this pull request
Apr 10, 2025
Summary: This diff forces casting the weight tensor index to `overflow_safe_int_t` whichi is int64_t to address the 32-bit int overflow issue when the TBE `indices` is int32_t. Before this diff, `idx_j` and `D_emb` are both int32_t when using 32-bit int `indices`. When accessing the embedding `weights` tensor, we computed an index by multiplying `idx_j` and `D_emb` together. Their product could be larger than the max value of int32_t. This led to an integer overflow problem, resulting in illegal memory access. By forcing the `idx_j` and `D_emb` to be int64_t, we can prevent the 32-bit int overflow problem. Note that using int64_t is safe since its max value is much larger than the memory sizes of modern GPUs and CPUs. We also add unit tests for this issue. X-link: pytorch#3877 Pull Request resolved: facebookresearch/FBGEMM#967 **Facebook:** This is the fix for S498528. The full root cause details are in https://fburl.com/gdoc/lhmvenw3. Reviewed By: brad-mengchi, spcyppt Differential Revision: D71796826 fbshipit-source-id: df7c8e06d36e2f06585a5812df8ae40863ea6253
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.
Differential Revision: D71796826