Skip to content

Do our own call into Python for DTensor dispatch#166370

Closed
swolchok wants to merge 17 commits intogh/swolchok/862/basefrom
gh/swolchok/862/head
Closed

Do our own call into Python for DTensor dispatch#166370
swolchok wants to merge 17 commits intogh/swolchok/862/basefrom
gh/swolchok/862/head

Conversation

@swolchok
Copy link
Contributor

@swolchok swolchok commented Oct 28, 2025

Stack from ghstack (oldest at bottom):

Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci

Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

[ghstack-poisoned]
@pytorch-bot
Copy link

pytorch-bot bot commented Oct 28, 2025

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

❌ 15 New Failures

As of commit 73ac204 with merge base 84776e1 (image):

NEW FAILURES - The following jobs have failed:

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

…spatch"

Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

[ghstack-poisoned]
Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Oct 28, 2025
Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

ghstack-source-id: 1327e9f
Pull Request resolved: #166370
Comment on lines 774 to 776
// REVIEW: without the device change here,
// test_dtensor_save_load breaks. I'm not at all sure why it was
// OK before.
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

based on test failures, changing the device to meta here clearly isn't correct . will go back.

Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Oct 29, 2025
Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

ghstack-source-id: c5cfbf4
Pull Request resolved: #166370
@swolchok
Copy link
Contributor Author

running into some weird coupling here. test_dtensor_save_load is broken. (repro: python test/distributed/tensor/test_dtensor.py -k test_dtensor_save_load, you need to patch the test to reduce world_size to 1 if you have only one GPU)
Apparently, the presence of DTensor.__torch_dispatch__ somehow causes DTensors' storage to have a valid data_ptr and thus be compatible with serialization. This PR deletes DTensor.__torch_dispatch__ and breaks it; I know it's the mere presence of DTensor.__torch_dispatch__ because 1) putting it back fixes the test 2) printing inside DTensor.__torch_dispatch__ shows that it's not called before or during serialization.

@swolchok swolchok added the topic: not user facing topic category label Oct 29, 2025
…l into Python for DTensor dispatch"

Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Oct 29, 2025
Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

ghstack-source-id: d19b694
Pull Request resolved: #166370
Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Oct 29, 2025
Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

ghstack-source-id: 45022fc
Pull Request resolved: #166370
…o our own call into Python for DTensor dispatch"

Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Oct 29, 2025
Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

ghstack-source-id: 8a4b52c
Pull Request resolved: #166370
@swolchok
Copy link
Contributor Author

Current failure looks like Dynamo needs to know about the DTensor dispatch key somehow; we're hitting __torch_dispatch__ path for a DTensor (AIUI it's actually a FakeTensor wrapping a DTensor somehow, so perhaps the way the FakeTensor is created is wrong? maybe we can make that impossible?)

torch/_tensor.py Outdated
Comment on lines 470 to 473
and (
type(self).__torch_dispatch__ is not torch.Tensor.__torch_dispatch__
or _is_dtensor(self)
)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not sure how to structurally avoid is_dtensor special cases in N things that currently just care about __torch_dispatch__

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thinking about whether we put DTensor.torch_dispatch back and have it just upcall to the DTensor dispatch key stuff

…call into Python for DTensor dispatch"

Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 5, 2025
This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

ghstack-source-id: 0937067
Pull Request resolved: #167051
swolchok added a commit that referenced this pull request Nov 6, 2025
…dd C++ fast path for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 6, 2025
This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

ghstack-source-id: 2f0ea67
Pull Request resolved: #167051
swolchok added a commit that referenced this pull request Nov 6, 2025
This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

ghstack-source-id: b3451db
Pull Request resolved: #167051
swolchok added a commit that referenced this pull request Nov 6, 2025
…r `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 6, 2025
…ch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 6, 2025
…n "Add C++ fast path for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 6, 2025
…path for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 6, 2025
This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

ghstack-source-id: 672fa5d
Pull Request resolved: #167051
swolchok added a commit that referenced this pull request Nov 6, 2025
…or.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 6, 2025
…tch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
@swolchok
Copy link
Contributor Author

swolchok commented Nov 6, 2025

superseded by #167051

@swolchok swolchok closed this Nov 6, 2025
swolchok added a commit that referenced this pull request Nov 7, 2025
… path for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 7, 2025
…or.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 10, 2025
…patch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 10, 2025
This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 12, 2025
…th for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 12, 2025
…__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 12, 2025
…nd add a couple comments on "Add C++ fast path for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 12, 2025
…nd add a couple comments on "Add C++ fast path for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 12, 2025
…comments on "Add C++ fast path for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 12, 2025
… gating in is_dtensor on "Add C++ fast path for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 12, 2025
…ensor on "Add C++ fast path for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 13, 2025
…C++ fast path for `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Nov 13, 2025
…r `DTensor.__torch_dispatch__`"

This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta msaroufim dcci

[ghstack-poisoned]
pytorchmergebot pushed a commit that referenced this pull request Nov 13, 2025
This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike #166370 and #166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of #166370.

Pull Request resolved: #167051
Approved by: https://github.com/ezyang
Khanaksahu pushed a commit to Khanaksahu/pytorch that referenced this pull request Nov 17, 2025
Another incremental step: don't deal with custom ops on the critical
path, let the dispatcher do that. Take control of calling into Python
on the critical path as well.

ghstack-source-id: 453ddca
Pull Request resolved: pytorch/pytorch#166370
Silv3S pushed a commit to Silv3S/pytorch that referenced this pull request Nov 18, 2025
This patches the `__torch_dispatch__` machinery to detect DTensor and hand over control to a C++ fast path. Unlike pytorch#166370 and pytorch#166369 (which added a DTensor dispatch key and are intended to be replaced by this PR), this approach fundamentally *is* `__torch_dispatch__`, hopefully sidestepping all manner of thorny "does it work just like `__torch_dispatch__`?" that came up during development and review of pytorch#166370.

Pull Request resolved: pytorch#167051
Approved by: https://github.com/ezyang
@github-actions github-actions bot deleted the gh/swolchok/862/head branch December 7, 2025 02:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/inductor oncall: distributed Add this issue/PR to distributed oncall triage queue topic: not user facing topic category

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants