add remove_duplicate flag to named_parameters (pytorch#759)

Summary:
Pull Request resolved: pytorch#759

Since the remove_duplicate flag was added to named_buffers in D39493161 (pytorch@672d105), this adds the same flag to named_parameters

Differential Revision: D40801899

fbshipit-source-id: a01807c81f0069ea02bdb95d3f937d458f3d9a5d

torchrec/distributed/batched_embedding_kernel.py

@@ -416,7 +416,7 @@ def named_buffers(
         return self.named_split_embedding_weights(prefix, recurse, remove_duplicate)
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, nn.Parameter]]:
         yield from ()
 
@@ -460,7 +460,7 @@ def named_buffers(
         yield from ()
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, nn.Parameter]]:
         combined_key = "/".join(
             [config.name for config in self._config.embedding_tables]
@@ -648,7 +648,7 @@ def named_buffers(
         return self.named_split_embedding_weights(prefix, recurse, remove_duplicate)
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, nn.Parameter]]:
         yield from ()
 
@@ -692,7 +692,7 @@ def named_buffers(
         yield from ()
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, nn.Parameter]]:
         combined_key = "/".join(
             [config.name for config in self._config.embedding_tables]

torchrec/distributed/embedding_lookup.py

@@ -176,8 +176,12 @@ def load_state_dict(
         return _IncompatibleKeys(missing_keys=m, unexpected_keys=u)
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
-    ) -> Iterator[Tuple[str, nn.Parameter]]:
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
+    ) -> Iterator[Tuple[str, torch.nn.Parameter]]:
+        assert remove_duplicate, (
+            "remove_duplicate=False in named_parameters for"
+            "GroupedEmbeddingsLookup is not supported"
+        )
         for emb_module in self._emb_modules:
             yield from emb_module.named_parameters(prefix, recurse)
 
@@ -346,8 +350,12 @@ def load_state_dict(
         return _IncompatibleKeys(missing_keys=m1 + m2, unexpected_keys=u1 + u2)
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
-    ) -> Iterator[Tuple[str, nn.Parameter]]:
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
+    ) -> Iterator[Tuple[str, torch.nn.Parameter]]:
+        assert remove_duplicate, (
+            "remove_duplicate=False in named_parameters for"
+            "GroupedPooledEmbeddingsLookup is not supported"
+        )
         for emb_module in self._emb_modules:
             yield from emb_module.named_parameters(prefix, recurse)
         for emb_module in self._score_emb_modules:
@@ -461,8 +469,12 @@ def load_state_dict(
         return _IncompatibleKeys(missing_keys=m, unexpected_keys=u)
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
-    ) -> Iterator[Tuple[str, nn.Parameter]]:
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
+    ) -> Iterator[Tuple[str, torch.nn.Parameter]]:
+        assert remove_duplicate, (
+            "remove_duplicate=False in named_buffers for"
+            "MetaInferGroupedEmbeddingsLookup is not supported"
+        )
         for emb_module in self._emb_modules:
             yield from emb_module.named_parameters(prefix, recurse)
 
@@ -617,8 +629,12 @@ def load_state_dict(
         return _IncompatibleKeys(missing_keys=m1 + m2, unexpected_keys=u1 + u2)
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
-    ) -> Iterator[Tuple[str, nn.Parameter]]:
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
+    ) -> Iterator[Tuple[str, torch.nn.Parameter]]:
+        assert remove_duplicate, (
+            "remove_duplicate=False in named_parameters for"
+            "MetaInferGroupedPooledEmbeddingsLookup is not supported"
+        )
         for emb_module in self._emb_modules:
             yield from emb_module.named_parameters(prefix, recurse)
         for emb_module in self._score_emb_modules:

torchrec/distributed/embeddingbag.py

@@ -531,11 +531,11 @@ def named_modules(
         yield from [(prefix, self)]
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, nn.Parameter]]:
         for lookup in self._lookups:
             yield from lookup.named_parameters(
-                append_prefix(prefix, "embedding_bags"), recurse
+                append_prefix(prefix, "embedding_bags"), recurse, remove_duplicate
             )
 
     def sharded_parameter_names(self, prefix: str = "") -> Iterator[str]:
@@ -789,8 +789,9 @@ def named_modules(
         yield from [(prefix, self)]
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, nn.Parameter]]:
+        # TODO: add remove_duplicate
         for name, parameter in self._lookup.named_parameters("", recurse):
             # update name to match embeddingBag parameter name
             yield append_prefix(prefix, name.split(".")[-1]), parameter

torchrec/distributed/grouped_position_weighted.py

@@ -70,7 +70,7 @@ def forward(self, features: KeyedJaggedTensor) -> KeyedJaggedTensor:
         )
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, nn.Parameter]]:
         for name, param in self.position_weights.items():
             yield append_prefix(prefix, f"position_weights.{name}"), param

torchrec/distributed/model_parallel.py

@@ -495,14 +495,15 @@ def _named_parameters(
                 )
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.nn.Parameter]]:
         gen = self._named_parameters(self.module, prefix, recurse)
         memo = set()
         for key, param in gen:
             if param in memo:
                 continue
-            memo.add(param)
+            if remove_duplicate:
+                memo.add(param)
             yield key, param
 
     def bare_named_parameters(

torchrec/modules/fused_embedding_modules.py

@@ -222,8 +222,11 @@ def parameters(
         yield cast(nn.Parameter, self._emb_module.weights)
 
     def named_parameters(
-        self, prefix: str = "", recurse: bool = True
-    ) -> Iterator[Tuple[str, nn.Parameter]]:
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
+    ) -> Iterator[Tuple[str, torch.nn.Parameter]]:
+        assert (
+            remove_duplicate
+        ), "remove_duplicate=False not supported in _BatchedFusedEmbeddingLookups.named_parameters"
         for table, weight in zip(
             self._embedding_tables, self.split_embedding_weights()
         ):