Fully native DTensor.new #162508

swolchok · 2025-09-09T17:44:41Z

Stack from ghstack (oldest at bottom):

Move the entirety of __new__ into C++, saving a layer of disable_dynamo and making progress toward all-C++.

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

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. [ghstack-poisoned]

pytorch-bot · 2025-09-09T17:44:46Z

🔗 Helpful Links

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

📄 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

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

✅ No Failures

As of commit 2c44479 with merge base a63221a ():
💚 Looks good so far! There are no failures yet. 💚

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

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. ghstack-source-id: 87f5032 Pull Request resolved: #162508

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. ghstack-source-id: e32eed6 Pull Request resolved: #162508

albanD

Not really sure what the end goal here is and the level of optimization we're aiming for. But pybind is most likely costing you a lot, especially for things like setting a slot attribute?

albanD · 2025-09-09T20:50:18Z

torch/csrc/autograd/python_variable.cpp

+    extra_dispatch_keys = extra_dispatch_keys.add(c10::DispatchKey::Negative);
+  }
+
+  py::handle spec = py::handle(r.pyobject(2));


Does .pyobject() return a new or borrowed reference?

borrowed -- it just vends from the input args array, which is borrowed.

albanD · 2025-09-09T20:51:25Z

torch/csrc/autograd/python_variable.cpp

+  }
+
+  py::handle spec = py::handle(r.pyobject(2));
+  const auto tensor_meta = spec.attr(dtensor_interned_strings.tensor_meta);


Some error checking from the old code is lost here when this is None?
Also what happens here when the attribute doesn't exist?

what happens here when the attribute doesn't exist?

PyObject_GetAttr() will set an error, and pybind11 will throw.

Some error checking from the old code is lost here when this is None

you can't get an attribute from None, so this change is safe. it's harder to debug, so I suppose we can TORCH_CHECK here.

albanD · 2025-09-09T20:52:30Z

torch/csrc/autograd/python_variable.cpp

+  py_tensor.attr(dtensor_interned_strings._spec) = spec;
+  py_tensor.attr(dtensor_interned_strings._local_tensor) = local_tensor;


I've never seen this before. Is the refcounting acutally correct on this?!

why wouldn't it be? it would be an obvious bug in pybind11 if this didn't work

swolchok · 2025-09-11T19:20:08Z

Not really sure what the end goal here is and the level of optimization we're aiming for.

Per discussions with @ezyang, the end goal is C++ DTensor. I've chosen to move toward that incrementally, since 1) it would be better if we don't actually have to stop and do a full rewrite 2) incremental software development is faster in general 3) it is unlikely that the incremental pieces will be entirely different from what we would write if we were sitting down with a blank slate.

But pybind is most likely costing you a lot, especially for things like setting a slot attribute?

pybind11's C++ wrappers for the CPython API seem mostly OK . The cost seems to center around per-bound-C++-function-call overhead. I have a pair of upstream pybind11 PRs that reduce that overhead somewhat waiting for reviews (pybind/pybind11#5824 and pybind/pybind11#5830).

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. ghstack-source-id: e1b6902 Pull Request resolved: #162508

…ly native DTensor.__new__" Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

swolchok · 2025-09-21T15:55:48Z

@pytorchbot merge

pytorchmergebot · 2025-09-21T15:58:11Z

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

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. Pull Request resolved: pytorch#162508 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#161695

…nsor_info (#162968) Next PR writes a C++ implementation. Seems good to have tests first. Pull Request resolved: #162968 Approved by: https://github.com/ezyang ghstack dependencies: #161695, #162508

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. Pull Request resolved: pytorch#162508 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#161695

…nsor_info (pytorch#162968) Next PR writes a C++ implementation. Seems good to have tests first. Pull Request resolved: pytorch#162968 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#161695, pytorch#162508

Fully native DTensor.__new__

af9b4f0

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. [ghstack-poisoned]

swolchok requested review from albanD and soulitzer as code owners September 9, 2025 17:44

This was referenced Sep 9, 2025

Add DISABLE_JUSTKNOBS to torch/_utils_internal.py and use it for dynamo _maybe_set_eval_frame #162298

Closed

Fix TODO in make_tensor_for_subclass_helper #162336

Closed

pytorch-bot bot added ciflow/inductor oncall: distributed Add this issue/PR to distributed oncall triage queue labels Sep 9, 2025

This was referenced Sep 8, 2025

Remove __torch_dispatch__ check in THPVariable_make_dtensor #162337

Closed

Port OpSchema.__post_init__ and OpSchema._recompute_comparison_key to C++ #161695

Closed

swolchok added a commit that referenced this pull request Sep 9, 2025

Fully native DTensor.__new__

7e1ff14

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. ghstack-source-id: 87f5032 Pull Request resolved: #162508

swolchok added the topic: not user facing topic category label Sep 9, 2025

Update on "Fully native DTensor.__new__"

ab0c152

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

swolchok added a commit that referenced this pull request Sep 9, 2025

Fully native DTensor.__new__

0259998

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. ghstack-source-id: e32eed6 Pull Request resolved: #162508

albanD reviewed Sep 9, 2025

View reviewed changes

rebase on "Fully native DTensor.__new__"

46fa23c

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

swolchok mentioned this pull request Sep 11, 2025

Add basic tests for torch.distributed.tensor._utils.compute_global_tensor_info #162773

Closed

Update on "Fully native DTensor.__new__"

5be3d2d

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

swolchok mentioned this pull request Sep 15, 2025

Add basic tests for torch.distributed.tensor._utils.compute_global_tensor_info #162968

Closed

Update on "Fully native DTensor.__new__"

3299cfa

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

swolchok mentioned this pull request Sep 15, 2025

DTensor: C++ compute_global_tensor_info #162990

Open

Update on "Fully native DTensor.__new__"

4db671f

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

This was referenced Sep 16, 2025

C++-accessible Placements via pybind11 #163030

Open

Use C++-accessible Placement in compute_global_tensor_info #163031

Closed

Update on "Fully native DTensor.__new__"

e712f48

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

swolchok added the release notes: distributed (dtensor) release notes category label Sep 17, 2025

swolchok removed the topic: not user facing topic category label Sep 17, 2025

swolchok added 2 commits September 17, 2025 11:41

Update on "Fully native DTensor.__new__"

9137a7f

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

Update on "Fully native DTensor.__new__"

081d867

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

swolchok added a commit that referenced this pull request Sep 17, 2025

Fully native DTensor.__new__

b441c98

Move the entirety of `__new__` into C++, saving a layer of disable_dynamo and making progress toward all-C++. ghstack-source-id: e1b6902 Pull Request resolved: #162508

swolchok requested review from albanD and ezyang September 18, 2025 20:10

ezyang approved these changes Sep 21, 2025

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Sep 21, 2025

pytorchmergebot added the merging label Sep 21, 2025

pytorchmergebot closed this in 5599f48 Sep 21, 2025

pytorchmergebot added Merged and removed merging labels Sep 21, 2025

swolchok mentioned this pull request Sep 23, 2025

Remove torch.distributed.tensor.OpSchema.has_symints #163667

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fully native DTensor.new #162508

Fully native DTensor.new #162508

Uh oh!

swolchok commented Sep 9, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Sep 9, 2025 •

edited

Loading

Uh oh!

albanD left a comment

Uh oh!

albanD Sep 9, 2025

Uh oh!

swolchok Sep 11, 2025

Uh oh!

albanD Sep 9, 2025

Uh oh!

swolchok Sep 11, 2025

Uh oh!

albanD Sep 9, 2025

Uh oh!

swolchok Sep 11, 2025

Uh oh!

swolchok commented Sep 11, 2025 •

edited

Loading

Uh oh!

swolchok commented Sep 21, 2025

Uh oh!

pytorchmergebot commented Sep 21, 2025

Uh oh!

Uh oh!

		py_tensor.attr(dtensor_interned_strings._spec) = spec;
		py_tensor.attr(dtensor_interned_strings._local_tensor) = local_tensor;

Fully native DTensor.__new__ #162508

Fully native DTensor.__new__ #162508

Uh oh!

Conversation

swolchok commented Sep 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

albanD Sep 9, 2025

Choose a reason for hiding this comment

Uh oh!

swolchok Sep 11, 2025

Choose a reason for hiding this comment

Uh oh!

albanD Sep 9, 2025

Choose a reason for hiding this comment

Uh oh!

swolchok Sep 11, 2025

Choose a reason for hiding this comment

Uh oh!

albanD Sep 9, 2025

Choose a reason for hiding this comment

Uh oh!

swolchok Sep 11, 2025

Choose a reason for hiding this comment

Uh oh!

swolchok commented Sep 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

swolchok commented Sep 21, 2025

Uh oh!

pytorchmergebot commented Sep 21, 2025

Merge started

Uh oh!

Uh oh!

Fully native DTensor.new #162508

Fully native DTensor.new #162508

swolchok commented Sep 9, 2025 •

edited

Loading

pytorch-bot bot commented Sep 9, 2025 •

edited

Loading

swolchok commented Sep 11, 2025 •

edited

Loading