-
Notifications
You must be signed in to change notification settings - Fork 528
Support DI sharding for FPE_EBC #2968
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
This pull request was exported from Phabricator. Differential Revision: D74671655 |
Summary: Support models that have FeatureProcessedEmbeddingBagCollection. These changes make sure we add QuantFeatureProcessedEmbeddingBagCollectionSharder as a recognized sharder, handle multiple envs needed for specifying DI sharding, and propagate TBE properly when processing the sharding plan. This doesn't support true hybrid sharding yet, so FPE_EBCs must be all sharded with the same (sharding_type, device). Differential Revision: D74671655
This pull request was exported from Phabricator. Differential Revision: D74671655 |
Summary: Support models that have FeatureProcessedEmbeddingBagCollection to be DI sharded. However, conservatively enforce that FPE itself can only be sharded on HBM and not across CPU as well. These changes make sure we add QuantFeatureProcessedEmbeddingBagCollectionSharder as a recognized sharder, handle multiple envs needed for specifying DI sharding, and propagate TBE properly when processing the sharding plan. This doesn't support true hybrid sharding for FPE. Differential Revision: D74671655
This pull request was exported from Phabricator. Differential Revision: D74671655 |
Summary: Support models that have FeatureProcessedEmbeddingBagCollection to be DI sharded. However, conservatively enforce that FPE itself can only be sharded on HBM and not across CPU as well. These changes make sure we add QuantFeatureProcessedEmbeddingBagCollectionSharder as a recognized sharder, handle multiple envs needed for specifying DI sharding, and propagate TBE properly when processing the sharding plan. This doesn't support true hybrid sharding for FPE. Differential Revision: D74671655
This pull request was exported from Phabricator. Differential Revision: D74671655 |
Summary: Support models that have FeatureProcessedEmbeddingBagCollection to be DI sharded. However, conservatively enforce that FPE itself can only be sharded on HBM and not across CPU as well. These changes make sure we add QuantFeatureProcessedEmbeddingBagCollectionSharder as a recognized sharder, handle multiple envs needed for specifying DI sharding, and propagate TBE properly when processing the sharding plan. This doesn't support true hybrid sharding for FPE. Differential Revision: D74671655
This pull request was exported from Phabricator. Differential Revision: D74671655 |
Summary: Pull Request resolved: pytorch#2968 Support models that have FeatureProcessedEmbeddingBagCollection to be DI sharded. However, conservatively enforce that FPE itself can only be sharded on HBM and not across CPU as well. These changes make sure we add QuantFeatureProcessedEmbeddingBagCollectionSharder as a recognized sharder, handle multiple envs needed for specifying DI sharding, and propagate TBE properly when processing the sharding plan. This doesn't support true hybrid sharding for FPE. Differential Revision: D74671655
Summary: Support models that have FeatureProcessedEmbeddingBagCollection to be DI sharded. However, conservatively enforce that FPE itself can only be sharded on HBM and not across CPU as well. These changes make sure we add QuantFeatureProcessedEmbeddingBagCollectionSharder as a recognized sharder, handle multiple envs needed for specifying DI sharding, and propagate TBE properly when processing the sharding plan. This doesn't support true hybrid sharding for FPE. Differential Revision: D74671655
This pull request was exported from Phabricator. Differential Revision: D74671655 |
assert ( | ||
shard_device_type == expected_device_type | ||
), f"Expected {expected_device_type} but got {shard_device_type} for FeatureProcessedEmbeddingBagCollection sharding device type" | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
since we assert here and the expected_device_type
is always cuda
, why don't we keep the signature of the method's env as a single ShardingEnv
, and just pass the cuda
's sharding evn to this?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
While right now FPE is used on cuda, we can support it on cuda/cpu in the future. Performing the check here allows us to de-duplicate the logic so we can ensure there's one sharding type for FPE for all sharding paths.
This pull request was exported from Phabricator. Differential Revision: D74671655 |
527d6cf
to
14c9ff0
Compare
Summary: Pull Request resolved: pytorch#2968 Support models that have FeatureProcessedEmbeddingBagCollection to be DI sharded. However, conservatively enforce that FPE itself can only be sharded on HBM and not across CPU as well. These changes make sure we add QuantFeatureProcessedEmbeddingBagCollectionSharder as a recognized sharder, handle multiple envs needed for specifying DI sharding, and propagate TBE properly when processing the sharding plan. This doesn't support true hybrid sharding for FPE. Reviewed By: faran928 Differential Revision: D74671655
Summary: Pull Request resolved: pytorch#2968 Support models that have FeatureProcessedEmbeddingBagCollection to be DI sharded. However, conservatively enforce that FPE itself can only be sharded on HBM and not across CPU as well. These changes make sure we add QuantFeatureProcessedEmbeddingBagCollectionSharder as a recognized sharder, handle multiple envs needed for specifying DI sharding, and propagate TBE properly when processing the sharding plan. This doesn't support true hybrid sharding for FPE. Reviewed By: faran928 Differential Revision: D74671655
This pull request was exported from Phabricator. Differential Revision: D74671655 |
14c9ff0
to
2ec92ba
Compare
Summary:
Support models that have FeatureProcessedEmbeddingBagCollection.
These changes make sure we add QuantFeatureProcessedEmbeddingBagCollectionSharder as a recognized sharder, handle multiple envs needed for specifying DI sharding, and propagate TBE properly when processing the sharding plan.
Differential Revision: D74671655