Adding new ID_COUNT eviction trigger condition (pytorch#4829)

Summary:

X-link: facebookresearch/FBGEMM#1855

Differential Revision: D81151216

Author: EddyLXJ

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py

@@ -62,7 +62,7 @@ def from_str(cls, key: str):
 
 class EvictionPolicy(NamedTuple):
     eviction_trigger_mode: int = (
-        0  # disabled, 0: disabled, 1: iteration, 2: mem_util, 3: manual
+        0  # disabled, 0: disabled, 1: iteration, 2: mem_util, 3: manual 4: id count
     )
     eviction_strategy: int = (
         0  # 0: timestamp, 1: counter , 2: counter + timestamp, 3: feature l2 norm 4: timestamp threshold 5: feature score
@@ -85,11 +85,11 @@ class EvictionPolicy(NamedTuple):
     feature_score_counter_decay_rates: Optional[List[float]] = (
         None  # feature_score_counter_decay_rates for each table if eviction strategy is feature score
     )
-    max_training_id_num_per_table: Optional[List[int]] = (
-        None  # max_training_id_num_per_table for each table
+    training_id_eviction_trigger_count: Optional[List[int]] = (
+        None  # training_id_eviction_trigger_count for each table
     )
-    target_eviction_percent_per_table: Optional[List[float]] = (
-        None  # target_eviction_percent_per_table for each table
+    training_id_keep_count: Optional[List[int]] = (
+        None  # training_id_keep_count for each table
     )
     l2_weight_thresholds: Optional[List[float]] = (
         None  # l2_weight_thresholds for each table if eviction strategy is feature l2 norm
@@ -116,8 +116,8 @@ class EvictionPolicy(NamedTuple):
     meta_header_lens: Optional[List[int]] = None  # metaheader length for each table
 
     def validate(self) -> None:
-        assert self.eviction_trigger_mode in [0, 1, 2, 3], (
-            "eviction_trigger_mode must be 0, 1, 2, or 3, "
+        assert self.eviction_trigger_mode in [0, 1, 2, 3, 4], (
+            "eviction_trigger_mode must be 0, 1, 2, 3 or 4 "
             f"actual {self.eviction_trigger_mode}"
         )
         if self.eviction_trigger_mode == 0:
@@ -139,6 +139,10 @@ def validate(self) -> None:
             assert (
                 self.eviction_mem_threshold_gb is not None
             ), "eviction_mem_threshold_gb must be set if eviction_trigger_mode is 2"
+        elif self.eviction_trigger_mode == 4:
+            assert (
+                self.training_id_eviction_trigger_count is not None
+            ), "training_id_eviction_trigger_count must be set if eviction_trigger_mode is 4"
 
         if self.eviction_strategy == 0:
             assert self.ttls_in_mins is not None, (
@@ -184,13 +188,13 @@ def validate(self) -> None:
                 "feature_score_counter_decay_rates must be set if eviction_strategy is 5, "
                 f"actual {self.feature_score_counter_decay_rates}"
             )
-            assert self.max_training_id_num_per_table is not None, (
-                "max_training_id_num_per_table must be set if eviction_strategy is 5,"
-                f"actual {self.max_training_id_num_per_table}"
+            assert self.training_id_eviction_trigger_count is not None, (
+                "training_id_eviction_trigger_count must be set if eviction_strategy is 5,"
+                f"actual {self.training_id_eviction_trigger_count}"
             )
-            assert self.target_eviction_percent_per_table is not None, (
-                "target_eviction_percent_per_table must be set if eviction_strategy is 5,"
-                f"actual {self.target_eviction_percent_per_table}"
+            assert self.training_id_keep_count is not None, (
+                "training_id_keep_count must be set if eviction_strategy is 5,"
+                f"actual {self.training_id_keep_count}"
             )
             assert self.threshold_calculation_bucket_stride is not None, (
                 "threshold_calculation_bucket_stride must be set if eviction_strategy is 5,"
@@ -201,12 +205,12 @@ def validate(self) -> None:
                 f"actual {self.threshold_calculation_bucket_num}"
             )
             assert (
-                len(self.target_eviction_percent_per_table)
+                len(self.training_id_keep_count)
                 == len(self.feature_score_counter_decay_rates)
-                == len(self.max_training_id_num_per_table)
+                == len(self.training_id_eviction_trigger_count)
             ), (
-                "feature_score_thresholds, max_training_id_num_per_table and target_eviction_percent_per_table must have the same length, "
-                f"actual {self.target_eviction_percent_per_table} vs {self.feature_score_counter_decay_rates} vs {self.max_training_id_num_per_table}"
+                "feature_score_thresholds, training_id_eviction_trigger_count and training_id_keep_count must have the same length, "
+                f"actual {self.training_id_keep_count} vs {self.feature_score_counter_decay_rates} vs {self.training_id_eviction_trigger_count}"
             )
 
 

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -690,16 +690,16 @@ def __init__(
                 )
                 # Please refer to https://fburl.com/gdoc/nuupjwqq for the following eviction parameters.
                 eviction_config = torch.classes.fbgemm.FeatureEvictConfig(
-                    self.kv_zch_params.eviction_policy.eviction_trigger_mode,  # eviction is disabled, 0: disabled, 1: iteration, 2: mem_util, 3: manual
+                    self.kv_zch_params.eviction_policy.eviction_trigger_mode,  # eviction is disabled, 0: disabled, 1: iteration, 2: mem_util, 3: manual, 4: id count
                     self.kv_zch_params.eviction_policy.eviction_strategy,  # evict_trigger_strategy: 0: timestamp, 1: counter, 2: counter + timestamp, 3: feature l2 norm, 4: timestamp threshold 5: feature score
                     self.kv_zch_params.eviction_policy.eviction_step_intervals,  # trigger_step_interval if trigger mode is iteration
                     eviction_mem_threshold_gb,  # mem_util_threshold_in_GB if trigger mode is mem_util
                     self.kv_zch_params.eviction_policy.ttls_in_mins,  # ttls_in_mins for each table if eviction strategy is timestamp
                     self.kv_zch_params.eviction_policy.counter_thresholds,  # counter_thresholds for each table if eviction strategy is counter
                     self.kv_zch_params.eviction_policy.counter_decay_rates,  # counter_decay_rates for each table if eviction strategy is counter
                     self.kv_zch_params.eviction_policy.feature_score_counter_decay_rates,  # feature_score_counter_decay_rates for each table if eviction strategy is feature score
-                    self.kv_zch_params.eviction_policy.max_training_id_num_per_table,  # max_training_id_num for each table
-                    self.kv_zch_params.eviction_policy.target_eviction_percent_per_table,  # target_eviction_percent for each table
+                    self.kv_zch_params.eviction_policy.training_id_eviction_trigger_count,  # training_id_eviction_trigger_count for each table
+                    self.kv_zch_params.eviction_policy.training_id_keep_count,  # training_id_keep_count for each table
                     self.kv_zch_params.eviction_policy.l2_weight_thresholds,  # l2_weight_thresholds for each table if eviction strategy is feature l2 norm
                     table_dims.tolist() if table_dims is not None else None,
                     self.kv_zch_params.eviction_policy.threshold_calculation_bucket_stride,  # threshold_calculation_bucket_stride if eviction strategy is feature score
@@ -1047,9 +1047,6 @@ def __init__(
                 self.stats_reporter.register_stats(
                     f"eviction.feature_table.{t}.processed_counts"
                 )
-                self.stats_reporter.register_stats(
-                    f"eviction.feature_table.{t}.eviction_threshold_with_dry_run"
-                )
                 self.stats_reporter.register_stats(
                     f"eviction.feature_table.{t}.evict_rate"
                 )
@@ -3926,11 +3923,6 @@ def _report_eviction_stats(self) -> None:
                 data_bytes=int(processed_counts[t].item()),
                 enable_tb_metrics=True,
             )
-            stats_reporter.report_data_amount(
-                iteration_step=self.step,
-                event_name=f"eviction.feature_table.{t}.eviction_threshold_with_dry_run",
-                data_bytes=float(eviction_threshold_with_dry_run[t].item()),
-            )
             if processed_counts[t].item() != 0:
                 stats_reporter.report_data_amount(
                     iteration_step=self.step,

fbgemm_gpu/src/dram_kv_embedding_cache/SynchronizedShardedMap.h

@@ -89,6 +89,19 @@ class SynchronizedShardedMap {
     return num_rows;
   }
 
+  auto getActualUsedIDCount() const {
+    size_t used_id_count = 0;
+    for (size_t i = 0; i < shards_.size(); ++i) {
+      int64_t mempool_idx = i % mempools_.size();
+      // only calculate the sizes of K, V and block that are used
+      if (mempools_[mempool_idx]->get_allocated_chunk_bytes() > 0) {
+        auto rlmap = shards_[i].rlock();
+        used_id_count += rlmap->size();
+      }
+    }
+    return used_id_count;
+  }
+
  private:
   std::vector<folly::Synchronized<folly::F14FastMap<K, V>, M>> shards_;
   std::vector<std::unique_ptr<FixedBlockPool>> mempools_;

fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h

@@ -1101,6 +1101,19 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
         }
         break;
       }
+      case EvictTriggerMode::ID_COUNT: {
+        auto used_id_count = get_map_actual_used_id_count();
+        auto eviction_threshold = 0;
+        for (const auto& eviction_trigger_count :
+             feature_evict_config_.value()
+                 ->training_id_eviction_trigger_count_.value()) {
+          eviction_threshold += eviction_trigger_count;
+        }
+        if (used_id_count > eviction_threshold) {
+          trigger_feature_evict();
+        }
+        break;
+      }
       default:
         break;
     }
@@ -1125,6 +1138,10 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
     return kv_store_.getNumRows();
   }
 
+  size_t get_map_actual_used_id_count() const {
+    return kv_store_.getActualUsedIDCount();
+  }
+
   void resume_ongoing_eviction(bool force_resume = false) override {
     if (!is_training_ && !force_resume) {
       return;

fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_inference_wrapper.cpp

@@ -56,8 +56,8 @@ void DramKVEmbeddingInferenceWrapper::init(
           std::nullopt /* counter_thresholds */,
           std::nullopt /* counter_decay_rates */,
           std::nullopt /* feature_score_counter_decay_rates */,
-          std::nullopt /* max_training_id_num_per_table */,
-          std::nullopt /* target_eviction_percent_per_table */,
+          std::nullopt /* training_id_eviction_trigger_count */,
+          std::nullopt /* training_id_keep_count */,
           std::nullopt /* l2_weight_thresholds */,
           std::nullopt /* embedding_dims */,
           std::nullopt /* threshold_calculation_bucket_stride */,