[DeviceMesh] Use _flatten_rank_map to replace _flatten_mesh_list so that we don't need to compare root mesh #166003

fduwjj · 2025-10-21T16:46:05Z

Stack from ghstack (oldest at bottom):

-> [DeviceMesh] Use _flatten_rank_map to replace _flatten_mesh_list so that we don't need to compare root mesh #166003

Since we are already share a flattened tensor _rank_map across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code.

We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate.

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

Differential Revision: D85394822

…hat we don't need to compare root mesh [ghstack-poisoned]

pytorch-bot · 2025-10-21T16:46:09Z

🔗 Helpful Links

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

📄 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 666b105 with merge base 516e589 ():
💚 Looks good so far! There are no failures yet. 💚

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

…hat we don't need to compare root mesh ghstack-source-id: b20c683 Pull Request resolved: #166003

test/distributed/test_device_mesh.py

fduwjj · 2025-10-23T18:03:54Z

@pytorchbot merge

fduwjj · 2025-10-23T18:59:38Z

@pytorchbot merge -f "bot not working and all CI green"

fduwjj · 2025-10-23T20:42:20Z

@pytorchbot merge

pytorchmergebot · 2025-10-23T20:44:21Z

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

…SPMD use case (#163358) Today FSDP needs to slicing out spmd mesh from root mesh here: https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_fully_shard/_fsdp_param.py#L301. But essentially, users want is a concatenate of some submesh into a big mesh and used as a spmd mesh. This PR is tentatively trying to implement this API for users. One thing to note is that, all sub-mesh needs to slicing/flatten or unflatten from same root mesh otherwise the indices make no sense when it comes to mesh indexing and device allocation. Pull Request resolved: #163358 Approved by: https://github.com/fegin ghstack dependencies: #166003

fduwjj · 2025-10-24T00:10:47Z

@fduwjj has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

clee2000 · 2025-10-24T15:12:43Z

@pytorchbot revert -m "failing internal tests D85405179 I believe there are uses of _flatten_mesh_list internally that need to be updated" -c ghfirst

pytorchmergebot · 2025-10-24T15:14:20Z

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

…ist so that we don't need to compare root mesh (#166003)" This reverts commit 8625ffb. Reverted #166003 on behalf of https://github.com/clee2000 due to failing internal tests D85405179 I believe there are uses of _flatten_mesh_list internally that need to be updated ([comment](#166003 (comment)))

pytorchmergebot · 2025-10-24T15:14:29Z

@fduwjj your PR has been successfully reverted.

…esh and SPMD use case (#163358)" This reverts commit 5a4997d. Reverted #163358 on behalf of https://github.com/clee2000 due to probably need to revert this one too, its stacked with #166003 (comment) ([comment](#163358 (comment)))

fduwjj · 2025-10-24T17:26:21Z

@fduwjj has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

…h_list so that we don't need to compare root mesh" Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta msaroufim dcci Differential Revision: [D85394822](https://our.internmc.facebook.com/intern/diff/D85394822) [ghstack-poisoned]

…hat we don't need to compare root mesh (#166003) Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. Pull Request resolved: #166003 Approved by: https://github.com/Skylion007, https://github.com/fegin Internal: << DO NOT EDIT BELOW THIS LINE >> **GitHub Author**: fduwjj <fduwjj@gmail.com> (Meta Employee) **GitHub Repo**: [pytorch/pytorch](https://github.com/pytorch/pytorch) **GitHub Pull Request**: [#166003](#166003) Initially generated by: https://www.internalfb.com/intern/sandcastle/job/9007201528851998/ This was imported as part of a Diff Train. Please review this as soon as possible. Since it is a direct copy of a commit on GitHub, there shouldn't be much to do. Below line forces Sandcastle to run only specified contbuilds. @build_only[github-export-checks,executorch,pytorch_benchmark,pytorch_benchmark_fb,pytorch_quantization,pytorch_distributed,pytorch_distributed_gpu,pytorch_dynamo,pytorch_inductor,pytorch_inductor_fb,pytorch_functorch,pytorch_fx2trt,pytorch_diff_train_tests_ads,glow_fb_pytorch_tests,training_platform,training_platform_compatibility,training_toolkit_applications,training_toolkit_examples,training_toolkit_model_optimization,dper3_pytorch,xplat_caffe2,pytorch_dev,android-pytorch-instrumentation-tests,smart__pytorch__github_first_try_merge,frl-target-determinator,f6-buck,training_platform_for_github,sigmoid_cpu,sigmoid_gpu,aiplatform_modelprocessing_for_github,accelerators_workloads_models_slimdsnn,ae_aotinductor_benchmark_test,aps_,apf,aps_deterministic_ne_tests,dper_lib_silvertorch,torchrec,torchrec_fb,deeplearning_aot_inductor,aiplatform_modelstore] #skipfbcodelongtail #disable_code_coverage @pytorch-oss-diff-train diff-train-source-id: 8625ffb Differential Revision: [D85394822](https://our.internmc.facebook.com/intern/diff/D85394822/) ghstack-source-id: 318594805

…h_list so that we don't need to compare root mesh" Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta msaroufim dcci Differential Revision: [D85394822](https://our.internmc.facebook.com/intern/diff/D85394822) [ghstack-poisoned]

…hat we don't need to compare root mesh (#166003) Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. Pull Request resolved: #166003 Approved by: https://github.com/Skylion007, https://github.com/fegin Internal: << DO NOT EDIT BELOW THIS LINE >> **GitHub Author**: fduwjj <fduwjj@gmail.com> (Meta Employee) **GitHub Repo**: [pytorch/pytorch](https://github.com/pytorch/pytorch) **GitHub Pull Request**: [#166003](#166003) Initially generated by: https://www.internalfb.com/intern/sandcastle/job/9007201528851998/ This was imported as part of a Diff Train. Please review this as soon as possible. Since it is a direct copy of a commit on GitHub, there shouldn't be much to do. Below line forces Sandcastle to run only specified contbuilds. @build_only[github-export-checks,executorch,pytorch_benchmark,pytorch_benchmark_fb,pytorch_quantization,pytorch_distributed,pytorch_distributed_gpu,pytorch_dynamo,pytorch_inductor,pytorch_inductor_fb,pytorch_functorch,pytorch_fx2trt,pytorch_diff_train_tests_ads,glow_fb_pytorch_tests,training_platform,training_platform_compatibility,training_toolkit_applications,training_toolkit_examples,training_toolkit_model_optimization,dper3_pytorch,xplat_caffe2,pytorch_dev,android-pytorch-instrumentation-tests,smart__pytorch__github_first_try_merge,frl-target-determinator,f6-buck,training_platform_for_github,sigmoid_cpu,sigmoid_gpu,aiplatform_modelprocessing_for_github,accelerators_workloads_models_slimdsnn,ae_aotinductor_benchmark_test,aps_,apf,aps_deterministic_ne_tests,dper_lib_silvertorch,torchrec,torchrec_fb,deeplearning_aot_inductor,aiplatform_modelstore] #skipfbcodelongtail #disable_code_coverage @pytorch-oss-diff-train diff-train-source-id: 8625ffb ghstack-source-id: 318681631 Differential Revision: [D85394822](https://our.internmc.facebook.com/intern/diff/D85394822/)

…h_list so that we don't need to compare root mesh" Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta msaroufim dcci Differential Revision: [D85394822](https://our.internmc.facebook.com/intern/diff/D85394822) [ghstack-poisoned]

…hat we don't need to compare root mesh (#166003) Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. Pull Request resolved: #166003 Approved by: https://github.com/Skylion007, https://github.com/fegin Internal: << DO NOT EDIT BELOW THIS LINE >> **GitHub Author**: fduwjj <fduwjj@gmail.com> (Meta Employee) **GitHub Repo**: [pytorch/pytorch](https://github.com/pytorch/pytorch) **GitHub Pull Request**: [#166003](#166003) Initially generated by: https://www.internalfb.com/intern/sandcastle/job/9007201528851998/ This was imported as part of a Diff Train. Please review this as soon as possible. Since it is a direct copy of a commit on GitHub, there shouldn't be much to do. Below line forces Sandcastle to run only specified contbuilds. @build_only[github-export-checks,executorch,pytorch_benchmark,pytorch_benchmark_fb,pytorch_quantization,pytorch_distributed,pytorch_distributed_gpu,pytorch_dynamo,pytorch_inductor,pytorch_inductor_fb,pytorch_functorch,pytorch_fx2trt,pytorch_diff_train_tests_ads,glow_fb_pytorch_tests,training_platform,training_platform_compatibility,training_toolkit_applications,training_toolkit_examples,training_toolkit_model_optimization,dper3_pytorch,xplat_caffe2,pytorch_dev,android-pytorch-instrumentation-tests,smart__pytorch__github_first_try_merge,frl-target-determinator,f6-buck,training_platform_for_github,sigmoid_cpu,sigmoid_gpu,aiplatform_modelprocessing_for_github,accelerators_workloads_models_slimdsnn,ae_aotinductor_benchmark_test,aps_,apf,aps_deterministic_ne_tests,dper_lib_silvertorch,torchrec,torchrec_fb,deeplearning_aot_inductor,aiplatform_modelstore] #skipfbcodelongtail #disable_code_coverage @pytorch-oss-diff-train diff-train-source-id: 8625ffb ghstack-source-id: 79a788c Differential Revision: [D85394822](https://our.internmc.facebook.com/intern/diff/D85394822/)

fduwjj · 2025-10-25T20:42:13Z

@fduwjj has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

fduwjj · 2025-10-26T03:37:52Z

@pytorchbot merge

pytorchmergebot · 2025-10-26T03:39:44Z

Merge failed

Reason: This PR has internal changes and must be landed via Phabricator! Please try reimporting/rexporting the PR!

Details for Dev Infra team

Raised by workflow job

…h_list so that we don't need to compare root mesh" Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta msaroufim dcci Differential Revision: [D85394822](https://our.internmc.facebook.com/intern/diff/D85394822) [ghstack-poisoned]

…hat we don't need to compare root mesh (#166003) Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. Pull Request resolved: #166003 Approved by: https://github.com/Skylion007, https://github.com/fegin Internal: << DO NOT EDIT BELOW THIS LINE >> **GitHub Author**: fduwjj <fduwjj@gmail.com> (Meta Employee) **GitHub Repo**: [pytorch/pytorch](https://github.com/pytorch/pytorch) **GitHub Pull Request**: [#166003](#166003) Initially generated by: https://www.internalfb.com/intern/sandcastle/job/9007201528851998/ This was imported as part of a Diff Train. Please review this as soon as possible. Since it is a direct copy of a commit on GitHub, there shouldn't be much to do. Below line forces Sandcastle to run only specified contbuilds. @build_only[github-export-checks,executorch,pytorch_benchmark,pytorch_benchmark_fb,pytorch_quantization,pytorch_distributed,pytorch_distributed_gpu,pytorch_dynamo,pytorch_inductor,pytorch_inductor_fb,pytorch_functorch,pytorch_fx2trt,pytorch_diff_train_tests_ads,glow_fb_pytorch_tests,training_platform,training_platform_compatibility,training_toolkit_applications,training_toolkit_examples,training_toolkit_model_optimization,dper3_pytorch,xplat_caffe2,pytorch_dev,android-pytorch-instrumentation-tests,smart__pytorch__github_first_try_merge,frl-target-determinator,f6-buck,training_platform_for_github,sigmoid_cpu,sigmoid_gpu,aiplatform_modelprocessing_for_github,accelerators_workloads_models_slimdsnn,ae_aotinductor_benchmark_test,aps_,apf,aps_deterministic_ne_tests,dper_lib_silvertorch,torchrec,torchrec_fb,deeplearning_aot_inductor,aiplatform_modelstore] #skipfbcodelongtail #disable_code_coverage @pytorch-oss-diff-train diff-train-source-id: 8625ffb ghstack-source-id: 318735710 Differential Revision: [D85394822](https://our.internmc.facebook.com/intern/diff/D85394822/)

…hat we don't need to compare root mesh (#166003) Summary: Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta msaroufim dcci imported-using-ghimport Test Plan: Imported from OSS Differential Revision: D85526705 Pulled By: fduwjj

…hat we don't need to compare root mesh (#166003) (#166264) Summary: Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta msaroufim dcci imported-using-ghimport Test Plan: Imported from OSS Differential Revision: D85526705 Pulled By: fduwjj Pull Request resolved: #166264 Approved by: https://github.com/XilunWu

fduwjj · 2025-10-27T03:45:40Z

merged in #166264

…hat we don't need to compare root mesh (#166003) (#166264) Summary: Since we are already share a flattened tensor `_rank_map` across all meshes from a same root mesh, we can just use a flattened list of it to replace the comparison of root_mesh and flattened_mesh_list (because with same _rank_map and layout, the mesh tensor is guaranteed to be the same). This way we can also give back the CPU overhead added in #164510 and further simply the code. We do have a more ambitious universe-based change here: #165680 but it needs more discussions and would lead to BC breaking. We might eventually merge that PR but probably not now and this is a change which is not BC breaking and will help concatenate and 2D integration with concatenate. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta msaroufim dcci imported-using-ghimport Test Plan: Imported from OSS Differential Revision: D85526705 Pulled By: fduwjj Pull Request resolved: #166264 Approved by: https://github.com/XilunWu

[DeviceMesh] Use _flatten_rank_map to replace _flatten_mesh_list so t…

3c09f69

…hat we don't need to compare root mesh [ghstack-poisoned]

pytorch-bot bot added the oncall: distributed Add this issue/PR to distributed oncall triage queue label Oct 21, 2025

fduwjj added a commit that referenced this pull request Oct 21, 2025

[DeviceMesh] Use _flatten_rank_map to replace _flatten_mesh_list so t…

a251742

…hat we don't need to compare root mesh ghstack-source-id: b20c683 Pull Request resolved: #166003

fduwjj requested review from fegin, lw and wconstab October 21, 2025 16:52

fduwjj added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 21, 2025

fduwjj mentioned this pull request Oct 17, 2025

[DeviceMesh] Implement a device mesh concatenate api for submesh and SPMD use case #163358

Closed

fduwjj added the release notes: DeviceMesh label Oct 21, 2025

fduwjj commented Oct 21, 2025

View reviewed changes

test/distributed/test_device_mesh.py Show resolved Hide resolved

This was referenced Oct 17, 2025

[DeviceMesh][2D] Use concatenate for 2D (FSDP+TP) instead of getting from root mesh #165492

Closed

[For discussion][DeviceMesh] Use a shared_state to cache pg per layout, root_mesh and rank_map #166010

Open

Skylion007 approved these changes Oct 23, 2025

View reviewed changes

pytorchmergebot added the merging label Oct 23, 2025

fegin approved these changes Oct 23, 2025

View reviewed changes

pytorchmergebot closed this in 8625ffb Oct 23, 2025

pytorchmergebot added Merged and removed merging labels Oct 23, 2025

pytorchmergebot added the ci-no-td Do not run TD on this PR label Oct 24, 2025

pytorchmergebot reopened this Oct 24, 2025

pytorch-auto-revert bot mentioned this pull request Oct 24, 2025

[DO NOT CLOSE] Autorevert actions shadow mode stream #163650

Open

meta-codesync bot added fb-exported meta-exported labels Oct 25, 2025

pytorchmergebot added the merging label Oct 26, 2025

pytorchmergebot removed the merging label Oct 26, 2025

fduwjj closed this Oct 27, 2025

github-actions bot deleted the gh/fduwjj/228/head branch November 28, 2025 02:17

[DeviceMesh] Use _flatten_rank_map to replace _flatten_mesh_list so that we don't need to compare root mesh #166003

[DeviceMesh] Use _flatten_rank_map to replace _flatten_mesh_list so that we don't need to compare root mesh #166003

Uh oh!

Conversation

fduwjj commented Oct 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Oct 21, 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/166003

✅ No Failures

Uh oh!

Uh oh!

fduwjj commented Oct 23, 2025

Uh oh!

fduwjj commented Oct 23, 2025

Uh oh!

fduwjj commented Oct 23, 2025

Uh oh!

pytorchmergebot commented Oct 23, 2025

Merge started

Uh oh!

fduwjj commented Oct 24, 2025

Uh oh!

clee2000 commented Oct 24, 2025

Uh oh!

pytorchmergebot commented Oct 24, 2025

Uh oh!

pytorchmergebot commented Oct 24, 2025

Uh oh!

fduwjj commented Oct 24, 2025

Uh oh!

fduwjj commented Oct 25, 2025

Uh oh!

fduwjj commented Oct 26, 2025

Uh oh!

pytorchmergebot commented Oct 26, 2025

Merge failed

Uh oh!

fduwjj commented Oct 27, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

fduwjj commented Oct 21, 2025 •

edited

Loading

pytorch-bot bot commented Oct 21, 2025 •

edited

Loading