[DTensor] require DeviceMesh size equals world size #91801
Conversation
[ghstack-poisoned]
XilunWu
requested review from
mrshenli,
pritamdamania87,
zhaojuanmao,
rohan-varma,
H-Huang,
awgu,
kwen2501 and
wanchaol
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/91801
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 387458b: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
[ghstack-poisoned]
|
Two meshes created over world in |
| @@ -143,6 +143,13 @@ def __init__( | |||
| f"Mesh should not be bigger than default world size, but found {self.mesh.numel()} ranks!" | |||
| ) | |||
|
|
|||
| # TODO: we will support mesh on a subset of WORLD in future | |||
| if self.mesh.numel() < world_size: | |||
There was a problem hiding this comment.
this check should not happen all the time, it should only happen when there's no default pg exist and we want to help user create a world_pg, this check should be only inside get_or_create_group I think
There was a problem hiding this comment.
Make sense because IIRC we can define for example mesh A on rank 0, 1 and mesh B on rank 2, 3. The example we discussed last time is actually about mesh is defined on rank 0, 1 and no mesh is defined on rank 2, 3. Right?
There was a problem hiding this comment.
yeah, it's possibly to create sub meshes, that's what 2-D did currently, so we should still allow such behavior.
[ghstack-poisoned]
|
@pytorchmergebot merge -g |
Merge startedYour change will be merged once all checks on your PR pass since you used the green (-g) flag (ETA: 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Stack from ghstack (oldest at bottom):