torchrec changes v4 (#3348)

Summary:
Pull Request resolved: #3348

add utility functions to re-initialize torch states of ShardedEmbeddingBag classes

Reviewed By: liangbeixu

Differential Revision: D81154653

fbshipit-source-id: 80e90575f66a01d127558420842686e45669a3ff

Author: jiayulu

torchrec/distributed/batched_embedding_kernel.py

@@ -2310,15 +2310,22 @@ def __init__(
             self._emb_module,
             pg,
         )
-        self._param_per_table: Dict[str, TableBatchedEmbeddingSlice] = dict(
+        self.init_parameters()
+
+    @property
+    def _param_per_table(self) -> Dict[str, TableBatchedEmbeddingSlice]:
+        return dict(
             _gen_named_parameters_by_table_fused(
                 emb_module=self._emb_module,
                 table_name_to_count=self.table_name_to_count.copy(),
                 config=self._config,
-                pg=pg,
+                pg=self._pg,
             )
         )
-        self.init_parameters()
+
+    @_param_per_table.setter
+    def _param_per_table(self, v: Dict[str, TableBatchedEmbeddingSlice]) -> None:
+        self.__dict__["_param_per_table"] = v
 
     @property
     def emb_module(
@@ -3169,15 +3176,22 @@ def __init__(
             self._emb_module,
             pg,
         )
-        self._param_per_table: Dict[str, TableBatchedEmbeddingSlice] = dict(
+        self.init_parameters()
+
+    @property
+    def _param_per_table(self) -> Dict[str, TableBatchedEmbeddingSlice]:
+        return dict(
             _gen_named_parameters_by_table_fused(
                 emb_module=self._emb_module,
                 table_name_to_count=self.table_name_to_count.copy(),
                 config=self._config,
-                pg=pg,
+                pg=self._pg,
             )
         )
-        self.init_parameters()
+
+    @_param_per_table.setter
+    def _param_per_table(self, v: Dict[str, TableBatchedEmbeddingSlice]) -> None:
+        self.__dict__["_param_per_table"] = v
 
     @property
     def emb_module(

torchrec/distributed/fp_embeddingbag.py

@@ -163,6 +163,9 @@ def sharded_parameter_names(self, prefix: str = "") -> Iterator[str]:
             if "_embedding_bag_collection" in fqn:
                 yield append_prefix(prefix, fqn)
 
+    def _initialize_torch_state(self, skip_registering: bool = False) -> None:  # noqa
+        self._embedding_bag_collection._initialize_torch_state(skip_registering)
+
 
 class FeatureProcessedEmbeddingBagCollectionSharder(
     BaseEmbeddingSharder[FeatureProcessedEmbeddingBagCollection]