[dtensor] directly return local_tensor under no_grad #128145
Conversation
as titled, skip the autograd function and directly return the local_tensor if it's under no_grad context, this would avoid creating views [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/128145
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (6 Unrelated Failures)As of commit 7a1cfdf with merge base 2fc9079 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
ciflow/inductor
oncall: distributed
| .. note:: `to_local` is differentiable, the `requires_grad` of the local tensor returned | ||
| will depend on if the `DTensor` requires_grad or not. | ||
| """ | ||
| if not torch.is_grad_enabled(): |
There was a problem hiding this comment.
Now the question is, how expensive is torch.is_grad_enabled() 😆
There was a problem hiding this comment.
I hope it's pretty cheap given that we are using it invasively in pytorch ;)
There was a problem hiding this comment.
tested and got a trace locally, I didn't even see the to_local or is_grad_enabled calls in the CPU trace :)
There was a problem hiding this comment.
I guess that's because there is no aten op in that call stack, so the profiler cannot see it? (but this could still contribute to gaps between other ops in the profiler trace?)
|
@pytorchbot merge |
Merge failedReason: This PR needs a If not, please add the To add a label, you can comment to pytorchbot, for example For more information, see Details for Dev Infra teamRaised by workflow job |
|
@pytorchbot merge |
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 |
|
@pytorchbot merge -f "queued job" |
|
The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command |
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 |
as titled, skip the autograd function and directly return the local_tensor if it's under no_grad context, this would avoid creating views Pull Request resolved: pytorch#128145 Approved by: https://github.com/awgu ghstack dependencies: pytorch#128112
3 participants
Stack from ghstack (oldest at bottom):
as titled, skip the autograd function and directly return the
local_tensor if it's under no_grad context, this would avoid creating
views
cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k