Support no eviction in Feature score eviction policy (meta-pytorch#3488)

Summary:
X-link: pytorch/FBGEMM#5059

Pull Request resolved: meta-pytorch#3488

X-link: https://github.com/facebookresearch/FBGEMM/pull/2068

As title
If one table is using feature score eviction in one tbe, then all tables in this tbe need to use the same policy. Feature score eviction can support ttl based eviction now. This diff is adding support no eviction in feature score eviction policy.

Reviewed By: emlin

Differential Revision: D84660528

fbshipit-source-id: 28d85126da30b2c357c7f99b6535be381a76f2ce

Author: EddyLXJ

torchrec/distributed/batched_embedding_kernel.py

@@ -354,6 +354,9 @@ def _populate_zero_collision_tbe_params(
         training_id_eviction_trigger_count = [0] * len(config.embedding_tables)
         training_id_keep_count = [0] * len(config.embedding_tables)
         l2_weight_thresholds = [0.0] * len(config.embedding_tables)
+        enable_eviction_for_feature_score_eviction_policy = [True] * len(
+            config.embedding_tables
+        )
         eviction_strategy = -1
         table_names = [table.name for table in config.embedding_tables]
         l2_cache_size = tbe_params["l2_cache_size"]
@@ -398,6 +401,9 @@ def _populate_zero_collision_tbe_params(
                     )
                     training_id_keep_count[i] = policy_t.training_id_keep_count
                     ttls_in_mins[i] = policy_t.eviction_ttl_mins
+                    enable_eviction_for_feature_score_eviction_policy[i] = (
+                        policy_t.enable_eviction
+                    )
                     if eviction_strategy == -1 or eviction_strategy == 5:
                         eviction_strategy = 5
                     else:
@@ -459,6 +465,7 @@ def _populate_zero_collision_tbe_params(
             eviction_free_mem_check_interval_batch=eviction_free_mem_check_interval_batch,
             threshold_calculation_bucket_stride=threshold_calculation_bucket_stride,
             threshold_calculation_bucket_num=threshold_calculation_bucket_num,
+            enable_eviction_for_feature_score_eviction_policy=enable_eviction_for_feature_score_eviction_policy,
         )
     else:
         eviction_policy = EvictionPolicy(meta_header_lens=meta_header_lens)

torchrec/modules/embedding_configs.py

@@ -243,6 +243,9 @@ class FeatureScoreBasedEvictionPolicy(VirtualTableEvictionPolicy):
     feature_score_mapping: Optional[Dict[str, float]] = None  # feature score mapping
     feature_score_default_value: Optional[float] = None  # default feature score value
     enable_auto_feature_score_collection: bool = False
+    enable_eviction: bool = (
+        True  # Currently we only support using same eviction policy in one tbe for mutiple tables, if one table doesn't need eviction we can set this flag to False
+    )
 
     def __post_init__(self) -> None:
         if self.inference_eviction_feature_score_threshold is None: