Enable Changing the # of shards for CW resharding #3188
Closed
+346
−106
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
facebook-github-bot
added
the
CLA Signed
|
This pull request was exported from Phabricator. Differential Revision: D78291717 |
isururanawaka
force-pushed
the
export-D78291717
branch
from
9c50718 to
1d44ff7
Compare
isururanawaka
added a commit
to isururanawaka/torchrec
that referenced
this pull request
Jul 15, 2025
Summary: Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
|
This pull request was exported from Phabricator. Differential Revision: D78291717 |
isururanawaka
force-pushed
the
export-D78291717
branch
from
1d44ff7 to
8411bc3
Compare
isururanawaka
added a commit
to isururanawaka/torchrec
that referenced
this pull request
Jul 15, 2025
Summary: Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
|
This pull request was exported from Phabricator. Differential Revision: D78291717 |
isururanawaka
added a commit
to isururanawaka/torchrec
that referenced
this pull request
Jul 15, 2025
Summary: Pull Request resolved: pytorch#3188 Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
isururanawaka
force-pushed
the
export-D78291717
branch
from
8411bc3 to
cae0e53
Compare
isururanawaka
added a commit
to isururanawaka/torchrec
that referenced
this pull request
Jul 15, 2025
Summary: Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
isururanawaka
force-pushed
the
export-D78291717
branch
from
cae0e53 to
5bd9cf7
Compare
|
This pull request was exported from Phabricator. Differential Revision: D78291717 |
isururanawaka
added a commit
to isururanawaka/torchrec
that referenced
this pull request
Jul 15, 2025
Summary: Pull Request resolved: pytorch#3188 Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
isururanawaka
force-pushed
the
export-D78291717
branch
from
5bd9cf7 to
d2c1ed9
Compare
isururanawaka
added a commit
to isururanawaka/torchrec
that referenced
this pull request
Jul 15, 2025
Summary: Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
isururanawaka
force-pushed
the
export-D78291717
branch
from
d2c1ed9 to
3abf13f
Compare
isururanawaka
added a commit
to isururanawaka/torchrec
that referenced
this pull request
Jul 15, 2025
Summary: Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
isururanawaka
force-pushed
the
export-D78291717
branch
from
3abf13f to
c85ea45
Compare
isururanawaka
added a commit
to isururanawaka/torchrec
that referenced
this pull request
Jul 16, 2025
Summary: Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
isururanawaka
force-pushed
the
export-D78291717
branch
from
c85ea45 to
99e2dfc
Compare
|
This pull request was exported from Phabricator. Differential Revision: D78291717 |
1 similar comment
|
This pull request was exported from Phabricator. Differential Revision: D78291717 |
isururanawaka
force-pushed
the
export-D78291717
branch
from
99e2dfc to
d1b7766
Compare
isururanawaka
added a commit
to isururanawaka/torchrec
that referenced
this pull request
Jul 16, 2025
Summary: Pull Request resolved: pytorch#3188 Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
isururanawaka
added a commit
to isururanawaka/torchrec
that referenced
this pull request
Jul 16, 2025
Summary: Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
isururanawaka
force-pushed
the
export-D78291717
branch
from
d1b7766 to
52f8a46
Compare
Summary: Pull Request resolved: pytorch#3188 Currently Dynamic Sharding assumes the # of shards per embedding table stays the same: - https://www.internalfb.com/code/fbsource/[6d270632037a1e8bca7f63500dd07fd0b213e572]/fbcode/torchrec/distributed/sharding/dynamic_sharding.py?lines=140 E.g. - `table_0` originally sharded on ranks: [0, 1] - Reshard API currently supports moving `table_0` shards to ranks [1, 2]. - Where `the shard` on rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2 We want to support changing the # of shards: - e.g. table_0 originally on ranks: [0, 1] --> reshard to [0] - Or reshard to [0, 1, 2, 3] Here's the unit test you can modify to check if your usecase passes: - https://www.internalfb.com/code/fbsource/[4d0d74b9f3c441e7aa35ce7102200fa0ca8c95cf]/fbcode/torchrec/distributed/tests/test_dynamic_sharding.py?lines=453-459 - Basically change the new sharding plan to be a different # of ranks than the original sharding plan. Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table - e.g. emb_table size: [4, 8], this can only be sharded on 1, 2, or 4 ranks. not 3 ranks Differential Revision: D78291717
|
This pull request was exported from Phabricator. Differential Revision: D78291717 |
isururanawaka
force-pushed
the
export-D78291717
branch
from
52f8a46 to
bbb5e37
Compare
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.
Summary:
Currently Dynamic Sharding assumes the # of shards per embedding table stays the same:
E.g.
table_0originally sharded on ranks: [0, 1]table_0shards to ranks [1, 2].the shardon rank 0 will move to rank 1, and the shard on rank 1 will move to rank 2We want to support changing the # of shards:
Here's the unit test you can modify to check if your usecase passes:
Note: the new total number of ranks for each embedding table should be a factor of the dimension 0 of that embedding table
Differential Revision: D78291717