DeprecationWarning: torch.distributed._sharded_tensor #13256
Description
🐛 Bug
The latest versions of pytorch-lightning and PyTorch give the following deprecation warning:
/opt/hostedtoolcache/Python/3.10.4/x64/lib/python3.10/site-packages/pytorch_lightning/core/lightning.py:118: in __init__
self._register_sharded_tensor_state_dict_hooks_if_available()
/opt/hostedtoolcache/Python/3.10.4/x64/lib/python3.10/site-packages/pytorch_lightning/core/lightning.py:2067: in _register_sharded_tensor_state_dict_hooks_if_available
from torch.distributed._sharded_tensor import pre_load_state_dict_hook, state_dict_hook
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
# Keep old package for BC purposes, this file should be removed once
# everything moves to the `torch.distributed._shard` package.
import sys
import torch
import warnings
from torch.distributed._shard.sharded_tensor import * # noqa: F403
> warnings.warn(
"torch.distributed._sharded_tensor will be deprecated, use torch.distributed._shard.sharded_tensor instead",
DeprecationWarning
)
E DeprecationWarning: torch.distributed._sharded_tensor will be deprecated, use torch.distributed._shard.sharded_tensor instead
This chunk of code has since moved, but torch.distributed._sharded_tensor is still being imported in pytorch_lightning/core/module.py and tests/core/test_lightning_module.py.
To Reproduce
Run code using LightningModule through pytest with the latest version of PyTorch.
Expected behavior
No DeprecationWarning.
Environment
- PyTorch Lightning Version (e.g., 1.5.0): 1.6.4
- PyTorch Version (e.g., 1.10): 1.11.0
- Python version (e.g., 3.9): 3.10.4, etc.
- OS (e.g., Linux): Linux, etc.
- CUDA/cuDNN version: N/A
- GPU models and configuration: N/A
- How you installed PyTorch (
conda,pip, source): pip - If compiling from source, the output of
torch.__config__.show(): N/A - Any other relevant information: N/A