Merge branch 'lite/debug' into lite/debug-distributed

Author: carmocca

src/pytorch_lightning/CHANGELOG.md

@@ -51,6 +51,13 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Removed the deprecated code in `pl.utilities.distributed` ([#16390](https://github.com/Lightning-AI/lightning/pull/16390))
 
+- Mark the `forward_module` argument as required ([#16386](https://github.com/Lightning-AI/lightning/pull/16386))
+  * Removed the deprecated `pl_module` argument from the distributed module wrappers
+  * Removed the deprecated `pytorch_lightning.overrides.base.unwrap_lightning_module` function
+  * Removed the `pytorch_lightning.overrides.distributed.LightningDistributedModule` class
+  * Removed the deprecated `pytorch_lightning.overrides.fairscale.unwrap_lightning_module_sharded` function
+  * Removed the `pytorch_lightning.overrides.fairscale.LightningDistributedModule` class
+
 - Removed the deprecated automatic GPU selection ([#16184](https://github.com/Lightning-AI/lightning/pull/16184))
   * Removed the `Trainer(auto_select_gpus=...)` argument
   * Removed the `pytorch_lightning.tuner.auto_gpu_select.{pick_single_gpu,pick_multiple_gpus}` functions

src/pytorch_lightning/overrides/__init__.py

@@ -1,2 +0,0 @@
-from pytorch_lightning.overrides.data_parallel import LightningParallelModule  # noqa: F401
-from pytorch_lightning.overrides.distributed import LightningDistributedModule  # noqa: F401

src/pytorch_lightning/overrides/base.py

@@ -11,16 +11,12 @@
 # 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.
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 import torch
-import torch.nn as nn
-from torch.nn import DataParallel
-from torch.nn.parallel import DistributedDataParallel
 
 import pytorch_lightning as pl
 from lightning_fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
-from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation
 
 
 class _LightningPrecisionModuleWrapperBase(_DeviceDtypeModuleMixin, torch.nn.Module):
@@ -55,9 +51,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
 
 
 class _LightningModuleWrapperBase(_DeviceDtypeModuleMixin, torch.nn.Module):
-    def __init__(
-        self, forward_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]]
-    ) -> None:
+    def __init__(self, forward_module: Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]) -> None:
         """Wraps the user's LightningModule and redirects the forward call to the appropriate method, either
         ``training_step``, ``validation_step``, ``test_step``, or ``predict_step``.
 
@@ -75,8 +69,6 @@ def __init__(
                 "`forward_module` must be a `LightningModule` instance or have an attribute `.module` pointing to one,"
                 f" got: {forward_module.__class__.__qualname__}"
             )
-        # TODO: In v2.0.0, remove the Optional type from forward_module and remove the assertion
-        assert forward_module is not None
         self._forward_module = forward_module
 
         # set the parameters_to_ignore from LightningModule.
@@ -111,47 +103,3 @@ def forward(self, *inputs: Any, **kwargs: Any) -> Any:
             if trainer.predicting:
                 return self._forward_module.predict_step(*inputs, **kwargs)
         return self._forward_module(*inputs, **kwargs)
-
-    @classmethod
-    def _validate_init_arguments(
-        cls,
-        pl_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-        forward_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-    ) -> None:
-        # TODO: In v2.0.0, remove this method and mark the forward_module init argument in all subclasses as required
-        if pl_module is not None:
-            rank_zero_deprecation(
-                f"The argument `pl_module` in `{cls.__name__}` is deprecated in v1.8.0 and will be removed in"
-                " v2.0.0. Please use `forward_module` instead."
-            )
-        elif forward_module is None:
-            raise ValueError("Argument `forward_module` is required.")
-
-
-def unwrap_lightning_module(wrapped_model: nn.Module, _suppress_warning: bool = False) -> "pl.LightningModule":
-    """Recursively unwraps a :class:`~pytorch_lightning.core.module.LightningModule` by following the ``.module``
-    attributes on the wrapper.
-
-    .. deprecated:: v1.8.0
-        The function ``unwrap_lightning_module`` is deprecated in v1.8.0 and will be removed in v2.0.0. Access the
-        ``LightningModule`` directly through the strategy attribute ``Strategy.lightning_module``.
-
-    Raises:
-        TypeError: If the unwrapping leads to a module that is not a LightningModule and that cannot be unwrapped
-            further.
-    """
-    if not _suppress_warning:
-        rank_zero_deprecation(
-            "The function `unwrap_lightning_module` is deprecated in v1.8.0 and will be removed in v2.0.0. Access the"
-            " `LightningModule` directly through the strategy attribute `Strategy.lightning_module`."
-        )
-    model = wrapped_model
-    if isinstance(model, (DistributedDataParallel, DataParallel)):
-        model = unwrap_lightning_module(model.module)
-    if isinstance(model, _LightningModuleWrapperBase):
-        model = model.lightning_module
-    if isinstance(model, _LightningPrecisionModuleWrapperBase):
-        model = model.module
-    if not isinstance(model, pl.LightningModule):
-        raise TypeError(f"Unwrapping the module did not yield a `LightningModule`, got {type(model)} instead.")
-    return model

src/pytorch_lightning/overrides/data_parallel.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import numbers
 import warnings
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 import torch
 from lightning_utilities.core.apply_func import apply_to_collection
@@ -52,23 +52,15 @@ class LightningParallelModule(_LightningModuleWrapperBase):
         )
 
     Args:
-        pl_module: The module to wrap. See description for `forward_module`.
-
-            .. deprecated:: v1.8.0
-                The argument ``pl_module`` is deprecated in v1.8.0 and will be removed in v2.0.0. Please use
-                ``forward_module`` instead.
-
         forward_module: The module to wrap. If it's not a ``LightningModule``, it must have an attribute ``.module``
             pointing to a ``LightningModule`` reference.
     """
 
     def __init__(
         self,
-        forward_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-        pl_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
+        forward_module: Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase],
     ) -> None:
-        self._validate_init_arguments(pl_module, forward_module)
-        super().__init__(forward_module=(pl_module or forward_module))
+        super().__init__(forward_module=forward_module)
         _ignore_scalar_return_in_dp()
 
     def forward(self, *inputs: Any, **kwargs: Any) -> Any:

src/pytorch_lightning/overrides/distributed.py

@@ -12,26 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import itertools
-from typing import Any, cast, Iterable, Iterator, List, Optional, Sized, Union
+from typing import Any, cast, Iterable, Iterator, List, Sized, Union
 
 import torch
 from torch import Tensor
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import BatchSampler, DistributedSampler, Sampler
 
-import pytorch_lightning as pl
 from lightning_fabric.utilities.distributed import _DatasetSamplerWrapper
-from pytorch_lightning.overrides.base import _LightningModuleWrapperBase, _LightningPrecisionModuleWrapperBase
-
-
-class LightningDistributedModule(_LightningModuleWrapperBase):
-    def __init__(
-        self,
-        forward_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-        pl_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-    ) -> None:
-        self._validate_init_arguments(pl_module, forward_module)
-        super().__init__(forward_module=(pl_module or forward_module))
 
 
 def _find_tensors(

src/pytorch_lightning/overrides/fairscale.py

@@ -11,21 +11,13 @@
 # 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.
-from typing import List, Optional, Union
+from typing import List
 
-import torch.nn as nn
 from lightning_utilities.core.imports import package_available
 from torch.optim import Optimizer
 
-import pytorch_lightning as pl
 from lightning_fabric.plugins import Precision
 from lightning_fabric.utilities.imports import _IS_WINDOWS
-from pytorch_lightning.overrides.base import (
-    _LightningModuleWrapperBase,
-    _LightningPrecisionModuleWrapperBase,
-    unwrap_lightning_module,
-)
-from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation
 
 _FAIRSCALE_AVAILABLE = not _IS_WINDOWS and package_available("fairscale")
 
@@ -35,30 +27,6 @@
     OSS = object
 
 
-class LightningShardedDataParallel(_LightningModuleWrapperBase):
-    def __init__(
-        self,
-        forward_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-        pl_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-    ) -> None:
-        self._validate_init_arguments(pl_module, forward_module)
-        super().__init__(forward_module=(pl_module or forward_module))
-
-
-def unwrap_lightning_module_sharded(wrapped_model: nn.Module) -> "pl.LightningModule":
-    from fairscale.nn.data_parallel.sharded_ddp import ShardedDataParallel
-
-    rank_zero_deprecation(
-        "The function `unwrap_lightning_module_sharded` is deprecated in v1.8.0 and will be removed in v2.0.0."
-        " Access the `LightningModule` directly through the strategy attribute `Strategy.lightning_module`."
-    )
-    model = wrapped_model
-    if isinstance(model, ShardedDataParallel):
-        model = model.module
-
-    return unwrap_lightning_module(model, _suppress_warning=True)
-
-
 def _reinit_optimizers_with_oss(optimizers: List[Optimizer], precision: Precision, num_nodes: int) -> List["OSS"]:
     for x, optimizer in enumerate(optimizers):
         if not isinstance(optimizer, OSS):

src/pytorch_lightning/strategies/bagua.py

@@ -64,11 +64,8 @@
 class LightningBaguaModule(_LightningModuleWrapperBase):
     def __init__(
         self,
-        forward_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-        pl_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
+        forward_module: Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase],
     ) -> None:
-        self._validate_init_arguments(pl_module, forward_module)
-        forward_module = pl_module or forward_module
         super().__init__(forward_module=forward_module)
         # Bagua use `bagua_module_name` to distinguish different modules
         self._bagua_module_name = f"{forward_module.__class__.__name__}{id(forward_module)}"

src/pytorch_lightning/strategies/ddp.py

@@ -43,8 +43,7 @@
 from lightning_fabric.utilities.seed import reset_seed
 from lightning_fabric.utilities.types import ReduceOp
 from pytorch_lightning.core.optimizer import LightningOptimizer
-from pytorch_lightning.overrides import LightningDistributedModule
-from pytorch_lightning.overrides.base import _LightningPrecisionModuleWrapperBase
+from pytorch_lightning.overrides.base import _LightningModuleWrapperBase, _LightningPrecisionModuleWrapperBase
 from pytorch_lightning.overrides.distributed import prepare_for_backward
 from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.plugins.precision import PrecisionPlugin
@@ -291,7 +290,7 @@ def configure_ddp(self) -> None:
         log.detail(f"{self.__class__.__name__}: configuring DistributedDataParallel")
         self.pre_configure_ddp()
         assert isinstance(self.model, (pl.LightningModule, _LightningPrecisionModuleWrapperBase))
-        self.model = self._setup_model(LightningDistributedModule(self.model))
+        self.model = self._setup_model(_LightningModuleWrapperBase(self.model))
         self._register_ddp_hooks()
 
     def determine_ddp_device_ids(self) -> Optional[List[int]]:

src/pytorch_lightning/strategies/ddp_spawn.py

@@ -36,8 +36,7 @@
 from lightning_fabric.utilities.imports import _TORCH_GREATER_EQUAL_1_11
 from lightning_fabric.utilities.optimizer import _optimizers_to_device
 from lightning_fabric.utilities.types import ReduceOp
-from pytorch_lightning.overrides import LightningDistributedModule
-from pytorch_lightning.overrides.base import _LightningPrecisionModuleWrapperBase
+from pytorch_lightning.overrides.base import _LightningModuleWrapperBase, _LightningPrecisionModuleWrapperBase
 from pytorch_lightning.overrides.distributed import prepare_for_backward
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.launchers.multiprocessing import _MultiProcessingLauncher
@@ -211,7 +210,7 @@ def _register_ddp_hooks(self) -> None:
     def configure_ddp(self) -> None:
         self.pre_configure_ddp()
         assert isinstance(self.model, (pl.LightningModule, _LightningPrecisionModuleWrapperBase))
-        self.model = self._setup_model(LightningDistributedModule(self.model))
+        self.model = self._setup_model(_LightningModuleWrapperBase(self.model))
         self._register_ddp_hooks()
 
         # set up optimizers after the wrapped module has been moved to the device

src/pytorch_lightning/strategies/hpu_parallel.py

@@ -22,7 +22,7 @@
 import pytorch_lightning as pl
 from lightning_fabric.plugins import CheckpointIO, ClusterEnvironment
 from lightning_fabric.utilities.distributed import group as _group
-from pytorch_lightning.overrides import LightningDistributedModule
+from pytorch_lightning.overrides.base import _LightningModuleWrapperBase
 from pytorch_lightning.overrides.torch_distributed import broadcast_object_list
 from pytorch_lightning.plugins.io.hpu_plugin import HPUCheckpointIO
 from pytorch_lightning.plugins.io.wrapper import _WrappingCheckpointIO
@@ -123,7 +123,7 @@ def configure_ddp(self) -> None:
         if _TORCH_LESSER_EQUAL_1_10_2:
             log.detail(f"{self.__class__.__name__}: configuring DistributedDataParallel")
             self._pre_configure_ddp()
-            self.model = self._setup_model(LightningDistributedModule(self.model))  # type: ignore
+            self.model = self._setup_model(_LightningModuleWrapperBase(self.model))  # type: ignore
             if self.root_device.type == "hpu" and self._static_graph:
                 self._model._set_static_graph()  # type: ignore
             self._register_ddp_hooks()