Split FPEBC (pytorch#2535)

Summary:
Pull Request resolved: pytorch#2535

Support API to split a FPEBC into the FP and the EBC counterparts. Furthermore, if FP is a nn.ModuleDict, wrap it over a nn.Module where the forward pass can be called. This helps support eager model split in IEN.

Reviewed By: iamzainhuda

Differential Revision: D65439160

fbshipit-source-id: 42cc148daefca22fce5d67b7adf10834732e6803

Author: PaulZhang12

torchrec/modules/fp_embedding_modules.py

@@ -7,7 +7,7 @@
 
 # pyre-strict
 
-from typing import Dict, List, Set, Union
+from typing import Dict, List, Set, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -55,6 +55,15 @@ def apply_feature_processors_to_kjt(
     )
 
 
+class FeatureProcessorDictWrapper(FeatureProcessorsCollection):
+    def __init__(self, feature_processors: nn.ModuleDict) -> None:
+        super().__init__()
+        self._feature_processors = feature_processors
+
+    def forward(self, features: KeyedJaggedTensor) -> KeyedJaggedTensor:
+        return apply_feature_processors_to_kjt(features, self._feature_processors)
+
+
 class FeatureProcessedEmbeddingBagCollection(nn.Module):
     """
     FeatureProcessedEmbeddingBagCollection represents a EmbeddingBagCollection module and a set of feature processor modules.
@@ -125,6 +134,18 @@ def __init__(
             feature_names_set.update(table_config.feature_names)
         self._feature_names: List[str] = list(feature_names_set)
 
+    def split(
+        self,
+    ) -> Tuple[FeatureProcessorsCollection, EmbeddingBagCollection]:
+        if isinstance(self._feature_processors, nn.ModuleDict):
+            return (
+                FeatureProcessorDictWrapper(self._feature_processors),
+                self._embedding_bag_collection,
+            )
+        else:
+            assert isinstance(self._feature_processors, FeatureProcessorsCollection)
+            return self._feature_processors, self._embedding_bag_collection
+
     def forward(
         self,
         features: KeyedJaggedTensor,

torchrec/modules/tests/test_fp_embedding_modules.py

@@ -95,6 +95,15 @@ def test_position_weighted_module_ebc_with_excessive_features(self) -> None:
         self.assertEqual(pooled_embeddings.values().size(), (3, 16))
         self.assertEqual(pooled_embeddings.offset_per_key(), [0, 8, 16])
 
+        # Test split method, FP then EBC
+        fp, ebc = fp_ebc.split()
+        fp_kjt = fp(features)
+        pooled_embeddings_split = ebc(fp_kjt)
+
+        self.assertEqual(pooled_embeddings_split.keys(), ["f1", "f2"])
+        self.assertEqual(pooled_embeddings_split.values().size(), (3, 16))
+        self.assertEqual(pooled_embeddings_split.offset_per_key(), [0, 8, 16])
+
 
 class PositionWeightedModuleCollectionEmbeddingBagCollectionTest(unittest.TestCase):
     def generate_fp_ebc(self) -> FeatureProcessedEmbeddingBagCollection:
@@ -144,3 +153,12 @@ def test_position_weighted_collection_module_ebc(self) -> None:
             pooled_embeddings_gm_script.offset_per_key(),
             pooled_embeddings.offset_per_key(),
         )
+
+        # Test split method, FP then EBC
+        fp, ebc = fp_ebc.split()
+        fp_kjt = fp(features)
+        pooled_embeddings_split = ebc(fp_kjt)
+
+        self.assertEqual(pooled_embeddings_split.keys(), ["f1", "f2"])
+        self.assertEqual(pooled_embeddings_split.values().size(), (3, 16))
+        self.assertEqual(pooled_embeddings_split.offset_per_key(), [0, 8, 16])