Skip to content

[DeviceMesh] Move DeviceMesh out from torch.distributed._tensor #112364

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 2 commits into from
Closed

[DeviceMesh] Move DeviceMesh out from torch.distributed._tensor #112364

wants to merge 2 commits into from

Conversation

wz337
Copy link
Contributor

@wz337 wz337 commented Oct 30, 2023 •
edited
Loading

Move DeviceMesh out as a standalone module. Once we make sure everything is migrated and doc is ready, we will make torch.distributed._device_mesh public in follow-up PRs.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng

@pytorch-bot
Copy link

pytorch-bot bot commented Oct 30, 2023 •
edited
Loading

🔗 Helpful Links

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

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit df6792f with merge base b91fcdf (image):

FLAKY - The following job failed but was likely due to flakiness present on trunk:

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

@wz337 wz337 force-pushed the move_device_mesh_out branch from 9f7556d to b02e839 Compare October 30, 2023 04:22
@wz337 wz337 added module: DeviceMesh topic: not user facing topic category ciflow/trunk Trigger trunk jobs on your pull request ciflow/periodic Trigger jobs ran periodically on master (periodic.yml) on the PR labels Oct 30, 2023
wanchaol
wanchaol requested changes Oct 30, 2023
View reviewed changes
Copy link
Collaborator

@wanchaol wanchaol left a comment

Choose a reason for hiding this comment

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

This looks wrong, we shouldn't move device mesh like this. I think it should be:
torch.distributed._tensor.device_mesh -> torch.distributed._device_mesh
test/distributed/_tensor/test_device_mesh -> test/distributed/test_device_mesh

No need to put deprecation warnings in _tensor.device_mesh, that file can just import the device mesh from torch.distributed, then we can fix the imports as a separate PR.

@wz337 wz337 force-pushed the move_device_mesh_out branch 13 times, most recently from 525d318 to bf61fe6 Compare October 31, 2023 22:52
@wz337 wz337 force-pushed the move_device_mesh_out branch from bf61fe6 to 2e7963c Compare November 1, 2023 00:13
@wz337 wz337 changed the title [WIP] move device mesh out [DeviceMesh] Move DeviceMesh out from torch.distributed._tensor Nov 1, 2023
@wz337 wz337 marked this pull request as ready for review November 1, 2023 16:04
@wz337 wz337 requested a review from wanchaol November 1, 2023 16:04
fegin
fegin reviewed Nov 1, 2023
View reviewed changes
torch/_dynamo/skipfiles.py Outdated
Copy link
Contributor

Choose a reason for hiding this comment

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

We can keep this line and can avoid confusion when compiling the model that still depend on this path.

fegin
fegin approved these changes Nov 1, 2023
View reviewed changes
Copy link
Contributor

@fegin fegin left a comment

Choose a reason for hiding this comment

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

LGTM

@wz337 wz337 force-pushed the move_device_mesh_out branch from fd55567 to 0c45d4d Compare November 1, 2023 19:46
wanchaol
wanchaol approved these changes Nov 1, 2023
View reviewed changes
Copy link
Collaborator

@wanchaol wanchaol left a comment

Choose a reason for hiding this comment

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

lgtm, please make sure CI passes :)

fduwjj
fduwjj reviewed Nov 1, 2023
View reviewed changes
torch/distributed/_tensor/device_mesh.py Outdated
Copy link
Contributor

Choose a reason for hiding this comment

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

shall we add a warning here telling users that the path has been moved?

fduwjj
fduwjj approved these changes Nov 1, 2023
View reviewed changes
Copy link
Contributor

@fduwjj fduwjj left a comment

Choose a reason for hiding this comment

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

Actually "we will make torch.distributed._device_mesh public." is not accurate because with _device_mesh it is still private? IIUC. But it's great that we decouple it from DTensor.

@wz337 wz337 force-pushed the move_device_mesh_out branch from 0c45d4d to df6792f Compare November 2, 2023 00:59
@wz337
Copy link
Contributor Author

wz337 commented Nov 2, 2023 •
edited
Loading

Actually "we will make torch.distributed._device_mesh public." is not accurate because with _device_mesh it is still private? IIUC. But it's great that we decouple it from DTensor.

Ye. We will in following PRs. Updated the summary as well.

@wz337
Copy link
Contributor Author

wz337 commented Nov 2, 2023

Actually "we will make torch.distributed._device_mesh public." is not accurate because with _device_mesh it is still private? IIUC. But it's great that we decouple it from DTensor.

I think it should be fine, as the previous path is still importing from the new path so nothing would be broken.
And since DTensor was prototype as well, we should be able to make any modification we want.

@fduwjj
Copy link
Contributor

fduwjj commented Nov 2, 2023

I think it should be fine, as the previous path is still importing from the new path so nothing would be broken.

I am not talking about BC breaking here. I mean for new users, they import from the new path right?

@wz337
Copy link
Contributor Author

wz337 commented Nov 2, 2023

I think it should be fine, as the previous path is still importing from the new path so nothing would be broken.

I am not talking about BC breaking here. I mean for new users, they import from the new path right?

Once the migration is done, we will add doc to the .rst file and make it public.

@wz337
Copy link
Contributor Author

wz337 commented Nov 2, 2023

@pytorchmergebot merge

@pytorchmergebot
Copy link
Collaborator

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

@wanchaol wanchaol added the release notes: distributed (dtensor) release notes category label Nov 2, 2023
xuhancn pushed a commit to xuhancn/pytorch that referenced this pull request Nov 7, 2023
@wz337 @xuhancn
[DeviceMesh] Move DeviceMesh out from torch.distributed._tensor (pyto…
1c0cb08 
…rch#112364)

Move DeviceMesh out as a standalone module. Once we make sure everything is migrated and doc is ready, we will make `torch.distributed._device_mesh` public in follow-up PRs.

Pull Request resolved: pytorch#112364
Approved by: https://github.com/wanchaol, https://github.com/fegin, https://github.com/fduwjj
This was referenced Nov 11, 2023
Closed
Closed
Skylion007 pushed a commit to Skylion007/pytorch that referenced this pull request Nov 14, 2023
@wz337 @Skylion007
[DeviceMesh] Move DeviceMesh out from torch.distributed._tensor (pyto…
4cfcadb 
…rch#112364)

Move DeviceMesh out as a standalone module. Once we make sure everything is migrated and doc is ready, we will make `torch.distributed._device_mesh` public in follow-up PRs.

Pull Request resolved: pytorch#112364
Approved by: https://github.com/wanchaol, https://github.com/fegin, https://github.com/fduwjj
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@fegin fegin fegin approved these changes

@fduwjj fduwjj fduwjj approved these changes

@wanchaol wanchaol wanchaol approved these changes

@mrshenli mrshenli Awaiting requested review from mrshenli

@zhaojuanmao zhaojuanmao Awaiting requested review from zhaojuanmao

@rohan-varma rohan-varma Awaiting requested review from rohan-varma

@H-Huang H-Huang Awaiting requested review from H-Huang

@awgu awgu Awaiting requested review from awgu

@kwen2501 kwen2501 Awaiting requested review from kwen2501

@kiukchung kiukchung Awaiting requested review from kiukchung

@d4l3k d4l3k Awaiting requested review from d4l3k

@albanD albanD Awaiting requested review from albanD albanD is a code owner

Assignees
No one assigned
Labels
ciflow/inductor ciflow/periodic Trigger jobs ran periodically on master (periodic.yml) on the PR ciflow/trunk Trigger trunk jobs on your pull request Merged module: DeviceMesh module: dynamo release notes: distributed (dtensor) release notes category topic: not user facing topic category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
5 participants
@wz337 @fduwjj @pytorchmergebot @fegin @wanchaol