uneven sharding for sharded tensor pool #3533
Open
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
Contributor
meta-cla
bot
added
the
CLA Signed
faran928
force-pushed
the
export-D79603009
branch
from
708ddb1 to
3f268a1
Compare
… pool (meta-pytorch#3533) Summary: A few changes in the diff: 1. Support to proportionally shard the tensor pool based on memory capacity per rank. 2. Using block_bucketize_sparse_features_inference to return bucket_mapping that can be used during request batching in inference w/ custom sigrid predictor engine 3. Wrapping some of the operations with fx wrappers to make it compatible with model split boundaries for DLRM serving where embeddings are sharded and split onto different pytorch modules 4. Exposing set_device() api to some of the modules if we want to place some shards to cpu while others to cuda. 5. Move _get_unbucketize_tensor_via_length_alignment to common util files. Differential Revision: D79603009
faran928
force-pushed
the
export-D79603009
branch
from
3f268a1 to
4e20f9e
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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: trochrec uneven sharding changes
Differential Revision: D79603009