Repository: Lightning-AI/pytorch-lightning
Pull request: Use base version when comparing torch versions
## What does this PR do?

Fixes #16644

## 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)
- [ ] Did you write any new necessary tests? (not for typos and docs)
- [ ] Did you verify new and existing tests pass locally with your changes?
- [x] Did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [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?
I made sure I had fun coding 🙃


cc @borda @tchaton @carmocca @justusschock @awaelchli

Problem statement:
TypeError: cannot create weak reference to 'weakcallableproxy' object after upgrading to 1.9.0
### Bug description

With the new pip package pytorch-lightning==1.9.0 I am unable to instantiate any PL module:

```python
import pytorch_lightning as pl
pl.LightningModule()
```

getting  error: `TypeError: cannot create weak reference to 'weakcallableproxy' object`

Downgrading to 1.8.6 helps.

### How to reproduce the bug

```python
!pip install pytorch-lightning==1.9.0
import pytorch_lightning as pl
pl.LightningModule()
```


### Error messages and logs

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In [2], line 1
----> 1 pl.LightningModule()

File /opt/conda/lib/python3.8/site-packages/pytorch_lightning/core/module.py:124, in LightningModule.__init__(self, *args, **kwargs)
    122 self._metric_attributes: Optional[Dict[int, str]] = None
    123 self._should_prevent_trainer_and_dataloaders_deepcopy: bool = False
--> 124 self._register_sharded_tensor_state_dict_hooks_if_available()
    125 self._compiler_ctx: Optional[Dict[str, Any]] = None
    127 # attributes only used when using fabric

File /opt/conda/lib/python3.8/site-packages/pytorch_lightning/core/module.py:2022, in LightningModule._register_sharded_tensor_state_dict_hooks_if_available(self)
   2019     self._register_load_state_dict_pre_hook(pre_load_state_dict_hook, True)
   2020 else:
   2021     # We need to make sure the self inside the method is a weakref proxy
-> 2022     self.__class__._register_load_state_dict_pre_hook(
   2023         weakref.proxy(self), pre_load_state_dict_hook, True  # type: ignore[arg-type]
   2024     )

File /opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1462, in Module._register_load_state_dict_pre_hook(self, hook, with_module)
   1447 r"""These hooks will be called with arguments: `state_dict`, `prefix`,
   1448 `local_metadata`, `strict`, `missing_keys`, `unexpected_keys`,
   1449 `error_msgs`, before loading `state_dict` into `self`. These arguments
   (...)
   1459         instance to the hook as the first parameter.
   1460 """
   1461 handle = hooks.RemovableHandle(self._load_state_dict_pre_hooks)
-> 1462 self._load_state_dict_pre_hooks[handle.id] = _WrappedHook(hook, self if with_module else None)
   1463 return handle

File /opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:52, in _WrappedHook.__init__(self, hook, module)
     49 self.with_module: bool = False
     51 if module is not None:
---> 52     self.module: weakref.ReferenceType["Module"] = weakref.ref(module)
     53     self.with_module = True

TypeError: cannot create weak reference to 'weakcallableproxy' object```
```

### Environment

<details>
  <summary>Current environment</summary>

```
* CUDA:
	- available:         True
	- version:           11.8
* Lightning:
	- functorch:         0.3.0a0
	- lightning-utilities: 0.6.0.post0
	- pytorch-lightning: 1.9.0
	- pytorch-quantization: 2.1.2
	- torch:             1.13.0a0+d0d6b1f
	- torch-tensorrt:    1.3.0a0
	- torchmetrics:      0.11.1
	- torchtext:         0.11.0a0
	- torchvision:       0.14.0a0
* Packages:
	- apex:              0.1
	- cffi:              1.15.1
	- conda:             22.9.0
	- conda-build:       3.22.0
	- conda-package-handling: 1.9.0
	- csvw:              3.1.3
	- cuda-python:       11.7.0+0.g95a2041.dirty
	- cudf:              22.8.0a0+304.g6ca81bbc78.dirty
	- cugraph:           22.8.0a0+132.g2daa31b6.dirty
	- cuml:              22.8.0a0+52.g73b8d00d0.dirty
	- cupy-cuda118:      11.0.0
	- cython:            0.29.21
	- functorch:         0.3.0a0
	- future:            0.18.2
	- ipykernel:         6.16.0
	- ipython:           8.5.0
	- ipython-genutils:  0.2.0
	- joblib:            1.2.0
	- keras:             2.9.0
	- keras-preprocessing: 1.1.2
	- libclang:          15.0.6.1
	- lightning-utilities: 0.6.0.post0
	- llvmlite:          0.39.1
	- numba:             0.56.4
	- numpy:             1.23.5
	- nvidia-dali-cuda110: 1.18.0
	- nvidia-pyindex:    1.0.9
	- nvtx:              0.2.5
	- onnx:              1.12.0
	- opt-einsum:        3.3.0
	- pip:               21.2.4
	- pkginfo:           1.8.3
	- pkgutil-resolve-name: 1.3.10
	- py:                1.11.0
	- pyarrow:           8.0.0
	- pynvml:            11.4.1
	- pytorch-lightning: 1.9.0
	- pytorch-quantization: 2.1.2
	- tensorboard:       2.9.1
	- tensorboard-data-server: 0.6.1
	- tensorboard-plugin-wit: 1.8.1
	- tensorboardx:      2.5.1
	- tensorflow:        2.9.1
	- tensorflow-addons: 0.19.0
	- tensorflow-estimator: 2.9.0
	- tensorflow-io-gcs-filesystem: 0.30.0
	- tensorflow-probability: 0.17.0
	- tensorrt:          8.5.0.12
	- torch:             1.13.0a0+d0d6b1f
	- torch-tensorrt:    1.3.0a0
	- torchmetrics:      0.11.1
	- torchtext:         0.11.0a0
	- torchvision:       0.14.0a0
* System:
	- OS:                Linux
	- architecture:
		- 64bit
		- ELF
	- processor:         x86_64
	- python:            3.8.13
	- version:           70-Ubuntu SMP Tue Jan 12 12:45:47 UTC 2021
```

</details>


### More info

I am using docker container based on this image: nvcr.io/nvidia/pytorch:22.10-py3

cc @tchaton @carmocca @justusschock @awaelchli @borda

Hints from issue discussion:
@radek555 This is reproducible. Thanks for pointing to the docker image.
I found out what's causing this issue. We run the following code conditionally:

https://github.com/Lightning-AI/lightning/blob/accd2b9e61063ba3c683764043030545ed87c71f/src/lightning/pytorch/core/module.py#L1635-L1641

You are using torch version `1.13.0a0+d0d6b1f` in that docker image. This is a torch version compiled from a specific commit of the pytorch repo. Our version comparison results in `_TORCH_GREATER_EQUAL_1_13` being False and in the above code, the wrong branch is taken. 

This is because the comparison `1.13.0a0+d0d6b1f >= 1.13` does not hold. To fix this, we would need to make our [compare_version](https://github.com/Lightning-AI/utilities/blob/d8f0b17e9ce64ec81caa514b3cd9d5301414c13a/src/lightning_utilities/core/imports.py#L60-L83) function work with version numbers that include commit tags.

To work around this issue, I suggest to "upgrade" your pytorch version to a stable version.
Thanks for the quick resolution @awaelchli !
I added some info to a relevant stackoverflow question: https://stackoverflow.com/questions/75356395/cannot-create-weak-reference-to-weakcallableproxy-object-in-pytorch-module/
@radek555 it looks like we can't just easily fix this directly in Lightning. We can't reliably have conditional logic based on unofficial/nightly PyTorch releases like the one used in your docker container. The PyTorch version in your container is `1.13.0a0+d0d6b1f` which is actually earlier than the stable 1.13 release. The best we could do is warn or error out if someone is trying to run with an unsupported version like this one. 

Do you know why this container has this particular version of PyTorch instead of the stable one? This is very odd IMO. 
I see. I was using this pytorch version, because that's the preinstalled one in the nvidia pytorch container, see here 
 https://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.html

Honestly, I have no clue why are they releasing images witch such a particular version. If it is not yet the 1.13 version, they should call it rather "1.12.0+d0d6b1f" IMO.
Strangely, I get this error also for pl 1.8.6 using torch '1.12.0.dev20220517+cu113'. Is this odd?
(Upgrading to torch 1.13 solves the issue on any version of pl).

Commit message: changelog

Patch to review:
diff --git a/src/lightning/fabric/utilities/imports.py b/src/lightning/fabric/utilities/imports.py
index b5b1441781640..85be277f35bb9 100644
--- a/src/lightning/fabric/utilities/imports.py
+++ b/src/lightning/fabric/utilities/imports.py
@@ -25,7 +25,10 @@
 # 2. The inspection mode via `python -i`: https://stackoverflow.com/a/6879085/1162383
 _IS_INTERACTIVE = hasattr(sys, "ps1") or bool(sys.flags.interactive)
 
-_TORCH_GREATER_EQUAL_1_12 = compare_version("torch", operator.ge, "1.12.0")
-_TORCH_GREATER_EQUAL_1_13 = compare_version("torch", operator.ge, "1.13.0")
+# We use "base_version" for non-nightly builds as well, because some environments like NVIDIA's PyTorch dockers
+# install PyTorch from source at a commit that doesn't align with the released version tag.
+# See: https://github.com/Lightning-AI/lightning/issues/16644
+_TORCH_GREATER_EQUAL_1_12 = compare_version("torch", operator.ge, "1.12.0", use_base_version=True)
+_TORCH_GREATER_EQUAL_1_13 = compare_version("torch", operator.ge, "1.13.0", use_base_version=True)
 _TORCH_GREATER_EQUAL_2_0 = compare_version("torch", operator.ge, "2.0.0", use_base_version=True)
 _TORCH_GREATER_EQUAL_2_1 = compare_version("torch", operator.ge, "2.1.0", use_base_version=True)
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
index 6aafd01fe9a29..c878f1d845a35 100644
--- a/src/lightning/pytorch/CHANGELOG.md
+++ b/src/lightning/pytorch/CHANGELOG.md
@@ -396,6 +396,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed an issue where `DistributedSampler.set_epoch` wasn't getting called during `trainer.predict` ([#16785](https://github.com/Lightning-AI/lightning/pull/16785), [#16826](https://github.com/Lightning-AI/lightning/pull/16826))
 
 
+- Fixed an issue with comparing torch versions when using a version of torch built from source ([#16657](https://github.com/Lightning-AI/lightning/pull/16657))
+
+
 ## [1.9.4] - 2023-03-01
 
 ### Added
diff --git a/src/lightning/pytorch/utilities/imports.py b/src/lightning/pytorch/utilities/imports.py
index 4a695d3b56fd2..1d606aa58df91 100644
--- a/src/lightning/pytorch/utilities/imports.py
+++ b/src/lightning/pytorch/utilities/imports.py
@@ -22,7 +22,7 @@
 _PYTHON_GREATER_EQUAL_3_10_0 = (sys.version_info.major, sys.version_info.minor) >= (3, 10)
 _PYTHON_GREATER_EQUAL_3_11_0 = (sys.version_info.major, sys.version_info.minor) >= (3, 11)
 # duplicated from fabric because HPU is patching it below
-_TORCH_GREATER_EQUAL_1_13 = compare_version("torch", operator.ge, "1.13.0")
+_TORCH_GREATER_EQUAL_1_13 = compare_version("torch", operator.ge, "1.13.0", use_base_version=True)
 _TORCHMETRICS_GREATER_EQUAL_0_9_1 = RequirementCache("torchmetrics>=0.9.1")
 _TORCHMETRICS_GREATER_EQUAL_0_11 = compare_version("torchmetrics", operator.ge, "0.11.0")  # using new API with task
 


Code review comments:
diff_hunk: @@ -22,7 +22,7 @@
 _PYTHON_GREATER_EQUAL_3_10_0 = (sys.version_info.major, sys.version_info.minor) >= (3, 10)
 _PYTHON_GREATER_EQUAL_3_11_0 = (sys.version_info.major, sys.version_info.minor) >= (3, 11)
 # duplicated from fabric because HPU is patching it below
path: src/lightning/pytorch/utilities/imports.py
text: @user1:
This no longer applies

```suggestion
```