Using eviction policy to populate _enable_feature_score_weight_accumulation (#3354)

Summary:
Pull Request resolved: #3354

As title, if kvzch table is enable feature score based eviction, we need to populate _enable_feature_score_weight_accumulation to true to collect feature score.

Reviewed By: emlin

Differential Revision: D81744008

fbshipit-source-id: 79b0b2e5dfb101387545e7d24c04ee133046431e

Author: EddyLXJ

torchrec/distributed/embedding.py

@@ -47,7 +47,6 @@
     ShardingType,
 )
 from torchrec.distributed.fused_params import (
-    ENABLE_FEATURE_SCORE_WEIGHT_ACCUMULATION,
     FUSED_PARAM_IS_SSD_TABLE,
     FUSED_PARAM_SSD_TABLE_LIST,
 )
@@ -91,6 +90,7 @@
 from torchrec.modules.embedding_configs import (
     EmbeddingConfig,
     EmbeddingTableConfig,
+    FeatureScoreBasedEvictionPolicy,
     PoolingType,
 )
 from torchrec.modules.embedding_modules import (
@@ -422,18 +422,6 @@ def __init__(
         super().__init__(qcomm_codecs_registry=qcomm_codecs_registry)
         self._enable_feature_score_weight_accumulation: bool = False
 
-        if (
-            fused_params is not None
-            and ENABLE_FEATURE_SCORE_WEIGHT_ACCUMULATION in fused_params
-        ):
-            self._enable_feature_score_weight_accumulation = cast(
-                bool, fused_params[ENABLE_FEATURE_SCORE_WEIGHT_ACCUMULATION]
-            )
-            fused_params.pop(ENABLE_FEATURE_SCORE_WEIGHT_ACCUMULATION)
-            logger.info(
-                f"EC feature score weight accumulation enabled: {self._enable_feature_score_weight_accumulation}."
-            )
-
         self._module_fqn = module_fqn
         self._embedding_configs: List[EmbeddingConfig] = module.embedding_configs()
         self._table_names: List[str] = [
@@ -492,6 +480,18 @@ def __init__(
         self._has_uninitialized_input_dist: bool = True
         logger.info(f"EC index dedup enabled: {self._use_index_dedup}.")
 
+        for config in self._embedding_configs:
+            virtual_table_eviction_policy = config.virtual_table_eviction_policy
+            if virtual_table_eviction_policy is not None and isinstance(
+                virtual_table_eviction_policy, FeatureScoreBasedEvictionPolicy
+            ):
+                self._enable_feature_score_weight_accumulation = True
+                break
+
+        logger.info(
+            f"EC feature score weight accumulation enabled: {self._enable_feature_score_weight_accumulation}."
+        )
+
         # Get all fused optimizers and combine them.
         optims = []
         for lookup in self._lookups:

torchrec/distributed/fused_params.py

@@ -33,9 +33,6 @@
 FUSED_PARAM_SSD_TABLE_LIST: str = "__register_ssd_table_list"
 # Bool fused param per table to check if the table is offloaded to SSD
 FUSED_PARAM_IS_SSD_TABLE: str = "__register_is_ssd_table"
-ENABLE_FEATURE_SCORE_WEIGHT_ACCUMULATION: str = (
-    "enable_feature_score_weight_accumulation"
-)
 
 
 class TBEToRegisterMixIn: