[dtensor] skip pytree when not necessary #110132

wanchaol · 2023-09-27T02:20:03Z

Stack from ghstack (oldest at bottom):

-> [dtensor] skip pytree when not necessary #110132

pytree is a great tool, but it sometimes considers to be evil for
tensor subclasses, it's useful to implement subclass quickly, but it:

exposes non-trival CPU overhead
many ops don't need pytree, only the one with list/dict ops needs
blindly use pytree to re-wrap have semantic issues for inplace/out
ops

This PR avoid using pytree for most ops during torch_dispatch and only
enable it for certain ops

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops [ghstack-poisoned]

pytorch-bot · 2023-09-27T02:20:05Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/110132

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 9be485d with merge base 40b83d9 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops ghstack-source-id: f35895d80052e81f37d539a6413e5da3d30378e8 Pull Request resolved: #110132

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops [ghstack-poisoned]

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops ghstack-source-id: 72f3370dd99e467b9647655f825efa0c1b8d9d35 Pull Request resolved: #110132

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops [ghstack-poisoned]

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops ghstack-source-id: d622e08c6218bb192d675bae2465cee7a6e0b4f8 Pull Request resolved: #110132

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops [ghstack-poisoned]

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops ghstack-source-id: 2c399f73c2b056d5dd5e0db22e3a852fa2b14852 Pull Request resolved: #110132

torch/distributed/_tensor/dispatch.py

fduwjj

LGTM, it's actually surprising to me that only two ops need treemap here..

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops [ghstack-poisoned]

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops ghstack-source-id: d506fdc7a843bc094e069316803001ab1abdb298 Pull Request resolved: #110132

Chillee · 2023-09-28T15:54:14Z

Have you tried using optree?

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops [ghstack-poisoned]

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops ghstack-source-id: f2947786ac40ae350eb712fbf48c8c7b322fcac2 Pull Request resolved: #110132

wanchaol · 2023-09-28T16:02:15Z

Have you tried using optree?

Oh just noticed that got landed!! I can try turn that on, is there a flag to turn it on?

My feeling is that even if that landed it would still have some overhead exposed (especially optree perf is most amazing on dicts, while operator arguments does not have too many dicts, maybe kwargs is a good try). This PR try to turn pytree/optree off completely for most ops which I thought it would be faster for either case?

Chillee · 2023-09-28T16:05:45Z

@wanchaol might be a little bit annoying to actually land but the idea was that we would incrementally port usages over incrementally.

wanchaol · 2023-09-28T16:19:03Z

@wanchaol might be a little bit annoying to actually land but the idea was that we would incrementally port usages over incrementally.

@Chillee sounds good, I'll give it a try in a follow up PR? the annoying part is that we need to guard the imports mostly so that all CI passes?

Chillee · 2023-09-28T16:23:08Z

Yeah, see #109684

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops [ghstack-poisoned]

pytree is a great tool, but it sometimes considers to be evil for tensor subclasses, it's useful to implement subclass quickly, but it: * exposes non-trival CPU overhead * many ops don't need pytree, only the one with list/dict ops needs * blindly use pytree to re-wrap have semantic issues for inplace/out ops This PR avoid using pytree for most ops during torch_dispatch and only enable it for certain ops ghstack-source-id: dfddc15534ce75369e914b7875511d991b29e762 Pull Request resolved: #110132

wanchaol · 2023-10-02T17:42:16Z

@pytorchbot merge

pytorchmergebot · 2023-10-02T17:44:10Z

Merge started

Your 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

Advanced Debugging

Check the merge workflow status
here

wanchaol requested review from mrshenli, zhaojuanmao, rohan-varma, H-Huang, awgu, kwen2501, fegin, fduwjj, kiukchung, d4l3k and wz337 as code owners September 27, 2023 02:20

fduwjj reviewed Sep 28, 2023

View reviewed changes

torch/distributed/_tensor/dispatch.py Show resolved Hide resolved

fduwjj approved these changes Sep 28, 2023

View reviewed changes

wanchaol requested review from ezyang, zou3519 and Chillee September 28, 2023 15:52

wanchaol added ciflow/trunk Trigger trunk jobs on your pull request release notes: distributed (dtensor) release notes category labels Sep 28, 2023

pytorchmergebot added the merging label Oct 2, 2023

pytorchmergebot added Merged and removed merging labels Oct 2, 2023

pytorchmergebot closed this in 26900d2 Oct 2, 2023

facebook-github-bot deleted the gh/wanchaol/357/head branch October 6, 2023 14:23

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[dtensor] skip pytree when not necessary #110132

[dtensor] skip pytree when not necessary #110132

wanchaol commented Sep 27, 2023 •

edited

pytorch-bot bot commented Sep 27, 2023 •

edited

fduwjj left a comment

Chillee commented Sep 28, 2023

wanchaol commented Sep 28, 2023

Chillee commented Sep 28, 2023

wanchaol commented Sep 28, 2023

Chillee commented Sep 28, 2023

wanchaol commented Oct 2, 2023

pytorchmergebot commented Oct 2, 2023

[dtensor] skip pytree when not necessary #110132

[dtensor] skip pytree when not necessary #110132

Conversation

wanchaol commented Sep 27, 2023 • edited

pytorch-bot bot commented Sep 27, 2023 • edited

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/110132

✅ No Failures

fduwjj left a comment

Choose a reason for hiding this comment

Chillee commented Sep 28, 2023

wanchaol commented Sep 28, 2023

Chillee commented Sep 28, 2023

wanchaol commented Sep 28, 2023

Chillee commented Sep 28, 2023

wanchaol commented Oct 2, 2023

pytorchmergebot commented Oct 2, 2023

Merge started

wanchaol commented Sep 27, 2023 •

edited

pytorch-bot bot commented Sep 27, 2023 •

edited