Add EC ZCH to default sharder (pytorch#1610)

Summary:

Added the EC ZCH module to the get_default_sharders method

Reviewed By: henrylhtsang

Differential Revision: D52541735

torchrec/distributed/sharding_plan.py

@@ -21,6 +21,7 @@
     FeatureProcessedEmbeddingBagCollectionSharder,
 )
 from torchrec.distributed.fused_embeddingbag import FusedEmbeddingBagCollectionSharder
+from torchrec.distributed.mc_embedding import ManagedCollisionEmbeddingCollectionSharder
 from torchrec.distributed.mc_embeddingbag import (
     ManagedCollisionEmbeddingBagCollectionSharder,
 )
@@ -47,6 +48,7 @@ def get_default_sharders() -> List[ModuleSharder[nn.Module]]:
         cast(ModuleSharder[nn.Module], QuantEmbeddingBagCollectionSharder()),
         cast(ModuleSharder[nn.Module], QuantEmbeddingCollectionSharder()),
         cast(ModuleSharder[nn.Module], ManagedCollisionEmbeddingBagCollectionSharder()),
+        cast(ModuleSharder[nn.Module], ManagedCollisionEmbeddingCollectionSharder()),
     ]
 
 

torchrec/distributed/tests/test_sharding_plan.py

@@ -23,6 +23,7 @@
     FusedEmbeddingBagCollectionSharder,
     get_module_to_default_sharders,
     ManagedCollisionEmbeddingBagCollectionSharder,
+    ManagedCollisionEmbeddingCollectionSharder,
     ParameterShardingGenerator,
     QuantEmbeddingBagCollectionSharder,
     QuantEmbeddingCollectionSharder,
@@ -52,7 +53,10 @@
 )
 from torchrec.modules.fp_embedding_modules import FeatureProcessedEmbeddingBagCollection
 from torchrec.modules.fused_embedding_modules import FusedEmbeddingBagCollection
-from torchrec.modules.mc_embedding_modules import ManagedCollisionEmbeddingBagCollection
+from torchrec.modules.mc_embedding_modules import (
+    ManagedCollisionEmbeddingBagCollection,
+    ManagedCollisionEmbeddingCollection,
+)
 from torchrec.quant.embedding_modules import (
     EmbeddingBagCollection as QuantEmbeddingBagCollection,
     EmbeddingCollection as QuantEmbeddingCollection,
@@ -685,12 +689,12 @@ def test_str(self) -> None:
         expected = """
 module: ebc
 
- param   | sharding type | compute kernel | ranks 
+ param   | sharding type | compute kernel | ranks
 -------- | ------------- | -------------- | ------
-user_id  | table_wise    | dense          | [0]   
+user_id  | table_wise    | dense          | [0]
 movie_id | row_wise      | dense          | [0, 1]
 
- param   | shard offsets | shard sizes |   placement  
+ param   | shard offsets | shard sizes |   placement
 -------- | ------------- | ----------- | -------------
 user_id  | [0, 0]        | [4096, 32]  | rank:0/cuda:0
 movie_id | [0, 0]        | [2048, 32]  | rank:0/cuda:0
@@ -710,6 +714,7 @@ def test_module_to_default_sharders(self) -> None:
                 QuantEmbeddingBagCollection,
                 QuantEmbeddingCollection,
                 ManagedCollisionEmbeddingBagCollection,
+                ManagedCollisionEmbeddingCollection,
             ],
         )
         self.assertIsInstance(
@@ -738,3 +743,8 @@ def test_module_to_default_sharders(self) -> None:
             default_sharder_map[ManagedCollisionEmbeddingBagCollection],
             ManagedCollisionEmbeddingBagCollectionSharder,
         )
+
+        self.assertIsInstance(
+            default_sharder_map[ManagedCollisionEmbeddingCollection],
+            ManagedCollisionEmbeddingCollectionSharder,
+        )