sys.getrefcount() of a LightningModule

[MGheini]

🐛 Bug

When creating a LightningModule, sys.getrefcount() returns 3 (vs. returning 2 for a conventional torch Module). In my particular case, for instance, this results in memory issues: even if I explicitly do del <LightningModule>, the memory is not released, and I believe this is the root of the issue.

I had originally discussed this problem on Slack. As discussed, I'm creating an issue for it and pinging @carmocca.

To Reproduce

Input:

import sys
import torch
import pytorch_lightning as pl

a = torch.nn.Module()
b = pl.LightningModule()
print(sys.getrefcount(b))
print(sys.getrefcount(a))

Output:

3
2

Expected behavior

Output:

2
2

Environment

• PyTorch Lightning Version: 1.5.9
• PyTorch Version: 1.10.2+cu113
• Python version: 3.8.12
• OS: Linux
• CUDA/cuDNN version: 11.6
• GPU models and configuration: Quadro RTX 8000
• How you installed PyTorch: pip

Additional context

Per @carmocca's advice on Slack, a workaround is to do b._load_state_dict_pre_hooks.clear() before deleting the LightningModule. This indeed seems to work. However, I'm not sure about an explanation as to why.

Thanks!

cc @carmocca @justusschock @awaelchli @Borda @ananthsub @ninginthecloud @jjenniferdai @rohitgr7

[carmocca]

This is caused by this call in the __init__: https://github.com/PyTorchLightning/pytorch-lightning/blob/de085f30f05512862d970e1675536b4c70199d27/pytorch_lightning/core/lightning.py#L115 concretely, this line where we call a sharded internal method: https://github.com/PyTorchLightning/pytorch-lightning/blob/de085f30f05512862d970e1675536b4c70199d27/pytorch_lightning/core/lightning.py#L2026 that keeps an owning reference to self: https://github.com/pytorch/pytorch/blob/76e9730d0273058297f6adce76e73ac7e3db25f0/torch/nn/modules/module.py#L1406

There are multiple ways to fix it, but I wonder if the proper fix should be done upstream with weakref.proxy(self) to avoid the owning reference.

As a workaround, we can clear the _load_state_dict_pre_hooks dictionary on __del__.

cc @yifuwang as the author of the sharded integration (#8944) and @pritamdamania87 as the author of the relevant piece of code upstream (pytorch/pytorch#62070)