Repository: rwth-i6/returnn

Issue:
PyTorch `weight_decay` ambiguity with shared params
The code logic of `_get_optimizer_param_groups` branch with activated weight decay can not run in its current form. Iterating over all named modules while at the same time iterating recursively over parameters will yield the same parameter multiple times with a different module reference. 

In my case this was:
```python
<class 'i6_experiments.users.rossenbach.experiments.librispeech.tts_architecture_improvement_23.pytorch_networks.ctc_aligner_v1.Model'>                                                                            
<class 'torch.nn.modules.sparse.Embedding'> 
```
for
`speaker_embedding.weight`

This means we need a completely new logic if we want to exclude some modules.



Discussion:
I don't understand. Why is it not possible to fix it? Why not make the params unique?

You would just keep a list of visited params, and if some param is already in there, skip it.

Can you please give a short example in which different modules share common parameters? I didn't know that this was possible (although it would make sense) and a brief demo would help me grasp the concept.

In principle we can do what @albertz mentioned, keeping a set of `Parameter` references that have been already visited; then the code wouldn't need much tinkering. Assuming that the `Parameter` reference is the same, of course, but it seems to be the issue you're having to begin with, so I think this would solve the issue.
> Can you please give a short example in which different modules share common parameters? I didn't know that this was possible (although it would make sense) and a brief demo would help me grasp the concept.
> 

Any nested modules will cause this problem, because the parameter access is recursive. Making it not recursive might already solve the problem, but then we can not use strings anymore like it is currently done, because you get name repetitions.

> In principle we can do what @albertz mentioned, keeping a set of `Parameter` references that have been already visited; then the code wouldn't need much tinkering. Assuming that the `Parameter` reference is the same, of course, but it seems to be the issue you're having to begin with, so I think this would solve the issue.

It does not solve this issue, because we are testing for the module name in the blacklist. Thus in my case the same parameter was added in both parameter groups and I got a crash.

For sure there are solutions, but as after some minutes I did not have one, I wanted to post this.


> It does not solve this issue

Sure, it will solve it. You will never visit the same param twice. Thus the problem is solved.

See the code in `rf.Module.named_parameters()` as an example.

I think there might be a solution if we check only for collisions with respect to the `Parameter` itself (`p` in the code), instead of the full parameter name (`fpn` in the code). I'm working on a fix, which you can check out in an isolated branch (an exception in this case for the "working in main" workflow, given the isolate fix).

Resolving patch:
diff --git a/returnn/torch/updater.py b/returnn/torch/updater.py
--- a/returnn/torch/updater.py
+++ b/returnn/torch/updater.py
@@ -7,8 +7,10 @@
 
 import torch
 import typing
+from typing import Set
 
 from returnn.log import log
+from returnn.util.basic import RefIdEq
 
 _OptimizerClassesDictInitialized = False
 _OptimizerClassesDict = {}
@@ -241,8 +243,14 @@ def _get_optimizer_param_groups(self, optim_class, optimizer_opts):
         wd_params = set()
         no_wd_params = set()
         blacklist_wd_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
+        # Tracker of visited parameters to only add each parameter once, in case two modules share common parameters.
+        # We need the wrapper class RefIdEq because Parameters are compared by value and not by reference.
+        visited_params: Set[RefIdEq[torch.nn.Parameter]] = set()
         for mn, m in self.network.named_modules():
             for pn, p in m.named_parameters():
+                if RefIdEq(p) in visited_params:
+                    continue
+                visited_params.add(RefIdEq(p))
                 fpn = "%s.%s" % (mn, pn) if mn else pn  # Full param name
                 if pn.endswith("bias"):
                     no_wd_params.add(fpn)


Test patch: