Report dry run eviction threshold metric (#4709)

Summary:
Pull Request resolved: #4709

X-link: facebookresearch/FBGEMM#1732

As title

Reviewed By: emlin

Differential Revision: D80198960

fbshipit-source-id: 58db1cb657d92d27eaa1d0e6d6edc93054d1d5d6

Author: EddyLXJ

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -1018,6 +1018,9 @@ 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"
                 )
@@ -3820,12 +3823,14 @@ def _report_eviction_stats(self) -> None:
         T = len(set(self.feature_table_map))
         evicted_counts = torch.zeros(T, dtype=torch.int64)
         processed_counts = torch.zeros(T, dtype=torch.int64)
+        eviction_threshold_with_dry_run = torch.zeros(T, dtype=torch.float)
         full_duration_ms = torch.tensor(0, dtype=torch.int64)
         exec_duration_ms = torch.tensor(0, dtype=torch.int64)
         dry_run_exec_duration_ms = torch.tensor(0, dtype=torch.int64)
         self.ssd_db.get_feature_evict_metric(
             evicted_counts,
             processed_counts,
+            eviction_threshold_with_dry_run,
             full_duration_ms,
             exec_duration_ms,
             dry_run_exec_duration_ms,
@@ -3860,6 +3865,11 @@ def _report_eviction_stats(self) -> None:
                 event_name=f"eviction.feature_table.{t}.processed_counts",
                 data_bytes=int(processed_counts[t].item()),
             )
+            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_wrapper.h

@@ -146,6 +146,7 @@ class DramKVEmbeddingCacheWrapper : public torch::jit::CustomClassHolder {
   void get_feature_evict_metric(
       at::Tensor evicted_counts,
       at::Tensor processed_counts,
+      at::Tensor eviction_threshold_with_dry_run,
       at::Tensor full_duration_ms,
       at::Tensor exec_duration_ms,
       at::Tensor dry_run_exec_duration_ms) {
@@ -155,6 +156,10 @@ class DramKVEmbeddingCacheWrapper : public torch::jit::CustomClassHolder {
           metrics.value().evicted_counts); // evicted_counts (Long)
       processed_counts.copy_(
           metrics.value().processed_counts); // processed_counts (Long)
+      eviction_threshold_with_dry_run.copy_(
+          metrics.value()
+              .eviction_threshold_with_dry_run); // eviction threshold with dry
+                                                 // run (float)
       full_duration_ms.copy_(
           metrics.value().full_duration_ms); // full duration (Long)
       exec_duration_ms.copy_(

fbgemm_gpu/src/dram_kv_embedding_cache/feature_evict.h

@@ -265,6 +265,7 @@ struct FeatureEvictMetrics {
   explicit FeatureEvictMetrics(int table_num) {
     evicted_counts.resize(table_num, 0);
     processed_counts.resize(table_num, 0);
+    eviction_threshold_with_dry_run.resize(table_num, 0.0);
     exec_duration_ms = 0;
     full_duration_ms = 0;
     dry_run_exec_duration_ms = 0;
@@ -273,6 +274,10 @@ struct FeatureEvictMetrics {
   void reset() {
     std::fill(evicted_counts.begin(), evicted_counts.end(), 0);
     std::fill(processed_counts.begin(), processed_counts.end(), 0);
+    std::fill(
+        eviction_threshold_with_dry_run.begin(),
+        eviction_threshold_with_dry_run.end(),
+        0.0);
     exec_duration_ms = 0;
     full_duration_ms = 0;
     dry_run_exec_duration_ms = 0;
@@ -296,6 +301,7 @@ struct FeatureEvictMetrics {
 
   std::vector<int64_t> evicted_counts;
   std::vector<int64_t> processed_counts;
+  std::vector<float> eviction_threshold_with_dry_run;
   int64_t exec_duration_ms;
   int64_t full_duration_ms;
   int64_t dry_run_exec_duration_ms;
@@ -307,6 +313,7 @@ struct FeatureEvictMetricTensors {
   explicit FeatureEvictMetricTensors(int64_t table_num)
       : evicted_counts(at::zeros({table_num}, at::kLong)),
         processed_counts(at::zeros({table_num}, at::kLong)),
+        eviction_threshold_with_dry_run(at::zeros({table_num}, at::kFloat)),
         exec_duration_ms(at::scalar_tensor(0, at::kLong)),
         dry_run_exec_duration_ms(at::scalar_tensor(0, at::kLong)),
         full_duration_ms(at::scalar_tensor(0, at::kLong)) {}
@@ -315,11 +322,14 @@ struct FeatureEvictMetricTensors {
   FeatureEvictMetricTensors(
       at::Tensor evicted,
       at::Tensor processed,
+      at::Tensor eviction_threshold_with_dry_run,
       at::Tensor exec_duration,
       at::Tensor dry_run_exec_duration_ms,
       at::Tensor full_duration)
       : evicted_counts(std::move(evicted)),
         processed_counts(std::move(processed)),
+        eviction_threshold_with_dry_run(
+            std::move(eviction_threshold_with_dry_run)),
         exec_duration_ms(std::move(exec_duration)),
         dry_run_exec_duration_ms(std::move(dry_run_exec_duration_ms)),
         full_duration_ms(std::move(full_duration)) {}
@@ -328,6 +338,7 @@ struct FeatureEvictMetricTensors {
     return FeatureEvictMetricTensors{
         evicted_counts.clone(),
         processed_counts.clone(),
+        eviction_threshold_with_dry_run.clone(),
         exec_duration_ms.clone(),
         dry_run_exec_duration_ms.clone(),
         full_duration_ms.clone()};
@@ -337,6 +348,8 @@ struct FeatureEvictMetricTensors {
   at::Tensor evicted_counts;
   // feature count before evict
   at::Tensor processed_counts;
+  // feature evict threshold with dry run
+  at::Tensor eviction_threshold_with_dry_run;
   // feature evict exec duration
   at::Tensor exec_duration_ms;
   // feature evict dry run exec duration
@@ -893,6 +906,14 @@ class FeatureEvict {
             at::kLong)
             .clone();
 
+    metric_tensors_.eviction_threshold_with_dry_run =
+        at::from_blob(
+            const_cast<float*>(metrics_.eviction_threshold_with_dry_run.data()),
+            {static_cast<int64_t>(
+                metrics_.eviction_threshold_with_dry_run.size())},
+            at::kFloat)
+            .clone();
+
     metric_tensors_.full_duration_ms =
         at::scalar_tensor(metrics_.full_duration_ms, at::kLong);
     metric_tensors_.exec_duration_ms =
@@ -913,14 +934,16 @@ class FeatureEvict {
         "  - dryrun Time taken: {}ms\n"
         "  - Total blocks processed: [{}]\n"
         "  - Blocks evicted: [{}]\n"
-        "  - Eviction rate: [{}]%\n",
+        "  - Eviction rate: [{}]%\n"
+        "  - Eviction threshold dry run: [{}]\n",
         metrics_.full_duration_ms,
         metrics_.exec_duration_ms,
         metrics_.exec_duration_ms * 100.0f / metrics_.full_duration_ms,
         metrics_.dry_run_exec_duration_ms,
         fmt::join(metrics_.processed_counts, ", "),
         fmt::join(metrics_.evicted_counts, ", "),
-        fmt::join(evict_rates, ", "));
+        fmt::join(evict_rates, ", "),
+        fmt::join(metrics_.eviction_threshold_with_dry_run, ", "));
   }
 
   // Thread pool.
@@ -1320,6 +1343,11 @@ class FeatureScoreBasedEvict : public FeatureEvict<weight_type> {
               std::chrono::high_resolution_clock::now().time_since_epoch())
               .count();
     }
+
+    for (int table_id = 0; table_id < num_tables_; ++table_id) {
+      this->metrics_.eviction_threshold_with_dry_run[table_id] =
+          thresholds_[table_id];
+    }
   }
 
  private:

fbgemm_gpu/test/tbe/ssd/kv_backend_test.py

@@ -821,6 +821,7 @@ def test_dram_kv_eviction(self) -> None:
 
         evicted_counts = torch.zeros(T, dtype=torch.int64)
         processed_counts = torch.zeros(T, dtype=torch.int64)
+        eviction_threshold_with_dry_run = torch.zeros(T, dtype=torch.float)
         full_duration_ms = torch.ones(1, dtype=torch.int64) * -1
         exec_duration_ms = torch.empty(1, dtype=torch.int64)
         dry_run_exec_duration_ms = torch.empty(1, dtype=torch.int64)
@@ -831,6 +832,7 @@ def test_dram_kv_eviction(self) -> None:
         dram_kv_backend.get_feature_evict_metric(  # pyre-ignore
             evicted_counts,
             processed_counts,
+            eviction_threshold_with_dry_run,
             full_duration_ms,
             exec_duration_ms,
             dry_run_exec_duration_ms,
@@ -840,10 +842,11 @@ def test_dram_kv_eviction(self) -> None:
             dram_kv_backend.set(indices, weights, count)
             time.sleep(0.01)  # 20ms, stimulate training forward time
             dram_kv_backend.set_cuda(indices, weights, count, 1, True)  # pyre-ignore
-            time.sleep(0.01)  # 20ms, stimulate training backward time
+            time.sleep(1)  # 20ms, stimulate training backward time
             dram_kv_backend.get_feature_evict_metric(  # pyre-ignore
                 evicted_counts,
                 processed_counts,
+                eviction_threshold_with_dry_run,
                 full_duration_ms,
                 exec_duration_ms,
                 dry_run_exec_duration_ms,
@@ -876,6 +879,7 @@ def test_dram_kv_eviction(self) -> None:
             dram_kv_backend.get_feature_evict_metric(
                 evicted_counts,
                 processed_counts,
+                eviction_threshold_with_dry_run,
                 full_duration_ms,
                 exec_duration_ms,
                 dry_run_exec_duration_ms,
@@ -946,6 +950,7 @@ def test_dram_kv_feature_score_eviction(self) -> None:
         full_duration_ms = torch.ones(1, dtype=torch.int64) * -1
         exec_duration_ms = torch.empty(1, dtype=torch.int64)
         dry_run_exec_duration_ms = torch.empty(1, dtype=torch.int64)
+        eviction_threshold_with_dry_run = torch.zeros(T, dtype=torch.float)
 
         shard_load = E / 4
         # init
@@ -969,6 +974,7 @@ def test_dram_kv_feature_score_eviction(self) -> None:
             dram_kv_backend.get_feature_evict_metric(  # pyre-ignore
                 evicted_counts,
                 processed_counts,
+                eviction_threshold_with_dry_run,
                 full_duration_ms,
                 exec_duration_ms,
                 dry_run_exec_duration_ms,
@@ -978,6 +984,7 @@ def test_dram_kv_feature_score_eviction(self) -> None:
         self.assertTrue(full_duration_ms.item() > 0)
         self.assertTrue(exec_duration_ms.item() >= 0)
         self.assertTrue(dry_run_exec_duration_ms.item() > 0)
+        self.assertTrue(all(eviction_threshold_with_dry_run > 0))
 
     @given(
         T=st.integers(min_value=2, max_value=10),