Thread raw embedding streamer to dram_kv_embedding_cache (#5432)

Summary:
Pull Request resolved: #5432

X-link: https://github.com/facebookresearch/FBGEMM/pull/2404

Thread RES (Raw Embedding Streaming) parameters through the DRAM KV
embedding cache constructor chain and pybind to enable streaming
for the embedding cache enrichment path.
Currently the feature is gated by `enable_raw_embedding_streaming`
Key changes:

- Thread 6 RES params (DramKVEmbeddingCache -> wrapper -> pybind -> Python)
- Make raw_embedding_streamer_ protected for subclass access

Reviewed By: FriedCosey

Differential Revision: D94431329

fbshipit-source-id: 7af10855718fe24b77de1d66d1437d681e47bd48

Author: chouxi

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -783,6 +783,12 @@ def __init__(
                 self.backend_return_whole_row,  # backend_return_whole_row
                 False,  # enable_async_update
                 self._embedding_cache_mode,  # disable_random_init
+                self.enable_raw_embedding_streaming,
+                self.res_params.res_store_shards,
+                self.res_params.res_server_port,
+                self.res_params.table_names,
+                self.res_params.table_offsets,
+                self.res_params.table_sizes,
             )
         else:
             raise AssertionError(f"Invalid backend type {self.backend_type}")

fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h

@@ -110,14 +110,26 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
       std::optional<at::Tensor> table_dims = std::nullopt,
       std::optional<at::Tensor> hash_size_cumsum = std::nullopt,
       bool is_training = true,
-      bool disable_random_init = false)
+      bool disable_random_init = false,
+      bool enable_raw_embedding_streaming = false,
+      int64_t res_store_shards = 0,
+      int64_t res_server_port = 0,
+      std::vector<std::string> table_names = {},
+      std::vector<int64_t> table_offsets = {},
+      std::vector<int64_t> table_sizes = {})
       : kv_db::EmbeddingKVDB(
             num_shards,
             max_D,
             0, // l2_cache_size_gb =0 to disable l2 cache
             0, // tbe_unqiue_id
             2, // ele_size_bytes
-            enable_async_update),
+            enable_async_update,
+            enable_raw_embedding_streaming,
+            res_store_shards,
+            res_server_port,
+            std::move(table_names),
+            std::move(table_offsets),
+            table_sizes),
         max_D_(max_D),
         num_shards_(num_shards),
         block_size_(FixedBlockPool::calculate_block_size<weight_type>(max_D)),

fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache_wrapper.h

@@ -33,7 +33,13 @@ class DramKVEmbeddingCacheWrapper : public torch::jit::CustomClassHolder {
       const std::optional<at::Tensor>& hash_size_cumsum = std::nullopt,
       bool backend_return_whole_row = false,
       bool enable_async_update = false,
-      bool disable_random_init = false) {
+      bool disable_random_init = false,
+      bool enable_raw_embedding_streaming = false,
+      int64_t res_store_shards = 0,
+      int64_t res_server_port = 0,
+      std::vector<std::string> table_names = {},
+      std::vector<int64_t> table_offsets = {},
+      std::vector<int64_t> table_sizes = {}) {
     if (row_storage_bitwidth == 16) {
       impl_ = std::make_shared<kv_mem::DramKVEmbeddingCache<at::Half>>(
           max_D,
@@ -48,7 +54,13 @@ class DramKVEmbeddingCacheWrapper : public torch::jit::CustomClassHolder {
           table_dims,
           hash_size_cumsum,
           true, // is_training
-          disable_random_init);
+          disable_random_init,
+          enable_raw_embedding_streaming,
+          res_store_shards,
+          res_server_port,
+          std::move(table_names),
+          std::move(table_offsets),
+          std::move(table_sizes));
     } else if (row_storage_bitwidth == 32) {
       impl_ = std::make_shared<kv_mem::DramKVEmbeddingCache<float>>(
           max_D,
@@ -63,7 +75,13 @@ class DramKVEmbeddingCacheWrapper : public torch::jit::CustomClassHolder {
           table_dims,
           hash_size_cumsum,
           true, // is_training
-          disable_random_init);
+          disable_random_init,
+          enable_raw_embedding_streaming,
+          res_store_shards,
+          res_server_port,
+          std::move(table_names),
+          std::move(table_offsets),
+          std::move(table_sizes));
     } else {
       throw std::runtime_error("Failed to create recording device");
     }

fbgemm_gpu/src/ssd_split_embeddings_cache/kv_db_table_batched_embeddings.h

@@ -528,6 +528,8 @@ class EmbeddingKVDB : public std::enable_shared_from_this<EmbeddingKVDB> {
 
   // -- commone path
   std::atomic<int64_t> total_cache_update_duration_{0};
+
+ protected:
   std::unique_ptr<fbgemm_gpu::RawEmbeddingStreamer> raw_embedding_streamer_;
 }; // class EmbeddingKVDB
 

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_split_table_batched_embeddings.cpp

@@ -978,7 +978,13 @@ static auto dram_kv_embedding_cache_wrapper =
                 std::optional<at::Tensor>,
                 bool,
                 bool,
-                bool>(),
+                bool,
+                bool,
+                int64_t,
+                int64_t,
+                std::vector<std::string>,
+                std::vector<int64_t>,
+                std::vector<int64_t>>(),
             "",
             {
                 torch::arg("max_D"),
@@ -993,6 +999,12 @@ static auto dram_kv_embedding_cache_wrapper =
                 torch::arg("backend_return_whole_row") = false,
                 torch::arg("enable_async_update") = false,
                 torch::arg("disable_random_init") = false,
+                torch::arg("enable_raw_embedding_streaming") = false,
+                torch::arg("res_store_shards") = 0,
+                torch::arg("res_server_port") = 0,
+                torch::arg("table_names") = std::vector<std::string>{},
+                torch::arg("table_offsets") = std::vector<int64_t>{},
+                torch::arg("table_sizes") = std::vector<int64_t>{},
             })
         .def(
             "set_cuda",