[aotd] Unwrap unseen AsyncCollectiveTensor tangents #138731
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
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/138731
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit c46cf0c with merge base 10a34dc ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
[ghstack-poisoned]
bdhirsh
reviewed
Oct 24, 2024
| return x | ||
|
|
||
| is_subclass: bool = is_traceable_wrapper_subclass(x) | ||
| mem_format = memory_format[0] if is_subclass else memory_format |
Contributor
There was a problem hiding this comment.
So right now we have two separate functions for doing "coercion" of tangents:
(1) coerce non-contiguous tangents to be contiguous
(2) coerce tangents with incorrect subclass metadata to have the right subclass metadata
It seems strange to add the logic for handling AsyncCollectiveTensor to the first one. So I would rather put it in the second (which might require tweaking the contiguous() coercion to run later), or in the worst case, merging the two pieces of coercion code.
bdhirsh
reviewed
Oct 24, 2024
| memory_format_for_dense_tensor = not isinstance(memory_format, list) | ||
| if isinstance(x, AsyncCollectiveTensor) and memory_format_for_dense_tensor: | ||
| # This is AsyncCollectiveTensor, that we have not seen during tracing time. | ||
| while True: |
Contributor
There was a problem hiding this comment.
hmm, I don't think it's possible to end up with an AsyncCollectiveTensor(AsyncCollectiveTensor(...) (that would probably a bug with the subclass itself). So I'd rather just not have a while loop here. Is there a place where you saw this show up?
Collaborator
|
I don't really like that we need to shove all of this stuff into AOTAutograd. Can we register this with the subclass in some manner? |
bdhirsh
reviewed
Oct 24, 2024
| mem_format = memory_format | ||
| if is_subclass: | ||
| memory_format_for_dense_tensor = not isinstance(memory_format, list) | ||
| if isinstance(x, AsyncCollectiveTensor) and memory_format_for_dense_tensor: |
Contributor
There was a problem hiding this comment.
once you do the refactor to merge all of the tangent coercion logic into one function, you should be able to use CompiledFunction.metadata.traced_tangent_metas to know when you need to do a wait() call here instead.
Contributor
Author
There was a problem hiding this comment.
Yes, we should use subclass_type information
bdhirsh
reviewed
Oct 24, 2024
test/functorch/test_aotdispatch.py
Outdated
| ga, gb, gc = torch.autograd.grad(out_buffer.sum(), (a, b, c)) | ||
|
|
||
| def test_unwrap_async_collective_tensor_tangent(self): | ||
| class WrapSC(torch.Tensor): |
Contributor
There was a problem hiding this comment.
nit: it's starting to feel like we have a lot of one-off subclasses across our tests. Any reason not to just re-use TwoTensor here?
Contributor
Author
There was a problem hiding this comment.
nit: it's starting to feel like we have a lot of one-off subclasses across our tests. Any reason not to just re-use
TwoTensorhere?
I was thinking to extract it to some utils module. Will try to use TwoTensor, I think it should work, I will try.
Contributor
Author
There was a problem hiding this comment.
TwoTensor works, removed redundant subclass in the test.
bdhirsh
approved these changes
Oct 24, 2024
Contributor
bdhirsh
left a comment
bdhirsh
left a comment
There was a problem hiding this comment.
Stamping to unblock Andrew, agreed offline that we should do a refactor to merge all of the tangent coercion logic into one place (it's a bit weird to be doing AsyncCollectiveTensor unwrapping right now in the code that is supposed to make tangents contiguous)
[ghstack-poisoned]
Contributor
Author
|
@pytorchbot merge |
Collaborator
Merge startedYour 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 |
Collaborator
Merge failedReason: 1 mandatory check(s) failed. The first few are: Dig deeper by viewing the failures on hud |
Contributor
Author
|
@pytorchbot merge -f "skip unrelated bazel build failures" |
Collaborator
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
IvanKobzarev
added a commit
that referenced
this pull request
Oct 29, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 29, 2024
…h expected_type for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 29, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 29, 2024
…h expected_type for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 29, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 29, 2024
…h expected_type for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 29, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 29, 2024
…h expected_type for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 29, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 30, 2024
…h expected_type for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 30, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 31, 2024
…h expected_type for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 31, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 31, 2024
…h expected_type for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 31, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 31, 2024
…h expected_type for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Oct 31, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Nov 1, 2024
…h expected_type for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Nov 1, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
Contributor
Author
|
Abandoning in favor of #139095 |
IvanKobzarev
added a commit
that referenced
this pull request
Nov 6, 2024
…h expected_type for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
IvanKobzarev
added a commit
that referenced
this pull request
Nov 6, 2024
…for e.g.AsyncCollectiveTensor" Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
pytorchmergebot
pushed a commit
that referenced
this pull request
Nov 7, 2024
…ncCollectiveTensor (#139095) Based on discussion here: #138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. Pull Request resolved: #139095 Approved by: https://github.com/bdhirsh
pobin6
pushed a commit
to pobin6/pytorch
that referenced
this pull request
Dec 5, 2024
…ncCollectiveTensor (pytorch#139095) Based on discussion here: pytorch#138731 Introducing ability for subclass implement type convertion to expected_type. ``` def __coerce_same_metadata_as_tangent__( self, expected_metadata: Any, expected_type: Optional[Type] = None ): ``` Here if `expected_type=None` means `SubclassClass` is expected. E.g. for `DTensor` we may find tangent `AsyncCollectiveTensor` where we expected `Tensor` - in this case `expected_type=Tensor` will be called during runtime Adding implementation to AsyncCollectiveTensor, that just triggers `wait()`. Pull Request resolved: pytorch#139095 Approved by: https://github.com/bdhirsh
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.
Stack from ghstack (oldest at bottom):