[dtensor] refactor view ops to use OpStrategy #126011
Closed
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
[ghstack-poisoned]
pytorch-bot
bot
added
ciflow/inductor
oncall: distributed
tianyu-l
added a commit
that referenced
this pull request
May 11, 2024
ghstack-source-id: 1767d9ed7dc1b77e7a83f7be3fa5caa158cfc5c2 Pull Request resolved: #126011
Previously, the rule-based view op sharding prop adjust a non-tensor arg `local_out_shape` within the rule itself. This was not viable in strategy-based sharding prop. Thus, this PR is adding a new option `non_tensor_arg_suggestions` into `PlacementStrategy` to address this problem. It also benefits the new factory ops in that we no longer need to compute their local shape and stride in `sharding_prop.py` in a customized way. Instead, we compute all such expected **tensor and non-tensor args** in `tensor_ops.py` and `view_ops.py`, which keeps the `sharding_prop.py` clean. cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 wconstab yf225 chauhang d4l3k [ghstack-poisoned]
tianyu-l
added a commit
that referenced
this pull request
May 13, 2024
ghstack-source-id: e7b1885f5fb24b8198e6e4b565d35cbc92a54b4d Pull Request resolved: #126011
tianyu-l
added
ciflow/trunk
wanchaol
reviewed
May 14, 2024
|
looks like CI is failing with circular deps, please fix |
As titled. Some ops require adjustment of output shape argument. In rule-based sharding prop, global output shape was inferred in the rule (in `view_ops.py`). In strategy-based sharding prop, it is now obtained from propagated out_tensor_meta (in `sharding_prop.py`). cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 wconstab yf225 chauhang d4l3k [ghstack-poisoned]
tianyu-l
added a commit
that referenced
this pull request
May 15, 2024
ghstack-source-id: 94c584701a1f79c1314fb4b2a602a5cb87cf0f28 Pull Request resolved: #126011
wanchaol
approved these changes
May 15, 2024
| self.assertEqual(view_as_complex_rule, expected_view_as_complex_rule) | ||
| expected_view_as_real_rule = ( | ||
| InputDim(0), | ||
| Split(InputDim(1), (13, 2), 0), | ||
| Split(InputDim(1), (13, 2), 1), | ||
| ) | ||
| view_as_real_rule = ops[torch.view_as_real].dim_map(intermediate) | ||
| view_as_real_rule = dim_maps[torch.view_as_real](intermediate) |
There was a problem hiding this comment.
I think we'll need to improve some of the existing test cases in this file to work with CommDebugMode, to make sure the communication that are happening is expected. This can be done in a follow up PR
As titled. Some ops require adjustment of output shape argument. In rule-based sharding prop, global output shape was inferred in the rule (in `view_ops.py`). In strategy-based sharding prop, it is now obtained from propagated out_tensor_meta (in `sharding_prop.py`). cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 wconstab yf225 chauhang d4l3k [ghstack-poisoned]
As titled. Some ops require adjustment of output shape argument. In rule-based sharding prop, global output shape was inferred in the rule (in `view_ops.py`). In strategy-based sharding prop, it is now obtained from propagated out_tensor_meta (in `sharding_prop.py`). cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 wconstab yf225 chauhang d4l3k [ghstack-poisoned]
XilunWu
approved these changes
May 16, 2024
|
@pytorchbot merge |
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 |
ZelboK
pushed a commit
to ZelboK/pytorch
that referenced
this pull request
May 19, 2024
As titled. Some ops require adjustment of output shape argument. In rule-based sharding prop, global output shape was inferred in the rule (in `view_ops.py`). In strategy-based sharding prop, it is now obtained from propagated out_tensor_meta (in `sharding_prop.py`). Pull Request resolved: pytorch#126011 Approved by: https://github.com/wanchaol, https://github.com/XilunWu
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
output_prefix from OpStrategy properties #126359As titled. Some ops require adjustment of output shape argument. In rule-based sharding prop, global output shape was inferred in the rule (in
view_ops.py). In strategy-based sharding prop, it is now obtained from propagated out_tensor_meta (insharding_prop.py).cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @wconstab @yf225 @chauhang @d4l3k