Skip to content

Conversation

awgu
Copy link
Collaborator

@awgu awgu commented Apr 23, 2024

Copy link

pytorch-bot bot commented Apr 23, 2024

🔗 Helpful Links

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

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

✅ No Failures

As of commit 0f814db with merge base c82fcb7 (image):
💚 Looks good so far! There are no failures yet. 💚

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

@pytorch-bot pytorch-bot bot added oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: distributed (fsdp) release notes category labels Apr 23, 2024
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 sounds good to me! have one suggestion

@@ -434,6 +440,23 @@ def get_group(
)
return dim_groups

@staticmethod
def from_group(group: ProcessGroup, device_type: str) -> "DeviceMesh":
Copy link
Collaborator

Choose a reason for hiding this comment

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

can we add some documentation?

Also I am thinking how we should support multi-dim from group, i.e. HSDP shard, replicate group to pass in to construct a DeviceMesh, I guess the group argument would likely need to be Union[ProcessGroup, List[ProcessGroup]]

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Also I am thinking how we should support multi-dim from group, i.e. HSDP shard, replicate group to pass in to construct a DeviceMesh, I guess the group argument would likely need to be Union[ProcessGroup, List[ProcessGroup]]

This makes sense to me. Shall we expand the API later when we need it in a separate PR?

Copy link
Collaborator

Choose a reason for hiding this comment

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

yeah we don't need to do it in this PR as it's BC compatible if later we expand the API

We need `DeviceMesh.from_group()` to allow FSDP2 to interoperate with distributed libraries that do not use `DeviceMesh` for all parallelisms.

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k

[ghstack-poisoned]
@awgu awgu changed the title [POC] Showed DeviceMesh.from_group() prototype [DeviceMesh] Added DeviceMesh.from_group() Apr 24, 2024
@awgu awgu marked this pull request as ready for review April 24, 2024 17:25
@awgu awgu requested a review from wz337 April 24, 2024 17:25
This PR adds a `DeviceMesh.from_group()` static method to convert an existing process group to a device mesh.

Motivation: We need `DeviceMesh.from_group()` to allow FSDP2 to interoperate with distributed libraries that do not use `DeviceMesh` for all parallelisms.

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k

[ghstack-poisoned]
awgu pushed a commit that referenced this pull request Apr 24, 2024
ghstack-source-id: 9c6b09f
Pull Request resolved: #124787
@awgu
Copy link
Collaborator Author

awgu commented Apr 24, 2024

cc: @wanchaol I added some documentation and a simple unit test in test_device_mesh.py too.

@awgu awgu added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 24, 2024
@awgu
Copy link
Collaborator Author

awgu commented Apr 24, 2024

@pytorchbot 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

alat-rights pushed a commit to alat-rights/pytorch that referenced this pull request Apr 26, 2024
This PR adds a `DeviceMesh.from_group()` static method to convert an existing process group to a device mesh.

Motivation: We need `DeviceMesh.from_group()` to allow FSDP2 to interoperate with distributed libraries that do not use `DeviceMesh` for all parallelisms.

Pull Request resolved: pytorch#124787
Approved by: https://github.com/wanchaol
ghstack dependencies: pytorch#124651, pytorch#124741, pytorch#124767, pytorch#124768, pytorch#124780
pytorchmergebot pushed a commit that referenced this pull request Apr 29, 2024
This PR renames the `FSDP` class to `FSDPModule`. This is a BC breaking change. The rationale is that `FSDPModule` is more descriptive since `fully_shard` is a module-level API (applied to a `module` arg), so the `FSDP` class will always correspond to a module.

Also, users commonly import `FullyShardedDataParallel` as `FSDP`, so this can help avoid some name conflict in some cases.

Pull Request resolved: #124955
Approved by: https://github.com/wanchaol, https://github.com/wconstab
ghstack dependencies: #124651, #124741, #124767, #124768, #124780, #124787
pytorch-bot bot pushed a commit that referenced this pull request May 3, 2024
This PR adds a `DeviceMesh.from_group()` static method to convert an existing process group to a device mesh.

Motivation: We need `DeviceMesh.from_group()` to allow FSDP2 to interoperate with distributed libraries that do not use `DeviceMesh` for all parallelisms.

Pull Request resolved: #124787
Approved by: https://github.com/wanchaol
ghstack dependencies: #124651, #124741, #124767, #124768, #124780
pytorch-bot bot pushed a commit that referenced this pull request May 3, 2024
This PR renames the `FSDP` class to `FSDPModule`. This is a BC breaking change. The rationale is that `FSDPModule` is more descriptive since `fully_shard` is a module-level API (applied to a `module` arg), so the `FSDP` class will always correspond to a module.

Also, users commonly import `FullyShardedDataParallel` as `FSDP`, so this can help avoid some name conflict in some cases.

Pull Request resolved: #124955
Approved by: https://github.com/wanchaol, https://github.com/wconstab
ghstack dependencies: #124651, #124741, #124767, #124768, #124780, #124787
@github-actions github-actions bot deleted the gh/awgu/571/head branch June 3, 2024 01:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ci-td-distributed ciflow/trunk Trigger trunk jobs on your pull request Merged oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: DeviceMesh
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants