Added remove_duplicate parameter to nn.Module (#39)

Summary:
Pull Request resolved: pytorch/torchrec#39

Pull Request resolved: pytorch/torchrec#6

This makes it so that shared parameters get their own entry in `named_parameters`.

More broadly, this makes it so that
```
params_and_buffers = {**mod.named_named_parameters(remove_duplicate=False), **mod.named_buffers(remove_duplicate=False)}
_stateless.functional_call(mod, params_and_buffers, args, kwargs)
```
is identical to calling the original module's forwards pass.

cc pietern mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse SciPioneer H-Huang

Pull Request resolved: pytorch#71542

Reviewed By: jbschlosser, albanD

Differential Revision: D33716716

Pulled By: Chillee

fbshipit-source-id: ff1ed9980bd1a3f7ebaf695ee5e401202b543213

Author: Chillee

test/test_nn.py

@@ -1118,6 +1118,24 @@ def names(named_buffers):
             names(s.named_buffers()),
             ['0.dummy_buf', '0.l1.layer_dummy_buf'])
 
+    def test_named_parameters_buffers_duplicates(self):
+        class Foo(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.bias = nn.Parameter(torch.randn(3))
+                self.linear = nn.Linear(3, 3)
+                self.linear.bias = self.bias
+                self.linear_cloned = self.linear
+                self.register_buffer('buffer', torch.randn(3))
+                self.register_buffer('buffer_cloned', self.buffer)
+
+        mod = Foo()
+        self.assertEqual(len(list(mod.named_parameters())), 2)
+        self.assertEqual(len(list(mod.named_parameters(remove_duplicate=False))), 5)
+
+        self.assertEqual(len(list(mod.named_buffers())), 1)
+        self.assertEqual(len(list(mod.named_buffers(remove_duplicate=False))), 2)
+
     def test_call_supports_python_dict_output(self):
         class Net(nn.Module):
             def __init__(self):

torch/distributed/nn/api/remote_module.py

@@ -374,15 +374,15 @@ def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
         )
 
     def named_parameters(  # type: ignore[return]
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, Parameter]]:
         _raise_not_supported(self.named_parameters.__name__)
 
     def buffers(self, recurse: bool = True) -> Iterator[Tensor]:  # type: ignore[return]
         _raise_not_supported(self.buffers.__name__)
 
     def named_buffers(  # type: ignore[return]
-        self, prefix: str = "", recurse: bool = True
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, Tensor]]:
         _raise_not_supported(self.named_buffers.__name__)
 

torch/nn/modules/module.py

@@ -1498,16 +1498,17 @@ def load(module, prefix=''):
                                self.__class__.__name__, "\n\t".join(error_msgs)))
         return _IncompatibleKeys(missing_keys, unexpected_keys)
 
-    def _named_members(self, get_members_fn, prefix='', recurse=True):
+    def _named_members(self, get_members_fn, prefix='', recurse=True, remove_duplicate=True):
         r"""Helper method for yielding various names + members of modules."""
         memo = set()
-        modules = self.named_modules(prefix=prefix) if recurse else [(prefix, self)]
+        modules = self.named_modules(prefix=prefix, remove_duplicate=remove_duplicate) if recurse else [(prefix, self)]
         for module_prefix, module in modules:
             members = get_members_fn(module)
             for k, v in members:
                 if v is None or v in memo:
                     continue
-                memo.add(v)
+                if remove_duplicate:
+                    memo.add(v)
                 name = module_prefix + ('.' if module_prefix else '') + k
                 yield name, v
 
@@ -1535,7 +1536,10 @@ def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
         for name, param in self.named_parameters(recurse=recurse):
             yield param
 
-    def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:
+    def named_parameters(self,
+                         prefix: str = '',
+                         recurse: bool = True,
+                         remove_duplicate: bool = True) -> Iterator[Tuple[str, Parameter]]:
         r"""Returns an iterator over module parameters, yielding both the
         name of the parameter as well as the parameter itself.
 
@@ -1544,6 +1548,9 @@ def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[T
             recurse (bool): if True, then yields parameters of this module
                 and all submodules. Otherwise, yields only parameters that
                 are direct members of this module.
+            remove_duplicate (bool): if True, then removes parameters
+                that are duplicates of each other. For example, if two
+                parameters are tied, it'll only return one.
 
         Yields:
             (string, Parameter): Tuple containing the name and parameter
@@ -1557,7 +1564,7 @@ def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[T
         """
         gen = self._named_members(
             lambda module: module._parameters.items(),
-            prefix=prefix, recurse=recurse)
+            prefix=prefix, recurse=recurse, remove_duplicate=remove_duplicate)
         for elem in gen:
             yield elem
 
@@ -1583,7 +1590,7 @@ def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
         for _, buf in self.named_buffers(recurse=recurse):
             yield buf
 
-    def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Tensor]]:
+    def named_buffers(self, prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) -> Iterator[Tuple[str, Tensor]]:
         r"""Returns an iterator over module buffers, yielding both the
         name of the buffer as well as the buffer itself.
 
@@ -1592,6 +1599,9 @@ def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tupl
             recurse (bool): if True, then yields buffers of this module
                 and all submodules. Otherwise, yields only buffers that
                 are direct members of this module.
+            remove_duplicate (bool): if True, then removes buffers
+                that are duplicates of each other. For example, if two
+                buffers are tied, it'll only return one.
 
         Yields:
             (string, torch.Tensor): Tuple containing the name and buffer
@@ -1605,7 +1615,7 @@ def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tupl
         """
         gen = self._named_members(
             lambda module: module._buffers.items(),
-            prefix=prefix, recurse=recurse)
+            prefix=prefix, recurse=recurse, remove_duplicate=remove_duplicate)
         for elem in gen:
             yield elem