generalize trec feature processor (pytorch#443)

Summary: Pull Request resolved: pytorch#443 torchrec feature processor oss change to support D37165122 Reviewed By: xing-liu, RenfeiChen-FB Differential Revision: D35735337 fbshipit-source-id: 0bf93776b0f3994962ae325b6105002284a589fd
classicvalues · Jun 21, 2022 · ca52aa0 · ca52aa0
1 parent 6da4b4c
commit ca52aa0
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/torchrec/distributed/embedding_lookup.py b/torchrec/distributed/embedding_lookup.py
@@ -288,9 +288,7 @@ def forward(
                 if (
                     config.has_feature_processor
                     and self._feature_processor is not None
-                    and isinstance(
-                        self._feature_processor, GroupedPositionWeightedModule
-                    )
+                    and isinstance(self._feature_processor, BaseGroupedFeatureProcessor)
                 ):
                     features = self._feature_processor(features)
                 embeddings.append(emb_op(features))
@@ -301,9 +299,17 @@ def forward(
                     self._id_score_list_feature_splits,
                 )
             )
-            for emb_op, features in zip(
-                self._score_emb_modules, id_score_list_features_by_group
+            for config, emb_op, features in zip(
+                self.grouped_score_configs,
+                self._score_emb_modules,
+                id_score_list_features_by_group,
             ):
+                if (
+                    config.has_feature_processor
+                    and self._feature_processor is not None
+                    and isinstance(self._feature_processor, BaseGroupedFeatureProcessor)
+                ):
+                    features = self._feature_processor(features)
                 embeddings.append(emb_op(features))
 
         if len(embeddings) == 0:
@@ -563,9 +569,7 @@ def forward(
                 if (
                     config.has_feature_processor
                     and self._feature_processor is not None
-                    and isinstance(
-                        self._feature_processor, GroupedPositionWeightedModule
-                    )
+                    and isinstance(self._feature_processor, BaseGroupedFeatureProcessor)
                 ):
                     features = self._feature_processor(features)
                 embeddings.append(emb_op(features))

diff --git a/torchrec/distributed/embedding_sharding.py b/torchrec/distributed/embedding_sharding.py
@@ -419,7 +419,7 @@ def _group_tables_per_rank(
         for data_type in DataType:
             for pooling in PoolingType:
                 for is_weighted in [True, False]:
-                    for has_feature_processor in [True, False]:
+                    for has_feature_processor in [False, True]:
                         for compute_kernel in [
                             EmbeddingComputeKernel.DENSE,
                             EmbeddingComputeKernel.FUSED,