Support tuple of tensors in estimate_strategy_runtime_cost #102
Conversation
This file contains hidden or 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
Previously, if we had tuple of tensors as an argument to a function, we wouldn't apply any sharding on it. This is split from #26 , where I originally found this issue
meta-cla
bot
added
the
CLA Signed
xmfan
approved these changes
Aug 21, 2025
wconstab
reviewed
Aug 21, 2025
| args = tree_map_only(torch.fx.Node, lambda x: x.meta["val"], node.args) | ||
| kwargs = tree_map_only(torch.fx.Node, lambda x: x.meta["val"], node.kwargs) | ||
|
|
||
| fake_mode = torch._guards.detect_fake_mode(args) |
Contributor
There was a problem hiding this comment.
is this bc we are already under a fake mode now, but we weren't in the initial autop?
Contributor
Author
There was a problem hiding this comment.
Yes, that's right, now the whole AutoParallel is running under fake mode, so we can remove it
xmfan
reviewed
Aug 21, 2025
| args = tree_map_only(torch.fx.Node, lambda x: x.meta["val"], node.args) | ||
| kwargs = tree_map_only(torch.fx.Node, lambda x: x.meta["val"], node.kwargs) | ||
|
|
||
| fake_mode = torch._guards.detect_fake_mode(args) |
Member
There was a problem hiding this comment.
this is removed because we're already running this in a fake mode?
wconstab
reviewed
Aug 21, 2025
| _get_sharded_shape_stride(spec) for spec in strategy.input_specs | ||
| ) | ||
|
|
||
| flat_args, treespec = tree_flatten(args) |
Contributor
There was a problem hiding this comment.
shouldn't we just call tree_map_only(Tensor, torch.empty) here instead of doing the for loop?
Contributor
Author
There was a problem hiding this comment.
We need to get the size from the args_sizes_strides (which comes from the spec), so I think we might need this indirection.
But if there is a cleaner way of doing this I'm happy to change the code!
wconstab
approved these changes
Aug 21, 2025
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.
Previously, if we had tuple of tensors as an argument to a function, we wouldn't apply any sharding on it. This is split from #26 , where I originally found this issue