pytorch
diff --git a/‎torchrec/distributed/batched_embedding_kernel.py
Lines changed: 14 additions & 8 deletions b/‎torchrec/distributed/batched_embedding_kernel.py
Lines changed: 14 additions & 8 deletions
diff --git a/‎torchrec/distributed/embedding_lookup.py
Lines changed: 20 additions & 4 deletions b/‎torchrec/distributed/embedding_lookup.py
Lines changed: 20 additions & 4 deletions
diff --git a/‎torchrec/distributed/embeddingbag.py
Lines changed: 4 additions & 3 deletions b/‎torchrec/distributed/embeddingbag.py
Lines changed: 4 additions & 3 deletions
diff --git a/‎torchrec/distributed/grouped_position_weighted.py
Lines changed: 1 addition & 1 deletion b/‎torchrec/distributed/grouped_position_weighted.py
Lines changed: 1 addition & 1 deletion
diff --git a/‎torchrec/distributed/model_parallel.py
Lines changed: 3 additions & 2 deletions b/‎torchrec/distributed/model_parallel.py
Lines changed: 3 additions & 2 deletions
diff --git a/‎torchrec/distributed/quant_embedding_kernel.py
Lines changed: 4 additions & 1 deletion b/‎torchrec/distributed/quant_embedding_kernel.py
Lines changed: 4 additions & 1 deletion
diff --git a/‎torchrec/modules/feature_processor.py
Lines changed: 1 addition & 1 deletion b/‎torchrec/modules/feature_processor.py
Lines changed: 1 addition & 1 deletion
diff --git a/‎torchrec/modules/fused_embedding_modules.py
Lines changed: 4 additions & 1 deletion b/‎torchrec/modules/fused_embedding_modules.py
Lines changed: 4 additions & 1 deletion
@@ -337,8 +337,11 @@ def flush(self) -> None:
         pass
 
     def named_split_embedding_weights(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
+        assert (
+            remove_duplicate
+        ), "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
         for config, param in zip(
             self._config.embedding_tables,
             self.emb_module.split_embedding_weights(),
@@ -404,13 +407,13 @@ def fused_optimizer(self) -> FusedOptimizer:
         return self._optim
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
         """
         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)
+        return self.named_split_embedding_weights(prefix, recurse, remove_duplicate)
 
     def named_parameters(
         self, prefix: str = "", recurse: bool = True
@@ -452,7 +455,7 @@ def emb_module(
         return self._emb_module
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
         yield from ()
 
@@ -562,8 +565,11 @@ def flush(self) -> None:
         pass
 
     def named_split_embedding_weights(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
+        assert (
+            remove_duplicate
+        ), "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
         for config, param in zip(
             self._config.embedding_tables,
             self.emb_module.split_embedding_weights(),
@@ -633,13 +639,13 @@ def fused_optimizer(self) -> FusedOptimizer:
         return self._optim
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
         """
         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)
+        return self.named_split_embedding_weights(prefix, recurse, remove_duplicate)
 
     def named_parameters(
         self, prefix: str = "", recurse: bool = True
@@ -681,7 +687,7 @@ def emb_module(
         return self._emb_module
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
         yield from ()
 
 
@@ -182,8 +182,12 @@ def named_parameters(
             yield from emb_module.named_parameters(prefix, recurse)
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
+        assert remove_duplicate, (
+            "remove_duplicate=False in named_buffers for"
+            "GroupedEmbeddingsLookup is not supported"
+        )
         for emb_module in self._emb_modules:
             yield from emb_module.named_buffers(prefix, recurse)
 
@@ -350,8 +354,12 @@ def named_parameters(
             yield from emb_module.named_parameters(prefix, recurse)
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
+        assert remove_duplicate, (
+            "remove_duplicate=False in named_buffers for"
+            "GroupedPooledEmbeddingsLookup is not supported"
+        )
         for emb_module in self._emb_modules:
             yield from emb_module.named_buffers(prefix, recurse)
         for emb_module in self._score_emb_modules:
@@ -459,8 +467,12 @@ def named_parameters(
             yield from emb_module.named_parameters(prefix, recurse)
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
+        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_buffers(prefix, recurse)
 
@@ -613,8 +625,12 @@ def named_parameters(
             yield from emb_module.named_parameters(prefix, recurse)
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
+        assert remove_duplicate, (
+            "remove_duplicate=False in named_buffers for"
+            "MetaInferGroupedPooledEmbeddingsLookup is not supported"
+        )
         for emb_module in self._emb_modules:
             yield from emb_module.named_buffers(prefix, recurse)
         for emb_module in self._score_emb_modules:
 
@@ -490,11 +490,11 @@ def sharded_parameter_names(self, prefix: str = "") -> Iterator[str]:
                 yield name
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
         for lookup in self._lookups:
             yield from lookup.named_buffers(
-                append_prefix(prefix, "embedding_bags"), recurse
+                append_prefix(prefix, "embedding_bags"), recurse, remove_duplicate
             )
 
     # pyre-fixme[14]: `load_state_dict` overrides method defined in `Module`
@@ -744,8 +744,9 @@ def sharded_parameter_names(self, prefix: str = "") -> Iterator[str]:
                 yield append_prefix(prefix, name.split(".")[-1])
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
+        # TODO: add remove_duplicate
         for name, buffer in self._lookup.named_buffers("", recurse):
             yield append_prefix(prefix, name.split(".")[-1]), buffer
 
 
@@ -76,7 +76,7 @@ def named_parameters(
             yield append_prefix(prefix, f"position_weights.{name}"), param
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
         yield from ()
 
 
@@ -550,14 +550,15 @@ def _named_buffers(
                 )
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
         gen = self._named_buffers(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
 
     @property
 
@@ -167,8 +167,11 @@ def forward(self, features: KeyedJaggedTensor) -> torch.Tensor:
         )
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
+        assert (
+            remove_duplicate
+        ), "remove_duplicate=False not supported in QuantBatchedEmbeddingBag.named_split_embedding_weights"
         for config, weight in zip(
             self._config.embedding_tables,
             self.emb_module.split_embedding_weights(),
 
@@ -277,7 +277,7 @@ def forward(self, features: KeyedJaggedTensor) -> KeyedJaggedTensor:
                 return features
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
         yield from ()
 
 
@@ -232,8 +232,11 @@ def named_parameters(
             yield key, cast(nn.Parameter, weight)
 
     def named_buffers(
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
+        assert (
+            remove_duplicate
+        ), "remove_duplicate=False not supported in _BatchedFusedEmbeddingLookups.named_buffers"
         for table, param in zip(self._embedding_tables, self.split_embedding_weights()):
             name = f"{table.name}.weight"
             key = f"{prefix}.{name}" if (prefix and name) else (prefix + name)