[DTensor] single-dim foreach strategy#170631
[DTensor] single-dim foreach strategy#170631pianpwk wants to merge 4 commits intogh/pianpwk/46/basefrom
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/170631
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit b2ec35b with merge base 61622da ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| return OpSchema( | ||
| target_op, # type: ignore[arg-type] | ||
| args_schema=tuple(target_args), | ||
| kwargs_schema=op_schema.kwargs_schema, |
There was a problem hiding this comment.
note: I don't know of if there are any cases where kwargs_schema could contain a TupleStrategy and that would need to also be translated.
There was a problem hiding this comment.
just changed to tree_map instead
| child_strategies: list[StrategyType] = [] | ||
| for tensorlist_i in range(tensorlist_len): | ||
| per_index_schema = _translate_foreach_op_schema(op_schema, tensorlist_i) | ||
| per_index_strategy = _expanded_strategy(per_index_schema) |
There was a problem hiding this comment.
wondering if it is worth doing LRU cache on this helper. It seems like a very common case to have 100s or 1000s of tensors in the foreach list, but all of them sharing the same placements, or only a handful of patterns for placements.
There was a problem hiding this comment.
ah, I tried this but it seems like it needs more work; the OpSchema._comparison_key that's used for caching doesn't include placements, so out-of-box it's over-caching and incorrect
I added a todo there
wconstab
left a comment
There was a problem hiding this comment.
this looks pretty good to me! thanks!
i think we can land it asap, i will work on landing the earlier PRs in the stack
|
|
||
| # TODO maybe this could be helped by adding a new 'tag' to the OpOverload? | ||
| # Also, i'm guessing that i'll need more info from the registration callsite | ||
| # about which inputs are expected to be lists vs tensors. But maybe I can just infer it all from the runtime |
There was a problem hiding this comment.
i guess you already handled this part (Also..) so you can delete it
wconstab
left a comment
There was a problem hiding this comment.
lgtm! i think we can land it and continue to iterate on the todos.
up to you whether to land my pointwise PR too, which I've rebased recently and made safe to land, or just copy its contents into your test for now
This PR adds a dummy version of a single-dim pointwise strategy, and it shouldn't be used yet for real but it is being used by unit tests in #170631 so it is useful to land. Pull Request resolved: #168115 Approved by: https://github.com/pianpwk Co-authored-by: Pian Pawakapan <pianpwk@meta.com>
|
@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 |
This PR adds a dummy version of a single-dim pointwise strategy, and it shouldn't be used yet for real but it is being used by unit tests in #170631 so it is useful to land. Pull Request resolved: #168115 Approved by: https://github.com/pianpwk Co-authored-by: Pian Pawakapan <pianpwk@meta.com>
This PR adds a dummy version of a single-dim pointwise strategy, and it shouldn't be used yet for real but it is being used by unit tests in pytorch#170631 so it is useful to land. Pull Request resolved: pytorch#168115 Approved by: https://github.com/pianpwk Co-authored-by: Pian Pawakapan <pianpwk@meta.com>
Pull Request resolved: pytorch#170631 Approved by: https://github.com/wconstab
Stack from ghstack (oldest at bottom):