[DTensor] implement dist_cat as a sharding prop rule #92677
Conversation
[ghstack-poisoned]
XilunWu
requested review from
mrshenli,
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/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-source-id: a27c5a3dffcb7b29dc8d22de45013a199f2ffb00 Pull Request resolved: #92677
[ghstack-poisoned]
ghstack-source-id: d4cbf117b4e630d80fe39229251d73724fae55cc Pull Request resolved: #92677
| 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-source-id: 8ec685e5998e48321a16a752d2b1a7c5a6c84ed4 Pull Request resolved: #92677
[ghstack-poisoned]
ghstack-source-id: f536f348cc4ed8e049eb9bdf1462415c18a89839 Pull Request resolved: #92677
| 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]
ghstack-source-id: 4c2f6291dbcdad2b1674b3db36dc46a21a9159ce Pull Request resolved: #92677
|
@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):