Repository: Lightning-AI/pytorch-lightning
Pull request: Fix number of references to LightningModule
## What does this PR do?

I know, this looks way too hacky, the proposal in bug report proposes to deal with it in `__del__` method. Problem is, the `__del__` method will never be called, exactly because of this issue, because there will always be an internal reference to self. There will be accompanying PR to pytorch to fix it on their side

Fixes #12622

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


cc @otaj @borda @carmocca @justusschock @awaelchli @ananthsub @ninginthecloud @jjenniferdai @rohitgr7

Problem statement:
sys.getrefcount() of a LightningModule
## 🐛 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

Hints from issue discussion:
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 (https://github.com/pytorch/pytorch/pull/62070)

Commit message: Update tests/core/test_lightning_module.py

Co-authored-by: Rohit Gupta <rohitgr1998@gmail.com>

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b84a24b71f3d9..41e434b0f6fa1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -157,6 +157,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Use only a single instance of `rich.console.Console` throughout codebase ([#12886](https://github.com/PyTorchLightning/pytorch-lightning/pull/12886))
 
 
+- Fixed the number of references of `LightningModule` so it can be deleted ([#12897](https://github.com/PyTorchLightning/pytorch-lightning/pull/12897))
+
+
 - Fixed an issue to ensure all the checkpoint states are saved in a common filepath with `DeepspeedStrategy` ([#12887](https://github.com/PyTorchLightning/pytorch-lightning/pull/12887))
 
 
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index a4f343fe68966..2d1a51bdf0753 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -19,6 +19,7 @@
 import numbers
 import os
 import tempfile
+import weakref
 from contextlib import contextmanager
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Mapping, Optional, overload, Sequence, Tuple, Union
@@ -1976,4 +1977,6 @@ def _register_sharded_tensor_state_dict_hooks_if_available(self) -> None:
         from torch.distributed._sharded_tensor import pre_load_state_dict_hook, state_dict_hook
 
         self._register_state_dict_hook(state_dict_hook)
-        self._register_load_state_dict_pre_hook(pre_load_state_dict_hook, True)
+
+        # We need to make sure the self inside the method is a weakref proxy
+        Module._register_load_state_dict_pre_hook(weakref.proxy(self), pre_load_state_dict_hook, True)
diff --git a/tests/core/test_lightning_module.py b/tests/core/test_lightning_module.py
index c7fee3b0d5fd0..822e7161ee7e6 100644
--- a/tests/core/test_lightning_module.py
+++ b/tests/core/test_lightning_module.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import sys
 from unittest.mock import Mock
 
 import pytest
@@ -399,3 +400,10 @@ def configure_gradient_clipping(self, optimizer, optimizer_idx, gradient_clip_va
         match=r"gradient_clip_algorithm='norm'\)` and have passed `clip_gradients\(gradient_clip_algorithm='foo'",
     ):
         trainer.fit(model)
+
+
+def test_proper_refcount():
+    torch_module = nn.Module()
+    lightning_module = LightningModule()
+
+    assert sys.getrefcount(torch_module) == sys.getrefcount(lightning_module)


Code review comments:
diff_hunk: @@ -1976,4 +1977,6 @@ def _register_sharded_tensor_state_dict_hooks_if_available(self) -> None:
         from torch.distributed._sharded_tensor import pre_load_state_dict_hook, state_dict_hook
 
         self._register_state_dict_hook(state_dict_hook)
-        self._register_load_state_dict_pre_hook(pre_load_state_dict_hook, True)
+
+        # We need to make sure the self inside the method is a weakref proxy
path: pytorch_lightning/core/lightning.py
text: @user1:
Let's wait until Torch merges upstream and we have and `if TORCH >= 1.12` to add the patch ourselves

@author:
This has been merged in upstream, although a little bit differently.