DTensor: avoid unnecessary DTensorSpec creation in _ToTorchTensor.backward #167588
Conversation
pytorch-bot
bot
added
ciflow/inductor
oncall: distributed
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/167588
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 035049e with merge base 573a79f ( UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
[ghstack-poisoned]
ezyang
left a comment
ezyang
left a comment
There was a problem hiding this comment.
Huh, I'm surprised this actual is an improvement lol
| ): | ||
| # Avoid actual sharing of specs in case they're modified during (e.g.) | ||
| # sharding propagation. | ||
| grad_spec = copy.copy(dtensor_spec) |
There was a problem hiding this comment.
Is copying really that much faster?
copy in this case I would expect is actually doing a pickle serialization/deserialization of the object to get a new copy.
Is that really faster than creating a clean object below?
There was a problem hiding this comment.
Is that really faster than creating a clean object below?
Yes. DTensorSpec does some relatively expensive computation in its constructor
|
@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 |
…kward (pytorch#167588) Looks like the check here is cheap and has a potentially large payoff. Pull Request resolved: pytorch#167588 Approved by: https://github.com/ezyang
Stack from ghstack (oldest at bottom):
Looks like the check here is cheap and has a potentially large payoff.
cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci