Differences in layer ordering causes errors when resuming from checkpoints

Bug description

When resuming from a checkpoint (or multi-gpu training with ddp, anything where you might construct multiple instances of your model), if the model layers are constructed in different orders, it will cause the optimizer stats to have parameters in different orders, which will cause a size mismatch error. In the cases where the sizes agree this error will go unnoticed (except with perhaps less optimal training dynamics). When the sizes disagree, then it explicitly errors.

Not sure if this is really a torch issue or a lightning issue. I'm submitting here first.

How to reproduce the bug

I have a MWE that reproduces the issue:

https://gist.github.com/Erotemic/dfaadf5cf9fa4910beb901ae6c93867b

It's about 400 lines that try to capture both the error and my use case that uncovered the error. I had to use this to even figure out what was going on, so I'm sure the example could be more minimal. I've highlighted where the part that causes the error. If you sort the set to make construction consistent the error goes away - so there is a user-side fix, but I'm not sure if it also warents a torch or lightning side fix.

The idea is we have multiple image sensors that observe multiple different bands at different resolutions and we are putting them all into a single network. I define a ModuleDict of stems to normalize the number of channels in each type of input so I can concatenate them into tokens for a transformer. The issue was that I was iterating over a set to construct entries in the ModuleDict, and that is non-determenistic between python runs. It does seem strange that either lightning / torch can't deal with this though. They modules are keyed, so you should be able to order them.

Error messages and logs

Epoch 1:  14%|████████████▊                                                                                 | 15/110 [00:00<00:06, 15.63it/s, loss=0.894, v_num=4, train_loss=0.955]Traceback (most recent call last):
  File "/home/joncrall/code/watch/dev/mwe/lightning_cli_ckpt_path_error.py", line 437, in <module>
    main()
  File "/home/joncrall/code/watch/dev/mwe/lightning_cli_ckpt_path_error.py", line 380, in main
    MWE_LightningCLI(

...

    step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
  File "/home/joncrall/.pyenv/versions/3.10.10/envs/pyenv3.10.10/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 234, in optimizer_step
    return self.precision_plugin.optimizer_step(
  File "/home/joncrall/.pyenv/versions/3.10.10/envs/pyenv3.10.10/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 119, in optimizer_step
    return optimizer.step(closure=closure, **kwargs)
  File "/home/joncrall/.pyenv/versions/3.10.10/envs/pyenv3.10.10/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 68, in wrapper
    return wrapped(*args, **kwargs)
  File "/home/joncrall/.pyenv/versions/3.10.10/envs/pyenv3.10.10/lib/python3.10/site-packages/torch/optim/optimizer.py", line 140, in wrapper
    out = func(*args, **kwargs)
  File "/home/joncrall/.pyenv/versions/3.10.10/envs/pyenv3.10.10/lib/python3.10/site-packages/torch/optim/optimizer.py", line 23, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/home/joncrall/.pyenv/versions/3.10.10/envs/pyenv3.10.10/lib/python3.10/site-packages/torch/optim/adam.py", line 234, in step
    adam(params_with_grad,
  File "/home/joncrall/.pyenv/versions/3.10.10/envs/pyenv3.10.10/lib/python3.10/site-packages/torch/optim/adam.py", line 300, in adam
    func(params,
  File "/home/joncrall/.pyenv/versions/3.10.10/envs/pyenv3.10.10/lib/python3.10/site-packages/torch/optim/adam.py", line 363, in _single_tensor_adam
    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 1

Environment

(pyenv3.10.10) joncrall@toothbrush:~/code/watch/dev/mwe$ pyversion pytorch_lightning
python -c "import pytorch_lightning; print('pytorch_lightning.__version__ = ' + str(pytorch_lightning.__version__))"
pytorch_lightning.__version__ = 1.9.4
(pyenv3.10.10) joncrall@toothbrush:~/code/watch/dev/mwe$ pyversion torch
python -c "import torch; print('torch.__version__ = ' + str(torch.__version__))"
torch.__version__ = 1.13.1+cu117
(pyenv3.10.10) joncrall@toothbrush:~/code/watch/dev/mwe$ python --version
Python 3.10.10

More info

No response

Start with the linked MWE and reproduce the failure using the reported PyTorch Lightning 1.9.4, PyTorch 1.13.1, and Python 3.10 environment. Compare checkpoint resumption with deterministic layer construction against construction from a set, then trace whether optimizer state restoration belongs to Lightning or PyTorch. Done means the ownership and expected behavior are established and a regression test or documented limitation is identified.