[DTensor] require DeviceMesh size equals world size #91801
Conversation
[ghstack-poisoned]
XilunWu
requested review from
H-Huang,
awgu,
kwen2501,
mrshenli,
pritamdamania87,
rohan-varma,
wanchaol and
zhaojuanmao
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 |
wanchaol
left a comment
wanchaol
left a comment
There was a problem hiding this comment.
Let's not always do the check and do it only when we initialize world_pg
| ) | ||
|
|
||
| # 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):