[DTensor] implement dist_cat as a sharding prop rule #92677
Conversation
[ghstack-poisoned]
XilunWu
requested review from
H-Huang,
awgu,
kwen2501,
mrshenli,
rohan-varma,
wanchaol and
zhaojuanmao
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/92677
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 3ea6b8c: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
[ghstack-poisoned]
wanchaol
left a comment
wanchaol
left a comment
There was a problem hiding this comment.
See comments inlined, we should make sure the xfail of cat to be removed so that it passes all the possible cases.
| return output_sharding | ||
|
|
||
|
|
||
| def _update_schema_suggestion_for_cat( |
There was a problem hiding this comment.
can you tell me what exactly this function is doing? it looks like a lot of duplicate logic with the rule itself and I am not quite sure what this function is used for.
There was a problem hiding this comment.
einop_rule expects the op_schema argument to have its args_schema in form [DTensorSpec, DTensorSpec, ...] but when it's passed into cat_rule the schema is actually [List[DTensorSpec]]. That's why I convert the args_schema at the beginning of cat_rule (https://github.com/pytorch/pytorch/pull/92677/files#diff-ebc7be1151cf411ce7edf46c4ca1cabb74cd953a2bdf47e04b4cc733c31f6085R492) before feeding it into einop_rule. Thus, we need to convert it back if a schema_suggestion is present here.
[ghstack-poisoned]
[ghstack-poisoned]
wanchaol
left a comment
wanchaol
left a comment
There was a problem hiding this comment.
lgtm, thanks for working on it! left a couple of suggestions and some question.
| dim_word = free_dim[:dim] + alphabet[i] + free_dim[dim:] | ||
| einop_notation_list.append(dim_word) | ||
| else: | ||
| einop_notation_list.append(alphabet[i]) |
There was a problem hiding this comment.
is this the empty tensor annotation where it have a single char?
There was a problem hiding this comment.
Not entirely for empty tensor but empty tensor whose ndim is smaller than other tensors. This is for case like concatenating Tensor([], shape=torch.Size([0])) with Tensor([[1, 2], [3, 4]], shape=torch.Size([2, 2])).
In this case, an empty annotation may still work but we want to ensure that the dim char for cat_dim in output tensor annotation must appear in input as well. Adding each input tensor's cat_dim dim char into annotation guarantees that.
[ghstack-poisoned]
|
@pytorchmergebot merge -g |
Merge startedYour 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 |
Stack from ghstack (oldest at bottom):