[dtensor][8/N] Introduce cost model for sharding #109145
Closed
+269
−9
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
This PR adds some basic comm cost model for sharding prop [ghstack-poisoned]
wanchaol
requested review from
mrshenli,
zhaojuanmao,
rohan-varma,
H-Huang,
awgu,
kwen2501,
fegin,
fduwjj and
kiukchung
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/109145
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 2192bdb with merge base 898482f ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This was referenced Sep 12, 2023
wanchaol
added a commit
that referenced
this pull request
Sep 12, 2023
This PR adds some basic comm cost model for sharding prop ghstack-source-id: 0efce2c591ff42d05e79a8fd03c365926c096cb4 Pull Request resolved: #109145
This PR adds some basic comm cost model for sharding prop [ghstack-poisoned]
wanchaol
added a commit
that referenced
this pull request
Sep 15, 2023
This PR adds some basic comm cost model for sharding prop ghstack-source-id: c45e652b428ee5ec1cb021975307e49dda6ffb7a Pull Request resolved: #109145
This PR adds some basic comm cost model for sharding prop [ghstack-poisoned]
wanchaol
added a commit
that referenced
this pull request
Oct 2, 2023
This PR adds some basic comm cost model for sharding prop ghstack-source-id: 1b99d3b8399e0944a2d849aedcd7b17d5bf0a70a Pull Request resolved: #109145
Chillee
reviewed
Oct 2, 2023
This PR adds some basic comm cost model for sharding prop. Why we need this? operators can generate multiple placement strategies, i.e. for matmul we have at least 4 possible shardings: `1. Shard(0), R 2. R, Shard(1) 3. Shard(1), Shard(0) 4. R, R` We need to be able to choose from one of these options during runtime, and perform resharding with reasonable choices. This is why we are building a cost model for sharding here. In this PR we associate each possible sharding strategy with redistribute costs. For eager mode since we run ops eagerly, we simply perform a min cost selection. One can imagine if we have some global information the strategy selection would become more intelligient. [ghstack-poisoned]
This PR adds some basic comm cost model for sharding prop. Why we need this? operators can generate multiple placement strategies, i.e. for matmul we have at least 4 possible shardings: `1. Shard(0), R 2. R, Shard(1) 3. Shard(1), Shard(0) 4. R, R` We need to be able to choose from one of these options during runtime, and perform resharding with reasonable choices. This is why we are building a cost model for sharding here. In this PR we associate each possible sharding strategy with redistribute costs. For eager mode since we run ops eagerly, we simply perform a min cost selection. One can imagine if we have some global information the strategy selection would become more intelligient. [ghstack-poisoned]
This PR adds some basic comm cost model for sharding prop. Why we need this? operators can generate multiple placement strategies, i.e. for matmul we have at least 4 possible shardings: `1. Shard(0), R 2. R, Shard(1) 3. Shard(1), Shard(0) 4. R, R` We need to be able to choose from one of these options during runtime, and perform resharding with reasonable choices. This is why we are building a cost model for sharding here. In this PR we associate each possible sharding strategy with redistribute costs. For eager mode since we run ops eagerly, we simply perform a min cost selection. One can imagine if we have some global information the strategy selection would become more intelligient. [ghstack-poisoned]
fduwjj
reviewed
Oct 13, 2023
fduwjj
reviewed
Oct 13, 2023
fduwjj
reviewed
Oct 14, 2023
fduwjj
reviewed
Oct 14, 2023
This PR adds some basic comm cost model for sharding prop [ghstack-poisoned]
This PR adds some basic comm cost model for sharding prop [ghstack-poisoned]
fduwjj
reviewed
Oct 14, 2023
This PR adds some basic comm cost model for sharding prop [ghstack-poisoned]
pytorchmergebot
pushed a commit
that referenced
this pull request
Oct 15, 2023
This PR switches matrix ops to generate the sharding strategies, and with the cost selection algorithm introduced in the previous PR we are able to enable this and more ops to leverage strategy based sharding prop This also fixes a bunch of corner cases that existing propagation does not cover, resulting in full coverage for baddbmm Pull Request resolved: #110717 Approved by: https://github.com/fduwjj ghstack dependencies: #109145
pytorchmergebot
pushed a commit
that referenced
this pull request
Oct 15, 2023
As titled, this also handles sth like [Shard(0), Shard(0)] correctly for pointwise ops, which was previously errored out Pull Request resolved: #111234 Approved by: https://github.com/fduwjj ghstack dependencies: #109145, #110717
pytorchmergebot
pushed a commit
that referenced
this pull request
Oct 15, 2023
This add __Str__ to op schema and dtensor spec for ease of reading Pull Request resolved: #111278 Approved by: https://github.com/fduwjj ghstack dependencies: #109145, #110717, #111234
yeounoh
pushed a commit
to yeounoh/pytorch
that referenced
this pull request
Oct 16, 2023
This PR adds some basic comm cost model for sharding prop Pull Request resolved: pytorch#109145 Approved by: https://github.com/fduwjj
yeounoh
pushed a commit
to yeounoh/pytorch
that referenced
this pull request
Oct 16, 2023
This PR switches matrix ops to generate the sharding strategies, and with the cost selection algorithm introduced in the previous PR we are able to enable this and more ops to leverage strategy based sharding prop This also fixes a bunch of corner cases that existing propagation does not cover, resulting in full coverage for baddbmm Pull Request resolved: pytorch#110717 Approved by: https://github.com/fduwjj ghstack dependencies: pytorch#109145
yeounoh
pushed a commit
to yeounoh/pytorch
that referenced
this pull request
Oct 16, 2023
As titled, this also handles sth like [Shard(0), Shard(0)] correctly for pointwise ops, which was previously errored out Pull Request resolved: pytorch#111234 Approved by: https://github.com/fduwjj ghstack dependencies: pytorch#109145, pytorch#110717
yeounoh
pushed a commit
to yeounoh/pytorch
that referenced
this pull request
Oct 16, 2023
This add __Str__ to op schema and dtensor spec for ease of reading Pull Request resolved: pytorch#111278 Approved by: https://github.com/fduwjj ghstack dependencies: pytorch#109145, pytorch#110717, pytorch#111234
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):
This PR adds some basic comm cost model for sharding prop