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

Summary:
X-link: pytorch/FBGEMM#5059


X-link: facebookresearch/FBGEMM#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.

Differential Revision: D84660528

Author: EddyLXJ

torchrec/distributed/batched_embedding_kernel.py

@@ -334,6 +334,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"]
@@ -393,6 +396,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:
@@ -454,6 +460,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

@@ -240,6 +240,9 @@ class FeatureScoreBasedEvictionPolicy(VirtualTableEvictionPolicy):
         None  # 0 means no eviction
     )
     inference_eviction_ttl_mins: Optional[int] = None  # 0 means no eviction
+    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: