Adding new ID_COUNT eviction trigger condition (#4829)

Summary:
X-link: meta-pytorch/torchrec#3357

Pull Request resolved: #4829

X-link: facebookresearch/FBGEMM#1855

ID COUNT Trigger is based on total num of id in all kvzch tables. Memory usage is cross all tables but eviction trigger is for tbe, like:
if target total memory is 10G tbe1 has table [table0, table1], training_id_eviction_trigger_count is [200, 300] tbe2 has table [table2], training_id_eviction_trigger_count is [500] then the target memory for table0 is 2G, table1 is 3G, table2 is 5G, but the trigger condition is tbe level.  If currently total id num in tbe1 is [100, 350], in this case tbe1 will not trigger evict even num in table1 is greater than 300 because total num. 100 + 350 = 450 < 500. If total id num in tbe1 is [100, 450], it will trigger eviction on tbe1 only for table1, because num in table0 100 < 200.

Reviewed By: emlin

Differential Revision: D81151216

fbshipit-source-id: 693bd47e950cc63498b4cdfc7691636197c832dc

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/dram_kv_embedding_cache.h

@@ -1101,6 +1101,14 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
         }
         break;
       }
+      case EvictTriggerMode::ID_COUNT: {
+        auto used_id_count = get_num_rows();
+        if (used_id_count > feature_evict_config_.value()
+                                ->total_id_eviction_trigger_count_.value()) {
+          trigger_feature_evict();
+        }
+        break;
+      }
       default:
         break;
     }

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 */,

fbgemm_gpu/src/dram_kv_embedding_cache/feature_evict.h

@@ -33,7 +33,8 @@ enum class EvictTriggerMode {
   DISABLED, // Do not use feature evict
   ITERATION, // Trigger based on iteration steps
   MEM_UTIL, // Trigger based on memory usage
-  MANUAL // Manually triggered by upstream
+  MANUAL, // Manually triggered by upstream
+  ID_COUNT // Trigger based on id count
 };
 inline std::string to_string(EvictTriggerMode mode) {
   switch (mode) {
@@ -45,6 +46,8 @@ inline std::string to_string(EvictTriggerMode mode) {
       return "MEM_UTIL";
     case EvictTriggerMode::MANUAL:
       return "MANUAL";
+    case EvictTriggerMode::ID_COUNT:
+      return "ID_COUNT";
   }
 }
 
@@ -102,8 +105,8 @@ struct FeatureEvictConfig : public torch::jit::CustomClassHolder {
       std::optional<std::vector<int64_t>> counter_thresholds,
       std::optional<std::vector<double>> counter_decay_rates,
       std::optional<std::vector<double>> feature_score_counter_decay_rates,
-      std::optional<std::vector<int64_t>> max_training_id_num_per_table,
-      std::optional<std::vector<double>> target_eviction_percent_per_table,
+      std::optional<std::vector<int64_t>> training_id_eviction_trigger_count,
+      std::optional<std::vector<int64_t>> training_id_keep_count,
       std::optional<std::vector<double>> l2_weight_thresholds,
       std::optional<std::vector<int64_t>> embedding_dims,
       std::optional<double> threshold_calculation_bucket_stride = 0.2,
@@ -120,10 +123,9 @@ struct FeatureEvictConfig : public torch::jit::CustomClassHolder {
         counter_decay_rates_(counter_decay_rates),
         feature_score_counter_decay_rates_(
             std::move(feature_score_counter_decay_rates)),
-        max_training_id_num_per_table_(
-            std::move(max_training_id_num_per_table)),
-        target_eviction_percent_per_table_(
-            std::move(target_eviction_percent_per_table)),
+        training_id_eviction_trigger_count_(
+            std::move(training_id_eviction_trigger_count)),
+        training_id_keep_count_(std::move(training_id_keep_count)),
         l2_weight_thresholds_(l2_weight_thresholds),
         embedding_dims_(embedding_dims),
         threshold_calculation_bucket_stride_(
@@ -160,6 +162,28 @@ struct FeatureEvictConfig : public torch::jit::CustomClassHolder {
       case EvictTriggerMode::MANUAL: {
         break;
       }
+      case EvictTriggerMode::ID_COUNT: {
+        CHECK(
+            training_id_eviction_trigger_count_.has_value() &&
+            !training_id_eviction_trigger_count_.value().empty());
+        const auto& vec = training_id_eviction_trigger_count_.value();
+        eviction_trigger_stats_log = ", training_id_eviction_trigger_count: [";
+        total_id_eviction_trigger_count_ = 0;
+        for (size_t i = 0; i < vec.size(); ++i) {
+          total_id_eviction_trigger_count_ =
+              total_id_eviction_trigger_count_.value() + vec[i];
+          if (vec[i] <= 0) {
+            throw std::runtime_error(
+                "Invalid training_id_eviction_trigger_count, must be positive if ID_COUNT trigger mode is used");
+          }
+          eviction_trigger_stats_log += std::to_string(vec[i]);
+          if (i + 1 < vec.size()) {
+            eviction_trigger_stats_log += ", ";
+          }
+        }
+        eviction_trigger_stats_log += "]";
+        break;
+      }
       default:
         throw std::runtime_error("Unknown evict trigger mode");
     }
@@ -178,18 +202,21 @@ struct FeatureEvictConfig : public torch::jit::CustomClassHolder {
 
       case EvictTriggerStrategy::BY_FEATURE_SCORE: {
         CHECK(feature_score_counter_decay_rates_.has_value());
-        CHECK(max_training_id_num_per_table_.has_value());
-        CHECK(target_eviction_percent_per_table_.has_value());
+        CHECK(training_id_eviction_trigger_count_.has_value());
+        CHECK(training_id_keep_count_.has_value());
+        CHECK(total_id_eviction_trigger_count_.has_value());
         CHECK(threshold_calculation_bucket_stride_.has_value());
         CHECK(threshold_calculation_bucket_num_.has_value());
         CHECK(ttls_in_mins_.has_value());
         LOG(INFO) << "eviction config, trigger mode:"
                   << to_string(trigger_mode_) << eviction_trigger_stats_log
                   << ", strategy: " << to_string(trigger_strategy_)
-                  << ", max_training_id_num_per_table: "
-                  << max_training_id_num_per_table_.value()
-                  << ", target_eviction_percent_per_table:"
-                  << target_eviction_percent_per_table_.value()
+                  << ", training_id_eviction_trigger_count: "
+                  << training_id_eviction_trigger_count_.value()
+                  << ", training_id_keep_count:"
+                  << training_id_keep_count_.value()
+                  << ", total_id_eviction_trigger_count: "
+                  << total_id_eviction_trigger_count_.value()
                   << ", ttls_in_mins: " << ttls_in_mins_.value()
                   << ", threshold_calculation_bucket_stride: "
                   << threshold_calculation_bucket_stride_.value()
@@ -252,8 +279,9 @@ struct FeatureEvictConfig : public torch::jit::CustomClassHolder {
   std::optional<std::vector<int64_t>> counter_thresholds_;
   std::optional<std::vector<double>> counter_decay_rates_;
   std::optional<std::vector<double>> feature_score_counter_decay_rates_;
-  std::optional<std::vector<int64_t>> max_training_id_num_per_table_;
-  std::optional<std::vector<double>> target_eviction_percent_per_table_;
+  std::optional<std::vector<int64_t>> training_id_eviction_trigger_count_;
+  std::optional<std::vector<int64_t>> training_id_keep_count_;
+  std::optional<int64_t> total_id_eviction_trigger_count_;
   std::optional<std::vector<double>> l2_weight_thresholds_;
   std::optional<std::vector<int64_t>> embedding_dims_;
   std::optional<double> threshold_calculation_bucket_stride_;
@@ -953,8 +981,8 @@ class FeatureScoreBasedEvict : public FeatureEvict<weight_type> {
       SynchronizedShardedMap<int64_t, weight_type*>& kv_store,
       const std::vector<int64_t>& sub_table_hash_cumsum,
       const std::vector<double>& decay_rates,
-      const std::vector<int64_t>& max_training_id_num_per_table,
-      const std::vector<double>& target_eviction_percent_per_table,
+      const std::vector<int64_t>& training_id_eviction_trigger_count,
+      const std::vector<int64_t>& training_id_keep_count,
       const std::vector<int64_t>& ttls_in_mins,
       const double threshold_calculation_bucket_stride,
       const int64_t threshold_calculation_bucket_num,
@@ -972,8 +1000,8 @@ class FeatureScoreBasedEvict : public FeatureEvict<weight_type> {
             is_training,
             test_mode),
         decay_rates_(decay_rates),
-        max_training_id_num_per_table_(max_training_id_num_per_table),
-        target_eviction_percent_per_table_(target_eviction_percent_per_table),
+        training_id_eviction_trigger_count_(training_id_eviction_trigger_count),
+        training_id_keep_count_(training_id_keep_count),
         ttls_in_mins_(ttls_in_mins),
         threshold_calculation_bucket_stride_(
             threshold_calculation_bucket_stride),
@@ -1007,6 +1035,13 @@ class FeatureScoreBasedEvict : public FeatureEvict<weight_type> {
       bool add_new_block = false) {
     int64_t key = FixedBlockPool::get_key(block);
     int sub_table_id = this->get_sub_table_id(key);
+
+    if (ttls_in_mins_[sub_table_id] > 0) {
+      // If ttl is enabled, we don't need to populate the feature score
+      // bucket.
+      return;
+    }
+
     if (add_new_block) {
       double ratio = FixedBlockPool::get_feature_score_rate(block);
       int64_t idx = get_bucket_id_from_ratio(ratio);
@@ -1124,11 +1159,7 @@ class FeatureScoreBasedEvict : public FeatureEvict<weight_type> {
             this->local_blocks_num_per_shard_per_table_[table_id][shard_id];
       }
 
-      const double target_keep_ratio =
-          1 - target_eviction_percent_per_table_[table_id];
-      int64_t max_id =
-          static_cast<int64_t>(max_training_id_num_per_table_[table_id]);
-      int64_t keep_count = static_cast<int64_t>(max_id * target_keep_ratio);
+      int64_t keep_count = training_id_keep_count_[table_id];
       int64_t evict_count = total - keep_count;
 
       int64_t acc_count = 0;
@@ -1202,10 +1233,12 @@ class FeatureScoreBasedEvict : public FeatureEvict<weight_type> {
   std::vector<double> thresholds_; // Threshold for eviction.
 
   const std::vector<int64_t>&
-      max_training_id_num_per_table_; // training max id for each table.
-  const std::vector<double>&
-      target_eviction_percent_per_table_; // target eviction percent for
-                                          // each table
+      training_id_eviction_trigger_count_; // training id num for trigger
+                                           // eviction.
+  const std::vector<int64_t>&
+      training_id_keep_count_; // target keep training id num per table after
+                               // eviction.
+
   const std::vector<int64_t>& ttls_in_mins_; // Time-to-live for eviction.
   std::vector<std::vector<std::vector<size_t>>>
       local_buckets_per_shard_per_table_;
@@ -1453,8 +1486,8 @@ std::unique_ptr<FeatureEvict<weight_type>> create_feature_evict(
           kv_store,
           sub_table_hash_cumsum,
           config->feature_score_counter_decay_rates_.value(),
-          config->max_training_id_num_per_table_.value(),
-          config->target_eviction_percent_per_table_.value(),
+          config->training_id_eviction_trigger_count_.value(),
+          config->training_id_keep_count_.value(),
           config->ttls_in_mins_.value(),
           config->threshold_calculation_bucket_stride_.value(),
           config->threshold_calculation_bucket_num_.value(),