convert stride_per_key_per_rank to tensor inside KJT #2959
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
facebook-github-bot
added
the
CLA Signed
|
This pull request was exported from Phabricator. Differential Revision: D74366343 |
TroyGarden
added a commit
to TroyGarden/torchrec
that referenced
this pull request
May 8, 2025
Summary: # context * this diff is part of the "variable-batch KJT refactoring" project ([doc](https://fburl.com/gdoc/svfysfai)) * previously the `stride_per_key_per_rank` variable is `List[List[int]] | None` which can't be handled correctly in PT2 IR (torch.export) * this change makes the KJT class variable `_stride_per_key_per_rank` as `torch.IntTensor | None` so it would be compatible with PT2 IR. # equivalency * to check if `self._stride_per_key_per_rank` is `None` this logic is used to differentiate variable_batch case, and should have the same behavior after this diff * to use `self._stride_per_key_per_rank` as `List[List[int]]` most of the callsite use the function to get the list: `def stride_per_key_per_rank(self) -> List[List[int]]:`, and this function is modified to covert the `torch.IntTensor` to list as ` _stride_per_key_per_rank.tolist()`, the results should be the same NOTE: this `self. _stride_per_key_per_rank.tolist()` tensor should always be on CPU since it's effective the meta data of a KJT. For generic torch APIs like `.to(...)`, `record_stream()`, etc. should in general avoid altering this variable. Differential Revision: D74366343
TroyGarden
force-pushed
the
export-D74366343
branch
from
d9ba5cf to
1ce1f10
Compare
|
This pull request was exported from Phabricator. Differential Revision: D74366343 |
TroyGarden
force-pushed
the
export-D74366343
branch
from
1ce1f10 to
c195982
Compare
TroyGarden
added a commit
to TroyGarden/torchrec
that referenced
this pull request
May 8, 2025
Summary: # context * this diff is part of the "variable-batch KJT refactoring" project ([doc](https://fburl.com/gdoc/svfysfai)) * previously the `stride_per_key_per_rank` variable is `List[List[int]] | None` which can't be handled correctly in PT2 IR (torch.export) * this change makes the KJT class variable `_stride_per_key_per_rank` as `torch.IntTensor | None` so it would be compatible with PT2 IR. # equivalency * to check if `self._stride_per_key_per_rank` is `None` this logic is used to differentiate variable_batch case, and should have the same behavior after this diff * to use `self._stride_per_key_per_rank` as `List[List[int]]` most of the callsite use the function to get the list: `def stride_per_key_per_rank(self) -> List[List[int]]:`, and this function is modified to covert the `torch.IntTensor` to list as ` _stride_per_key_per_rank.tolist()`, the results should be the same NOTE: this `self. _stride_per_key_per_rank.tolist()` tensor should always be on CPU since it's effective the meta data of a KJT. For generic torch APIs like `.to(...)`, `record_stream()`, etc. should in general avoid altering this variable. Reviewed By: jd7-tr Differential Revision: D74366343
|
This pull request was exported from Phabricator. Differential Revision: D74366343 |
Summary: # context * this diff is part of the "variable-batch KJT refactoring" project ([doc](https://fburl.com/gdoc/svfysfai)) * previously the `stride_per_key_per_rank` variable is `List[List[int]] | None` which can't be handled correctly in PT2 IR (torch.export) * this change makes the KJT class variable `_stride_per_key_per_rank` as `torch.IntTensor | None` so it would be compatible with PT2 IR. # equivalency * to check if `self._stride_per_key_per_rank` is `None` this logic is used to differentiate variable_batch case, and should have the same behavior after this diff * to use `self._stride_per_key_per_rank` as `List[List[int]]` most of the callsite use the function to get the list: `def stride_per_key_per_rank(self) -> List[List[int]]:`, and this function is modified to covert the `torch.IntTensor` to list as ` _stride_per_key_per_rank.tolist()`, the results should be the same NOTE: this `self. _stride_per_key_per_rank.tolist()` tensor should always be on CPU since it's effective the meta data of a KJT. For generic torch APIs like `.to(...)`, `record_stream()`, etc. should in general avoid altering this variable. Reviewed By: jd7-tr Differential Revision: D74366343
TroyGarden
force-pushed
the
export-D74366343
branch
from
c195982 to
7b44f11
Compare
|
This pull request was exported from Phabricator. Differential Revision: D74366343 |
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.
Summary:
context
stride_per_key_per_rankvariable isList[List[int]] | Nonewhich can't be handled correctly in PT2 IR (torch.export)_stride_per_key_per_rankastorch.IntTensor | Noneso it would be compatible with PT2 IR.equivalency
self._stride_per_key_per_rankisNonethis logic is used to differentiate variable_batch case, and should have the same behavior after this diff
self._stride_per_key_per_rankasList[List[int]]most of the callsite use the function to get the list:
def stride_per_key_per_rank(self) -> List[List[int]]:, and this function is modified to covert thetorch.IntTensorto list as_stride_per_key_per_rank.tolist(), the results should be the sameNOTE: this
self. _stride_per_key_per_rank.tolist()tensor should always be on CPU since it's effective the meta data of a KJT. For generic torch APIs like.to(...),record_stream(), etc. should in general avoid altering this variable.Differential Revision: D74366343