Add raw embedding streaming needed params in trec and mvai (#2935)

Summary:
Pull Request resolved: #2935

Add the variables needed in D73792631 to mvai and torch rec to be able to control them via config.

Reviewed By: aliafzal

Differential Revision: D74086201

fbshipit-source-id: 53fb269c17f08d87589a837d2049b733db0d665e

Author: chouxi

torchrec/distributed/batched_embedding_kernel.py

@@ -37,6 +37,7 @@
     DenseTableBatchedEmbeddingBagsCodegen,
     EmbeddingLocation,
     PoolingMode,
+    RESParams,
     SparseType,
     SplitTableBatchedEmbeddingBagsCodegen,
 )
@@ -171,6 +172,23 @@ def _populate_ssd_tbe_params(config: GroupedEmbeddingConfig) -> Dict[str, Any]:
             )
             ssd_tbe_params["cache_sets"] = int(max_cache_sets)
 
+    # populate res_params, which is used for raw embedding streaming
+    # here only populates the params available in fused_params and TBE configs
+    res_params: RESParams = RESParams()
+    res_params.table_names = [table.name for table in config.embedding_tables]
+    res_params.table_offsets = []
+    for emb_tbl in config.embedding_tables:
+        local_metadata = emb_tbl.local_metadata
+        if (
+            local_metadata is not None
+            and local_metadata.shard_offsets is not None
+            and len(local_metadata.shard_offsets) >= 1
+        ):
+            res_params.table_offsets.append(local_metadata.shard_offsets[0])
+    if "res_store_shards" in fused_params:
+        res_params.res_store_shards = fused_params.get("res_store_shards")
+    ssd_tbe_params["res_params"] = res_params
+
     return ssd_tbe_params
 
 

torchrec/distributed/types.py

@@ -640,6 +640,7 @@ class KeyValueParams:
         bulk_init_chunk_size: int: number of rows to insert into rocksdb in each chunk
         lazy_bulk_init_enabled: bool: whether to enable lazy(async) bulk init for SSD TBE
         enable_raw_embedding_streaming: Optional[bool]: enable raw embedding streaming for SSD TBE
+        res_store_shards: Optional[int] = None: the number of shards to store the raw embeddings
 
         # Parameter Server (PS) Attributes
         ps_hosts (Optional[Tuple[Tuple[str, int]]]): List of PS host ip addresses
@@ -664,6 +665,7 @@ class KeyValueParams:
     enable_raw_embedding_streaming: Optional[bool] = (
         None  # enable raw embedding streaming for SSD TBE
     )
+    res_store_shards: Optional[int] = None  # shards to store the raw embeddings
 
     # Parameter Server (PS) Attributes
     ps_hosts: Optional[Tuple[Tuple[str, int], ...]] = None
@@ -691,6 +693,7 @@ def __hash__(self) -> int:
                 self.bulk_init_chunk_size,
                 self.lazy_bulk_init_enabled,
                 self.enable_raw_embedding_streaming,
+                self.res_store_shards,
             )
         )