Move instance of params from TBE fused kernels to named_parameters (p…

…ytorch#893)

Summary:
Pull Request resolved: pytorch#893

Fixing long standing nit bug with TorchRec's named parameters not returning fused parameters.

Model code needs to change to account for this, we do this by adding the filter_optimizer_in_backward_named_parameters

This sets the stage for the composability change

The important parts to change, being MVAI/O3/PyPer should have this change. Tried to get everything else,

Most important files to look at IMO are

* full_sync_optimizer.py
* fbcode/torchrec/optim/optimizers.py
* torchrec/fb/optim/module_optimizer.py

Reviewed By: YLGH

Differential Revision: D41964643

fbshipit-source-id: 28f82a56a2147d10389ad4fd0e8dbaef78f12e3c

torchrec/distributed/batched_embedding_kernel.py

@@ -413,12 +413,20 @@ def named_buffers(
         By convention, fused parameters are designated as buffers because they no longer
         have gradients available to external optimizers.
         """
-        return self.named_split_embedding_weights(prefix, recurse, remove_duplicate)
+        # TODO can delete this override once SEA is removed
+        yield from ()
 
     def named_parameters(
         self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, nn.Parameter]]:
-        yield from ()
+        for name, tensor in self.named_split_embedding_weights(
+            prefix, recurse, remove_duplicate
+        ):
+            # hack before we support optimizer on sharded parameter level
+            param = nn.Parameter(tensor)
+            # pyre-ignore
+            param._overlapped_optimizer = True
+            yield name, param
 
     def flush(self) -> None:
         self._emb_module.flush()
@@ -570,12 +578,12 @@ def named_split_embedding_weights(
         assert (
             remove_duplicate
         ), "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
-        for config, param in zip(
+        for config, tensor in zip(
             self._config.embedding_tables,
             self.emb_module.split_embedding_weights(),
         ):
             key = append_prefix(prefix, f"{config.name}.weight")
-            yield key, param
+            yield key, tensor
 
 
 class BatchedFusedEmbeddingBag(BaseBatchedEmbeddingBag, FusedOptimizerModule):
@@ -645,12 +653,20 @@ def named_buffers(
         By convention, fused parameters are designated as buffers because they no longer
         have gradients available to external optimizers.
         """
-        return self.named_split_embedding_weights(prefix, recurse, remove_duplicate)
+        # TODO can delete this override once SEA is removed
+        yield from ()
 
     def named_parameters(
         self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, nn.Parameter]]:
-        yield from ()
+        for name, tensor in self.named_split_embedding_weights(
+            prefix, recurse, remove_duplicate
+        ):
+            # hack before we support optimizer on sharded parameter level
+            param = nn.Parameter(tensor)
+            # pyre-ignore
+            param._overlapped_optimizer = True
+            yield name, param
 
     def flush(self) -> None:
         self._emb_module.flush()