Adding remove_caches API to Float8Tensor class #1425
Conversation
This file contains 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
Signed-off-by: Youngeun Kwon <[email protected]>
|
/te-ci pytorch L1 |
timmoon10
approved these changes
Jan 28, 2025
| @@ -919,6 +919,14 @@ def _reset_caches(self) -> None: | |||
| """ | |||
| self._transpose_invalid = True | |||
|
|
|||
| def remove_caches(self) -> None: | |||
There was a problem hiding this comment.
Nit: I'd prefer if we treated this as an internal function and not a public API that can be relied on.
Suggested change
| def remove_caches(self) -> None: | |
| def _remove_caches(self) -> None: |
There was a problem hiding this comment.
Hi @timmoon10, ADLR asked for a public API that is able to remove the transpose cache, which is the reason I am working on this PR. I guess it will be better to have public API to meet their request if you are also okay with that.
Co-authored-by: Tim Moon <[email protected]> Signed-off-by: Youngeun Kwon <[email protected]>
Thank you for the review! I will wait for the 2.0 release to be merged. Can I know a rough ETA of that? |
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.
Description
Add remove_caches API for Float8Tensor.
This API frees the allocated memory for the transpose cache.
This API is required for FSDP use cases.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: