[DTensor] Disallow redistribution to mixed partial types by wconstab · Pull Request #172609 · pytorch/pytorch

wconstab · 2026-01-16T00:49:29Z

Stack from ghstack (oldest at bottom):

[ghstack-poisoned]

pytorch-bot · 2026-01-16T00:49:32Z

🔗 Helpful Links

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

📄 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.

✅ You can merge normally! (3 Unrelated Failures)

As of commit c4a0c7c with merge base 4ac0624 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

inductor / inductor-test / test (inductor_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
shufflenet_v2_x1_0
inductor / unit-test / inductor-test / test (inductor, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (disabled by #137684)
test/inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCUDA::test_comprehensive_linalg_lu_cuda_float32
trunk / linux-jammy-rocm-py3.10 / test (distributed, 1, 3, linux.rocm.gpu.gfx942.4) (gh) (similar failure)
test/distributed/_composable/fsdp/test_fully_shard_training.py::TestFullyShardCudaGraph::test_two_layer_fully_shard_cudagraph

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

zpcore · 2026-01-16T17:42:51Z

+    partial_reduce_ops: set[str] = set()
+    for p in placements:
+        if p.is_partial():
+            partial_reduce_ops.add(p.reduce_op)  # type: ignore[union-attr]


Not sure if this make code simpler. We can compare Partial object directly.
partial_ops = set([p for p in placement if p.is_partial()])

zpcore

LGTM!

Curious, do we still plan to support mixed Partial in the future?

wconstab · 2026-01-17T22:42:04Z

No, can't imagine a use case, so unless someone finds a compelling use case let's simplify our lives.

ghstack-source-id: b13b323 Pull Request resolved: pytorch/pytorch#172609

tianyu-l · 2026-01-19T20:19:25Z

+    # We do not see a valid use case for mixing different partial types in the same DTensor.
+    # in principle it could be supported, but since nonlinear reductions (e.g. max) exist, relative ordering
+    # of different partials would become semantically critical.  Without a motivating use case, we prohibit this.
+    _assert_no_mixed_partial_types(current_spec.placements)
+    _assert_no_mixed_partial_types(target_spec.placements)


Is it true that there are several places where mixed partial could show up?

distribute_tensor placements arg (not banned)

DTensor.from_local placements arg (not banned)

sharding propagation (banned here if needs_redistribute; o/w not banned although I can't think of a real case where it happens from any op)

DTensor.redistribute (banned here)

From a completeness perspective, do you think it makes sense to make _assert_no_mixed_partial_types into utils file, and call it from case 1 and case 2?

yea, I was planning to add (1) and (2) in a later PR and update docs accordingly, but if folks are in agreement about this I can just include it in this PR. And I agree about (3), it is not strictly banned but it seems practically impossible to happen.

[ghstack-poisoned]

wconstab · 2026-01-20T19:25:22Z

@pytorchbot merge

pytorchmergebot · 2026-01-20T19:28:00Z

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

pytorchmergebot · 2026-01-20T20:31:48Z

Merge failed

Reason: 2 jobs have failed, first few of them are: trunk / linux-jammy-rocm-py3.10 / test (distributed, 1, 3, linux.rocm.gpu.gfx942.4), trunk / linux-jammy-rocm-py3.10 / test (distributed, 2, 3, linux.rocm.gpu.gfx942.4)

Details for Dev Infra team

Raised by workflow job

[ghstack-poisoned]

wconstab · 2026-01-21T00:14:24Z

@pytorchbot merge -i

pytorchmergebot · 2026-01-21T00:16:28Z

Merge started

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

pytorchmergebot · 2026-01-21T00:22:32Z

Merge failed

Reason: 1 jobs have failed, first few of them are: Check Labels / Check labels

Details for Dev Infra team

Raised by workflow job

wconstab · 2026-01-21T04:06:55Z

@pytorchbot merge -i

pytorchmergebot · 2026-01-21T04:09:00Z

Merge started

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

pytorchmergebot · 2026-01-26T04:07:40Z

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

…orch#172609)" This reverts commit eba22f9. Reverted pytorch#172609 on behalf of https://github.com/wconstab due to Breaking internal use case https://www.internalfb.com/intern/testinfra/testconsole/testrun/10696049259246146 ([comment](pytorch#172609 (comment)))

) Pull Request resolved: pytorch#172609 Approved by: https://github.com/zpcore, https://github.com/tianyu-l ghstack dependencies: pytorch#172924

Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. ghstack-source-id: eb4a927 Pull Request resolved: #173614

…filters mixed partials" Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

…rtials" Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. ghstack-source-id: 8cca0a6 Pull Request resolved: #173614

Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. ghstack-source-id: 2984751 Pull Request resolved: #173614

…filters mixed partials" Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

…rtials" Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

…filters mixed partials" Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

…rtials" Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. ghstack-source-id: 2483827 Pull Request resolved: #173614

Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. ghstack-source-id: e4407b4 Pull Request resolved: #173614

…filters mixed partials" Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

…rtials" Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. ghstack-source-id: 38e8412 Pull Request resolved: #173614

…filters mixed partials" Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

…rtials" Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. [ghstack-poisoned]

…3614) Per #172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. Pull Request resolved: #173614 Approved by: https://github.com/pianpwk, https://github.com/zpcore

…orch#173614) Per pytorch#172609, we do not allow mixed partial placements in one DTensor spec. It is therefore pointless to generate strategies that encourage mixed partial types. Before this PR, any OP strategy that has a rule for more than one partial type and uses the expand util would have combined them inappropriately. After this PR, the expansion util filters out any such combinations. Pull Request resolved: pytorch#173614 Approved by: https://github.com/pianpwk, https://github.com/zpcore

[DTensor] Disallow redistribution to mixed partial types

0162ee2

[ghstack-poisoned]

pytorch-bot Bot added ciflow/inductor release notes: distributed (dtensor) release notes category labels Jan 16, 2026

wconstab mentioned this pull request Jan 16, 2026

[DTensor] Optimize redistribute comms using flattened meshes #172610

Closed

zpcore reviewed Jan 16, 2026

View reviewed changes

zpcore approved these changes Jan 16, 2026

View reviewed changes

SergeyTyshkevich pushed a commit to SergeyTyshkevich/chart2 that referenced this pull request Jan 19, 2026

[DTensor] Disallow redistribution to mixed partial types

ed1c627

ghstack-source-id: b13b323 Pull Request resolved: pytorch/pytorch#172609

tianyu-l reviewed Jan 19, 2026

View reviewed changes

Update on "[DTensor] Disallow redistribution to mixed partial types"

cf1690b

[ghstack-poisoned]

tianyu-l approved these changes Jan 20, 2026

View reviewed changes

Update on "[DTensor] Disallow redistribution to mixed partial types"

f717fab

[ghstack-poisoned]

wconstab mentioned this pull request Jan 20, 2026

[DTensor] Fix partial redistribution order #172277

Closed

pytorch-bot Bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jan 20, 2026

pytorchmergebot added the merging label Jan 20, 2026

pytorchmergebot removed the merging label Jan 20, 2026

Update on "[DTensor] Disallow redistribution to mixed partial types"

22432bc

[ghstack-poisoned]

pytorchmergebot added the merging label Jan 21, 2026

pytorchmergebot removed the merging label Jan 21, 2026

pytorchmergebot added the merging label Jan 21, 2026

pytorchmergebot added the merging label Jan 26, 2026

pytorchmergebot closed this in a560e01 Jan 26, 2026

pytorchmergebot removed the merging label Jan 26, 2026

This was referenced Jan 26, 2026

[DTensor] Make RedistributionPlanner handle all partials #172479

Closed

[DTensor] make debugmode print optimized transforminfos #173436

Closed

wconstab mentioned this pull request Jan 28, 2026

[DTensor] expand_to_full_mesh_op_strategy filters mixed partials #173614

Closed

Conversation

wconstab commented Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Jan 16, 2026 • 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/172609

✅ You can merge normally! (3 Unrelated Failures)

Uh oh!

zpcore Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

zpcore left a comment

Choose a reason for hiding this comment

Uh oh!

wconstab commented Jan 17, 2026

Uh oh!

tianyu-l Jan 19, 2026

Choose a reason for hiding this comment

Uh oh!

wconstab Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

wconstab commented Jan 20, 2026

Uh oh!

pytorchmergebot commented Jan 20, 2026

Merge started

Uh oh!

pytorchmergebot commented Jan 20, 2026

Merge failed

Uh oh!

wconstab commented Jan 21, 2026

Uh oh!

pytorchmergebot commented Jan 21, 2026

Merge started

Uh oh!

pytorchmergebot commented Jan 21, 2026

Merge failed

Uh oh!

wconstab commented Jan 21, 2026

Uh oh!

pytorchmergebot commented Jan 21, 2026

Merge started

Uh oh!

pytorchmergebot commented Jan 26, 2026

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

wconstab commented Jan 16, 2026 •

edited

Loading

pytorch-bot Bot commented Jan 16, 2026 •

edited

Loading