[coor-slicing] extract _get_mesh_tensor_from_full_mesh by aorenste · Pull Request #169550 · pytorch/pytorch

aorenste · 2025-12-04T06:29:01Z

For compile on one rank we need to be able to compute the DeviceMesh rank Tensor based on the raw Tensor and current rank. So this PR factors out DeviceMesh._get_mesh_tensor_from_full_mesh() into a static method.

Stack from ghstack (oldest at bottom):

[ghstack-poisoned]

pytorch-bot · 2025-12-04T06:29:05Z

🔗 Helpful Links

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

📄 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, 91 Pending

As of commit 159b331 with merge base dc48fef ():
💚 Looks good so far! There are no failures yet. 💚

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

⏳ pull / linux-docs / build-docs-python-false (gh) (#111256)

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

[ghstack-poisoned]

ghstack-source-id: 2f72461 Pull Request resolved: pytorch/pytorch#169550

[ghstack-poisoned]

For compile on one rank we need to be able to compute the DeviceMesh rank Tensor based on the raw Tensor and current rank. So this PR factors out `DeviceMesh._get_mesh_tensor_from_full_mesh()` into a static method. [ghstack-poisoned]

pytorchmergebot · 2026-01-09T02:28:57Z

Starting merge as part of PR stack under #169551

aorenste · 2026-01-09T03:35:18Z

@pytorchbot merge

pytorchmergebot · 2026-01-09T03:37:13Z

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-09T04:03:54Z

Merge failed

Reason: 2 jobs have failed, first few of them are: trunk / linux-jammy-cuda13.0-py3.10-gcc11 / test (distributed, 1, 3, lf.linux.g4dn.12xlarge.nvidia.gpu), trunk / linux-jammy-cuda12.8-py3.10-gcc11 / test (distributed, 1, 3, lf.linux.g4dn.12xlarge.nvidia.gpu)

Details for Dev Infra team

Raised by workflow job

For compile on one rank we need to be able to compute the DeviceMesh rank Tensor based on the raw Tensor and current rank. So this PR factors out `DeviceMesh._get_mesh_tensor_from_full_mesh()` into a static method. [ghstack-poisoned]

pytorchmergebot · 2026-01-10T17:09:01Z

Starting merge as part of PR stack under #169551

`Placement._split_tensor()` computes and returns too much information - in general most callers call it and then throw away most of the results. Added `Placement._select_split_tensor()` which allows the caller to say which parts they want so we can compute only those bits - in essence it is the combination of `Placement._split_tensor()` and `Shard._select_shard()`. Pull Request resolved: #169551 Approved by: https://github.com/ezyang ghstack dependencies: #169549, #169550

jeanschmidt · 2026-01-12T06:47:07Z

@pytorchbot revert -m "seems to be breaking internal signals, see D90448078" -c ghfirst

pytorchmergebot · 2026-01-12T06:48:40Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

pytorchmergebot · 2026-01-12T06:48:49Z

@aorenste your PR has been successfully reverted.

…)" This reverts commit 649d9b3. Reverted #169550 on behalf of https://github.com/jeanschmidt due to seems to be breaking internal signals, see D90448078 ([comment](#169550 (comment)))

pytorchmergebot · 2026-01-12T07:05:21Z

Starting merge as part of PR stack under #169551

`Placement._split_tensor()` computes and returns too much information - in general most callers call it and then throw away most of the results. Added `Placement._select_split_tensor()` which allows the caller to say which parts they want so we can compute only those bits - in essence it is the combination of `Placement._split_tensor()` and `Shard._select_shard()`. Pull Request resolved: #169551 Approved by: https://github.com/ezyang ghstack dependencies: #169549, #169550

For compile on one rank we need to be able to compute the DeviceMesh rank Tensor based on the raw Tensor and current rank. So this PR factors out `DeviceMesh._get_mesh_tensor_from_full_mesh()` into a static method. Pull Request resolved: pytorch#169550 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#169549

`Placement._split_tensor()` computes and returns too much information - in general most callers call it and then throw away most of the results. Added `Placement._select_split_tensor()` which allows the caller to say which parts they want so we can compute only those bits - in essence it is the combination of `Placement._split_tensor()` and `Shard._select_shard()`. Pull Request resolved: pytorch#169551 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#169549, pytorch#169550

…ch#169550)" This reverts commit 649d9b3. Reverted pytorch#169550 on behalf of https://github.com/jeanschmidt due to seems to be breaking internal signals, see D90448078 ([comment](pytorch#169550 (comment)))

For compile on one rank we need to be able to compute the DeviceMesh rank Tensor based on the raw Tensor and current rank. So this PR factors out `DeviceMesh._get_mesh_tensor_from_full_mesh()` into a static method. Pull Request resolved: pytorch#169550 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#169549

`Placement._split_tensor()` computes and returns too much information - in general most callers call it and then throw away most of the results. Added `Placement._select_split_tensor()` which allows the caller to say which parts they want so we can compute only those bits - in essence it is the combination of `Placement._split_tensor()` and `Shard._select_shard()`. Pull Request resolved: pytorch#169551 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#169549, pytorch#169550

For compile on one rank we need to be able to compute the DeviceMesh rank Tensor based on the raw Tensor and current rank. So this PR factors out `DeviceMesh._get_mesh_tensor_from_full_mesh()` into a static method. Pull Request resolved: pytorch#169550 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#169549

`Placement._split_tensor()` computes and returns too much information - in general most callers call it and then throw away most of the results. Added `Placement._select_split_tensor()` which allows the caller to say which parts they want so we can compute only those bits - in essence it is the combination of `Placement._split_tensor()` and `Shard._select_shard()`. Pull Request resolved: pytorch#169551 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#169549, pytorch#169550

…ch#169550)" This reverts commit 649d9b3. Reverted pytorch#169550 on behalf of https://github.com/jeanschmidt due to seems to be breaking internal signals, see D90448078 ([comment](pytorch#169550 (comment)))

For compile on one rank we need to be able to compute the DeviceMesh rank Tensor based on the raw Tensor and current rank. So this PR factors out `DeviceMesh._get_mesh_tensor_from_full_mesh()` into a static method. Pull Request resolved: pytorch#169550 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#169549

`Placement._split_tensor()` computes and returns too much information - in general most callers call it and then throw away most of the results. Added `Placement._select_split_tensor()` which allows the caller to say which parts they want so we can compute only those bits - in essence it is the combination of `Placement._split_tensor()` and `Shard._select_shard()`. Pull Request resolved: pytorch#169551 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#169549, pytorch#169550

ghstack-source-id: 34bf40c Pull Request resolved: pytorch/pytorch#169550

ghstack-source-id: 96485cf Pull Request resolved: pytorch/pytorch#169550

WIP: extract _get_mesh_tensor_from_full_mesh

098f053

[ghstack-poisoned]

aorenste added 4 commits December 4, 2025 08:38

Update on "WIP: extract _get_mesh_tensor_from_full_mesh"

73f874b

[ghstack-poisoned]

Update on "WIP: extract _get_mesh_tensor_from_full_mesh"

0b14851

[ghstack-poisoned]

Update on "WIP: extract _get_mesh_tensor_from_full_mesh"

f5667df

[ghstack-poisoned]

Update on "WIP: extract _get_mesh_tensor_from_full_mesh"

d913b0f

[ghstack-poisoned]

tiendatngcs pushed a commit to tiendatngcs/pytorch-Dec25 that referenced this pull request Dec 10, 2025

WIP: extract _get_mesh_tensor_from_full_mesh

43351e7

ghstack-source-id: 2f72461 Pull Request resolved: pytorch/pytorch#169550

Update on "WIP: extract _get_mesh_tensor_from_full_mesh"

946fda9

[ghstack-poisoned]

aorenste added the topic: not user facing topic category label Dec 10, 2025

aorenste changed the title ~~WIP: extract _get_mesh_tensor_from_full_mesh~~ [coor-slicing] extract _get_mesh_tensor_from_full_mesh Jan 6, 2026

aorenste marked this pull request as ready for review January 7, 2026 14:35

aorenste requested a review from ezyang January 7, 2026 14:35

ezyang approved these changes Jan 8, 2026

View reviewed changes

ezyang requested a review from fduwjj January 8, 2026 15:13

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

pytorchmergebot added the merging label Jan 9, 2026

pytorchmergebot closed this in 649d9b3 Jan 10, 2026

pytorchmergebot added the Merged label Jan 10, 2026

pytorchmergebot added Reverted ci-no-td Do not run TD on this PR labels Jan 12, 2026

pytorchmergebot reopened this Jan 12, 2026

pytorchmergebot closed this in 876ee1e Jan 12, 2026

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

WIP: extract _get_mesh_tensor_from_full_mesh

c9f7edb

ghstack-source-id: 34bf40c Pull Request resolved: pytorch/pytorch#169550

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

[coor-slicing] extract _get_mesh_tensor_from_full_mesh

4a18232

ghstack-source-id: 96485cf Pull Request resolved: pytorch/pytorch#169550

github-actions Bot deleted the gh/aorenste/157/head branch February 12, 2026 02:23

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[coor-slicing] extract _get_mesh_tensor_from_full_mesh#169550

[coor-slicing] extract _get_mesh_tensor_from_full_mesh#169550
aorenste wants to merge 10 commits intogh/aorenste/157/basefrom
gh/aorenste/157/head

aorenste commented Dec 4, 2025 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Dec 4, 2025 •

edited

Loading

Uh oh!

pytorchmergebot commented Jan 9, 2026

Uh oh!

aorenste commented Jan 9, 2026

Uh oh!

pytorchmergebot commented Jan 9, 2026

Uh oh!

pytorchmergebot commented Jan 9, 2026

Uh oh!

pytorchmergebot commented Jan 10, 2026

Uh oh!

jeanschmidt commented Jan 12, 2026

Uh oh!

pytorchmergebot commented Jan 12, 2026

Uh oh!

pytorchmergebot commented Jan 12, 2026

Uh oh!

pytorchmergebot commented Jan 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

aorenste commented Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Dec 4, 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/169550

⏳ No Failures, 91 Pending

Uh oh!

pytorchmergebot commented Jan 9, 2026

Uh oh!

aorenste commented Jan 9, 2026

Uh oh!

pytorchmergebot commented Jan 9, 2026

Merge started

Uh oh!

pytorchmergebot commented Jan 9, 2026

Merge failed

Uh oh!

pytorchmergebot commented Jan 10, 2026

Uh oh!

jeanschmidt commented Jan 12, 2026

Uh oh!

pytorchmergebot commented Jan 12, 2026

Uh oh!

pytorchmergebot commented Jan 12, 2026

Uh oh!

pytorchmergebot commented Jan 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

aorenste commented Dec 4, 2025 •

edited

Loading

pytorch-bot Bot commented Dec 4, 2025 •

edited

Loading