Remove the deprecated LightningDeepSpeedModule (#16041)

src/pytorch_lightning/CHANGELOG.md

@@ -89,6 +89,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Removed deprecated code in `pytorch_lightning.utilities.meta` ([#16038](https://github.com/Lightning-AI/lightning/pull/16038))
 
 
+- Removed the deprecated `LightningDeepSpeedModule` ([#16041](https://github.com/Lightning-AI/lightning/pull/16041))
+
+
 ### Fixed
 
 - Enhanced `reduce_boolean_decision` to accommodate `any`-analogous semantics expected by the `EarlyStopping` callback ([#15253](https://github.com/Lightning-AI/lightning/pull/15253))

src/pytorch_lightning/_graveyard/__init__.py

@@ -17,6 +17,7 @@
 import pytorch_lightning._graveyard.legacy_import_unpickler
 import pytorch_lightning._graveyard.loggers
 import pytorch_lightning._graveyard.profiler
+import pytorch_lightning._graveyard.strategies
 import pytorch_lightning._graveyard.trainer
 import pytorch_lightning._graveyard.training_type
 import pytorch_lightning._graveyard.utilities  # noqa: F401

src/pytorch_lightning/_graveyard/strategies.py

@@ -0,0 +1,15 @@
+from typing import Any
+
+import pytorch_lightning.strategies.deepspeed as deepspeed
+
+
+class _LightningDeepSpeedModule:
+    # TODO: Remove in v2.0.0
+    def __init__(self, *_: Any, **__: Any) -> None:
+        raise RuntimeError(
+            "`pytorch_lightning.strategies.deepspeed.LightningDeepSpeedModule` was deprecated in v1.7.1 and is no"
+            " longer supported as of v1.9.0."
+        )
+
+
+deepspeed.LightningDeepSpeedModule = _LightningDeepSpeedModule

src/pytorch_lightning/strategies/deepspeed.py

@@ -45,7 +45,7 @@
 from pytorch_lightning.utilities import GradClipAlgorithmType
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.model_helpers import is_overridden
-from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation, rank_zero_info, rank_zero_only, rank_zero_warn
+from pytorch_lightning.utilities.rank_zero import rank_zero_info, rank_zero_only, rank_zero_warn
 from pytorch_lightning.utilities.types import LRSchedulerConfig, STEP_OUTPUT
 
 log = logging.getLogger(__name__)
@@ -67,28 +67,6 @@ def remove_module_hooks(model: torch.nn.Module) -> None:
         module._load_state_dict_pre_hooks = OrderedDict()
 
 
-class LightningDeepSpeedModule(_LightningModuleWrapperBase):
-    """
-    .. deprecated:: v1.7.1
-        ``LightningDeepSpeedModule`` has been deprecated in v1.7.1 and will be removed in v1.9.0.
-    """
-
-    def __init__(
-        self,
-        forward_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-        precision: Union[str, int] = 32,
-        pl_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-    ) -> None:
-        rank_zero_deprecation("`LightningDeepSpeedModule` has been deprecated in v1.7.1 and will be removed in v1.9.0")
-        self._validate_init_arguments(pl_module, forward_module)
-        super().__init__(forward_module=(pl_module or forward_module))
-        self.precision = precision
-
-    def forward(self, *inputs: Any, **kwargs: Any) -> Any:
-        inputs = apply_to_collection(inputs, Tensor, function=_fp_to_half, precision=self.precision)
-        return super().forward(*inputs, **kwargs)
-
-
 class DeepSpeedStrategy(DDPStrategy):
     strategy_name = "deepspeed"
     DEEPSPEED_ENV_VAR = "PL_DEEPSPEED_CONFIG_PATH"

tests/tests_pytorch/deprecated_api/test_remove_1-9.py

@@ -30,7 +30,6 @@
 from pytorch_lightning.profiler.pytorch import PyTorchProfiler, RegisterRecordFunction, ScheduleWrapper
 from pytorch_lightning.profiler.simple import SimpleProfiler
 from pytorch_lightning.profiler.xla import XLAProfiler
-from pytorch_lightning.strategies.deepspeed import LightningDeepSpeedModule
 from pytorch_lightning.utilities.imports import _KINETO_AVAILABLE
 from pytorch_lightning.utilities.rank_zero import rank_zero_only
 from tests_pytorch.helpers.runif import RunIf
@@ -218,8 +217,3 @@ def test_gpu_accelerator_deprecation_warning():
         )
     ):
         GPUAccelerator()
-
-
-def test_v1_9_0_deprecated_lightning_deepspeed_module():
-    with pytest.deprecated_call(match=r"has been deprecated in v1.7.1 and will be removed in v1.9."):
-        _ = LightningDeepSpeedModule(BoringModel(), 32)

tests/tests_pytorch/graveyard/test_strategies.py

@@ -0,0 +1,8 @@
+import pytest
+
+
+def test_v2_0_0_lightningdeepspeedmodule():
+    from pytorch_lightning.strategies.deepspeed import LightningDeepSpeedModule
+
+    with pytest.raises(RuntimeError, match="LightningDeepSpeedModule` was deprecated in v1.7.1"):
+        LightningDeepSpeedModule()

tests/tests_pytorch/strategies/test_deepspeed_strategy.py

@@ -32,7 +32,7 @@
 from pytorch_lightning.loggers import CSVLogger
 from pytorch_lightning.plugins import DeepSpeedPrecisionPlugin
 from pytorch_lightning.strategies import DeepSpeedStrategy
-from pytorch_lightning.strategies.deepspeed import _DEEPSPEED_AVAILABLE, LightningDeepSpeedModule
+from pytorch_lightning.strategies.deepspeed import _DEEPSPEED_AVAILABLE
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests_pytorch.helpers.datamodules import ClassifDataModule
 from tests_pytorch.helpers.runif import RunIf
@@ -84,27 +84,6 @@ def automatic_optimization(self) -> bool:
         return False
 
 
-@RunIf(min_cuda_gpus=1)
-def test_deepspeed_lightning_module_precision():
-    """Test to ensure that a model wrapped in `LightningDeepSpeedModule` moves tensors to half when precision
-    16."""
-    model = BoringModel()
-    with pytest.deprecated_call(match="`LightningDeepSpeedModule` has been deprecated in v1.7.1"):
-        module = LightningDeepSpeedModule(model, precision=16)
-
-    module.to(device="cuda", dtype=torch.half)
-    assert module.dtype == torch.half
-    assert model.dtype == torch.half
-
-    x = torch.randn((1, 32), device="cuda", dtype=torch.float)
-    out = module(x)
-    assert out.dtype == torch.half
-
-    module.to(torch.double)
-    assert module.dtype == torch.double
-    assert model.dtype == torch.double
-
-
 @pytest.fixture
 def deepspeed_config():
     return {