Avoid nn.Linear decomposition by replacing view -> mm -> view with einsum #26
+133
−1
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
This tries to support CP-style sharding, by overcoming a limitation of DTensor. Doesn't yet work as _mm_strategy is failing
facebook-github-bot
added
the
CLA Signed
wconstab
reviewed
Jul 2, 2025
autoparallel/propagation_rules.py
Outdated
|
|
||
| @register_opschema_rule(torch.ops.aten.matmul.default) | ||
| def matmul_rule(mesh, op_schema): | ||
| # from torch.distributed.tensor._ops._einsum_strategy import gen_einsum_strategies |
Contributor
There was a problem hiding this comment.
i would have thought to use the einsum strategies here. for my education, what is the difference between einsum and mm_like in this context? cc @XilunWu
Contributor
Author
There was a problem hiding this comment.
I'll have to end-up using the einsum strategies, because mm_strategies fail :-)
The difference I believe is just that the mm_strategy filter out invalid strategies already, while einsum strategies don't as it doesn't know the size of the tensors, only the number of dimensions
…sa/replace_view_mm_view
Somethings are starting to work, but we are not yet there
…sa/replace_view_mm_view
…sa/replace_view_mm_view
Before this, if we had a list of tensors we wouldn't shard the tensors inside the list
ezyang
reviewed
Aug 13, 2025
ezyang
approved these changes
Aug 13, 2025
This removes a long-standing hack to tell the solver that S(1) -> R is more expensive than S(0) -> R because of an additional data movement
bdhirsh
reviewed
Aug 13, 2025
This removes a long-standing hack to tell the solver that S(1) -> R is more expensive than S(0) -> R because of an additional data movement
…sa/compute_cost_in_comms
…sa/compute_cost_in_comms
…h/autoparallel into fmassa/replace_view_mm_view
fmassa
added a commit
that referenced
this pull request
Aug 21, 2025
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
This was referenced Aug 21, 2025
fmassa
added a commit
that referenced
this pull request
Aug 21, 2025
* Support tuple of tensors in estimate_strategy_runtime_cost 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 * Fix bad copy-paste
…sa/compute_cost_in_comms
…h/autoparallel into fmassa/replace_view_mm_view
…sa/compute_cost_in_comms
…h/autoparallel into fmassa/replace_view_mm_view
fmassa
changed the title
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.
PyTorch currently decomposes any 3d-input
nn.Linear(and matmul) into a sequence ofview -> mm -> viewoperations.This has as a consequence of breaking any type of sharding on both the batch and the sequence dimension, because the flattening that happens doesn't allow to preserve this sharding.
While we wait for PyTorch to avoid decomposing
nn.Linear, we instead take the route of pattern-matching thenn.Linearspecific occurences, and we replace them with aneinsumoperator.We perform this pattern-matching replacement for both the forward as well as the backward pass.
For now, the pass is disabled by default, and can be enabled via a global flag. I'm leaving it disabled for now by default because this actually requires changing some other things like improving the cost model as in #94, so I'm keeping the behavior the same for now while I experiment with the other things more easily