Changing Backend Tensor initialization (#5056)

Summary:

X-link: meta-pytorch/torchrec#3484

X-link: facebookresearch/FBGEMM#2066

**Context:**
Currently, we are enabling SSD optimizer offloading for the ssd tbe kernel

**In this diff:**
We retrieve the newly added parameters from the tbe config and pass it down to the tbe

Differential Revision: D85353134

Author: Raahul Kalyaan Jakka

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -179,6 +179,7 @@ def __init__(
         table_names: Optional[list[str]] = None,
         use_rowwise_bias_correction: bool = False,  # For Adam use
         optimizer_state_dtypes: dict[str, SparseType] = {},  # noqa: B006
+        enable_optimizer_offloading: bool = False,
     ) -> None:
         super(SSDTableBatchedEmbeddingBags, self).__init__()
 

fbgemm_gpu/src/ssd_split_embeddings_cache/embedding_rocksdb_wrapper.h

@@ -45,7 +45,9 @@ class EmbeddingRocksDBWrapper : public torch::jit::CustomClassHolder {
       std::optional<at::Tensor> table_dims = std::nullopt,
       std::optional<at::Tensor> hash_size_cumsum = std::nullopt,
       int64_t flushing_block_size = 2000000000 /*2GB*/,
-      bool disable_random_init = false)
+      bool disable_random_init = false,
+      std::optional<bool> enable_optimizer_offloading = std::nullopt,
+      std::optional<int64_t> optimizer_D = std::nullopt)
       : impl_(
             std::make_shared<ssd::EmbeddingRocksDB>(
                 path,
@@ -77,7 +79,9 @@ class EmbeddingRocksDBWrapper : public torch::jit::CustomClassHolder {
                 table_dims,
                 hash_size_cumsum,
                 flushing_block_size,
-                disable_random_init)) {}
+                disable_random_init,
+                enable_optimizer_offloading,
+                optimizer_D)) {}
 
   void set_cuda(
       at::Tensor indices,

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_split_table_batched_embeddings.cpp

@@ -809,7 +809,9 @@ static auto embedding_rocks_db_wrapper =
                 std::optional<at::Tensor>,
                 std::optional<at::Tensor>,
                 int64_t,
-                bool>(),
+                bool,
+                std::optional<bool>,
+                std::optional<int64_t>>(),
             "",
             {
                 torch::arg("path"),
@@ -842,6 +844,8 @@ static auto embedding_rocks_db_wrapper =
                 torch::arg("hash_size_cumsum") = std::nullopt,
                 torch::arg("flushing_block_size") = 2000000000 /* 2GB */,
                 torch::arg("disable_random_init") = false,
+                torch::arg("enable_optimizer_offloading") = std::nullopt,
+                torch::arg("optimizer_D") = std::nullopt,
             })
         .def(
             "set_cuda",

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_table_batched_embeddings.h

@@ -122,8 +122,8 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
       std::optional<at::Tensor> hash_size_cumsum = std::nullopt,
       int64_t flushing_block_size = 2000000000 /*2GB*/,
       bool disable_random_init = false,
-      bool enable_optimizer_offloading = false,
-      int64_t optimizer_D = 0)
+      std::optional<bool> enable_optimizer_offloading = std::nullopt,
+      std::optional<int64_t> optimizer_D = std::nullopt)
       : kv_db::EmbeddingKVDB(
             num_shards,
             max_D,
@@ -426,8 +426,8 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
       float uniform_init_upper,
       int64_t row_storage_bitwidth,
       bool disable_random_init,
-      bool enable_optimizer_offloading = false,
-      int64_t optimizer_D = 0) {
+      std::optional<bool> enable_optimizer_offloading = std::nullopt,
+      std::optional<int64_t> optimizer_D = std::nullopt) {
     for (auto i = 0; i < num_shards; ++i) {
       auto* gen = at::check_generator<at::CPUGeneratorImpl>(
           at::detail::getDefaultCPUGenerator());
@@ -442,9 +442,13 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
 
         // When Optimizer offloading is enabled, we want to initialize the last
         // optimizer_D columns(optimizer values) to zero
-        if (enable_optimizer_offloading) {
+        if (enable_optimizer_offloading.has_value() &&
+            enable_optimizer_offloading.value() && optimizer_D.has_value()) {
           auto& tensor = initializer->row_storage_;
-          tensor.index({"...", at::indexing::Slice(max_D - optimizer_D, max_D)})
+          tensor
+              .index(
+                  {"...",
+                   at::indexing::Slice(max_D - optimizer_D.value(), max_D)})
               .zero_();
         }
         initializers_.push_back(std::move(initializer));
@@ -1378,6 +1382,7 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
   std::vector<std::string> db_paths_;
 
   bool disable_random_init_;
+  std::optional<bool> enable_optimizer_offloading = std::nullopt;
 }; // class EmbeddingRocksDB
 
 /// @ingroup embedding-ssd