add read load metrics (#4956)

Summary:
Pull Request resolved: #4956

Add DRAM metrics to tbe stats report:
- the read load cnts: measure loading amount

Reviewed By: emlin

Differential Revision: D83506472

fbshipit-source-id: eb36b588d0ecfa1da998f855169f5b092ff473ff

Author: Kathy Xu

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -3971,8 +3971,8 @@ def _report_dram_kv_perf_stats(self) -> None:
             self.step, stats_reporter.report_interval  # pyre-ignore
         )
 
-        if len(dram_kv_perf_stats) != 23:
-            logging.error("dram cache perf stats should have 23 elements")
+        if len(dram_kv_perf_stats) != 24:
+            logging.error("dram cache perf stats should have 24 elements")
             return
 
         dram_read_duration = dram_kv_perf_stats[0]
@@ -4001,6 +4001,7 @@ def _report_dram_kv_perf_stats(self) -> None:
         dram_kv_allocated_bytes = dram_kv_perf_stats[20]
         dram_kv_actual_used_chunk_bytes = dram_kv_perf_stats[21]
         dram_kv_num_rows = dram_kv_perf_stats[22]
+        dram_kv_read_counts = dram_kv_perf_stats[23]
 
         stats_reporter.report_duration(
             iteration_step=self.step,
@@ -4142,6 +4143,13 @@ def _report_dram_kv_perf_stats(self) -> None:
             enable_tb_metrics=True,
         )
 
+        stats_reporter.report_data_amount(
+            iteration_step=self.step,
+            event_name="dram_kv.perf.get.dram_kv_read_counts",
+            data_bytes=dram_kv_read_counts,
+            enable_tb_metrics=True,
+        )
+
         stats_reporter.report_data_amount(
             iteration_step=self.step,
             event_name=self.dram_kv_allocated_bytes_stats_name,

fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h

@@ -777,6 +777,10 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
       const at::Tensor& count,
       int64_t width_offset = 0,
       std::optional<int64_t> width_length = std::nullopt) {
+    auto read_count = count.scalar_type() == at::ScalarType::Long
+        ? *(count.data_ptr<int64_t>())
+        : *(count.data_ptr<int32_t>());
+    read_num_counts_ += read_count;
     // assuming get is called once each iteration and only by train
     // iteration(excluding state_dict)
     auto start_ts = facebook::WallClockUtil::NowInUsecFast();
@@ -1330,7 +1334,7 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
   std::vector<double> get_dram_kv_perf(
       const int64_t step,
       const int64_t interval) {
-    std::vector<double> ret(23, 0); // num metrics
+    std::vector<double> ret(24, 0); // num metrics
     if (step > 0 && step % interval == 0) {
       int reset_val = 0;
 
@@ -1377,6 +1381,8 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
       auto dram_bwd_l1_cnflct_miss_write_missing_load_ =
           bwd_l1_cnflct_miss_write_missing_load_avg_.exchange(reset_val);
 
+      auto read_num_counts = read_num_counts_.exchange(reset_val);
+
       ret[0] = dram_read_total_duration / interval;
       ret[1] = dram_read_sharding_total_duration / interval;
       ret[2] = dram_read_cache_hit_copy_duration / interval;
@@ -1404,6 +1410,7 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
       ret[21] = get_map_actual_used_chunk_in_bytes();
 
       ret[22] = get_num_rows();
+      ret[23] = read_num_counts / interval;
     }
     return ret;
   }
@@ -1659,6 +1666,8 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
   std::atomic<int64_t> inplace_update_hit_cnt_{0};
   std::atomic<int64_t> inplace_update_miss_cnt_{0};
 
+  std::atomic<int64_t> read_num_counts_{0};
+
   bool disable_random_init_;
 }; // class DramKVEmbeddingCache