Fixes

Author: carmocca

src/pytorch_lightning/lite/lite.py

@@ -38,7 +38,7 @@
 from pytorch_lightning.accelerators import Accelerator as PLAccelerator
 from pytorch_lightning.plugins import DeepSpeedPrecisionPlugin as PLDeepSpeedPrecisionPlugin
 from pytorch_lightning.plugins import DoublePrecisionPlugin as PLDoublePrecisionPlugin
-from pytorch_lightning.plugins import NativeMixedPrecisionPlugin as PLNativeMixedPrecisionPlugin
+from pytorch_lightning.plugins import MixedPrecisionPlugin as PLMixedPrecisionPlugin
 from pytorch_lightning.plugins import PrecisionPlugin as PLPrecisionPlugin
 from pytorch_lightning.plugins import TPUBf16PrecisionPlugin as PLTPUBf16PrecisionPlugin
 from pytorch_lightning.plugins import TPUPrecisionPlugin as PLTPUPrecisionPlugin
@@ -284,7 +284,7 @@ def _to_lite_precision(plugin: Optional[PLPrecisionPlugin]) -> LitePrecision:
     if type(plugin) is PLPrecisionPlugin:
         return LitePrecision()
 
-    if type(plugin) is PLNativeMixedPrecisionPlugin:
+    if type(plugin) is PLMixedPrecisionPlugin:
         return LiteMixedPrecision(
             precision=plugin.precision, device=plugin.device, scaler=plugin.scaler  # type: ignore[arg-type]
         )

src/pytorch_lightning/plugins/precision/__init__.py

@@ -19,7 +19,7 @@
 from pytorch_lightning.plugins.precision.fully_sharded_native_amp import FullyShardedNativeMixedPrecisionPlugin
 from pytorch_lightning.plugins.precision.hpu import HPUPrecisionPlugin
 from pytorch_lightning.plugins.precision.ipu import IPUPrecisionPlugin
-from pytorch_lightning.plugins.precision.native_amp import NativeMixedPrecisionPlugin
+from pytorch_lightning.plugins.precision.native_amp import MixedPrecisionPlugin, NativeMixedPrecisionPlugin
 from pytorch_lightning.plugins.precision.precision_plugin import PrecisionPlugin
 from pytorch_lightning.plugins.precision.sharded_native_amp import ShardedNativeMixedPrecisionPlugin
 from pytorch_lightning.plugins.precision.tpu import TPUPrecisionPlugin
@@ -35,6 +35,7 @@
     "HPUPrecisionPlugin",
     "IPUPrecisionPlugin",
     "NativeMixedPrecisionPlugin",
+    "MixedPrecisionPlugin",
     "PrecisionPlugin",
     "ShardedNativeMixedPrecisionPlugin",
     "TPUPrecisionPlugin",

src/pytorch_lightning/plugins/precision/sharded_native_amp.py

@@ -14,7 +14,7 @@
 from typing import Optional, Union
 
 from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
-from pytorch_lightning.plugins.precision.native_amp import NativeMixedPrecisionPlugin
+from pytorch_lightning.plugins.precision.native_amp import MixedPrecisionPlugin
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 if _FAIRSCALE_AVAILABLE:
@@ -24,7 +24,7 @@
     OSS = ShardedGradScaler = object
 
 
-class ShardedNativeMixedPrecisionPlugin(NativeMixedPrecisionPlugin):
+class ShardedNativeMixedPrecisionPlugin(MixedPrecisionPlugin):
     """Native AMP for Sharded Training."""
 
     def __init__(self, precision: Union[str, int], device: str, scaler: Optional[ShardedGradScaler] = None) -> None:

src/pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -29,7 +29,7 @@
 from lightning_lite.utilities.cloud_io import get_filesystem
 from lightning_lite.utilities.types import _PATH
 from pytorch_lightning.callbacks import ModelCheckpoint
-from pytorch_lightning.plugins.precision import ApexMixedPrecisionPlugin, NativeMixedPrecisionPlugin
+from pytorch_lightning.plugins.precision import ApexMixedPrecisionPlugin, MixedPrecisionPlugin
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities import _OMEGACONF_AVAILABLE
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
@@ -295,9 +295,7 @@ def restore_precision_plugin_state(self) -> None:
         # old checkpoints compatibility
         if "amp_scaling_state" in self._loaded_checkpoint and isinstance(prec_plugin, ApexMixedPrecisionPlugin):
             prec_plugin.load_state_dict(self._loaded_checkpoint["amp_scaling_state"])
-        if "native_amp_scaling_state" in self._loaded_checkpoint and isinstance(
-            prec_plugin, NativeMixedPrecisionPlugin
-        ):
+        if "native_amp_scaling_state" in self._loaded_checkpoint and isinstance(prec_plugin, MixedPrecisionPlugin):
             prec_plugin.load_state_dict(self._loaded_checkpoint["native_amp_scaling_state"])
 
     def _restore_quantization_callbacks(self) -> None:

tests/tests_lite/plugins/precision/test_deepspeed.py

@@ -23,27 +23,11 @@
 
 def test_invalid_precision_with_deepspeed_precision():
     with pytest.raises(ValueError, match="is not supported in DeepSpeed. `precision` must be one of"):
-        DeepSpeedPrecision(precision=64, amp_type="native")
-
-
-def test_deepspeed_precision_apex_not_installed(monkeypatch):
-    import lightning_lite.plugins.precision.deepspeed as deepspeed
-
-    monkeypatch.setattr(deepspeed, "_APEX_AVAILABLE", False)
-    with pytest.raises(ImportError, match="You have asked for Apex AMP but `apex` is not installed."):
-        DeepSpeedPrecision(precision=16, amp_type="apex")
-
-
-@mock.patch("lightning_lite.plugins.precision.deepspeed._APEX_AVAILABLE", return_value=True)
-def test_deepspeed_precision_apex_default_level(_):
-    with pytest.deprecated_call(match="apex AMP implementation has been deprecated"):
-        precision = DeepSpeedPrecision(precision=16, amp_type="apex")
-    assert isinstance(precision, DeepSpeedPrecision)
-    assert precision.amp_level == "O2"
+        DeepSpeedPrecision(precision=64)
 
 
 def test_deepspeed_precision_backward():
-    precision = DeepSpeedPrecision(precision=32, amp_type="native")
+    precision = DeepSpeedPrecision(precision=32)
     tensor = Mock()
     model = Mock()
     precision.backward(tensor, model, "positional-arg", keyword="arg")
@@ -61,7 +45,7 @@ def test_deepspeed_engine_is_steppable(engine):
 
 
 def test_deepspeed_precision_optimizer_step():
-    precision = DeepSpeedPrecision(precision=32, amp_type="native")
+    precision = DeepSpeedPrecision(precision=32)
     optimizer = model = Mock()
     precision.optimizer_step(optimizer, lr_kwargs=dict())
     model.step.assert_called_once_with(lr_kwargs=dict())

tests/tests_pytorch/deprecated_api/test_remove_1-10.py

@@ -355,6 +355,7 @@ def test_profiler_classes_deprecated_warning(cls):
         cls()
 
 
+@RunIf(amp_apex=True)
 def test_apex_deprecation_warnings():
     class MyModel(BoringModel):
         def optimizer_step(

tests/tests_pytorch/models/test_ddp_fork_amp.py

@@ -15,7 +15,7 @@
 
 import torch
 
-from pytorch_lightning.plugins import NativeMixedPrecisionPlugin
+from pytorch_lightning.plugins import MixedPrecisionPlugin
 from tests_pytorch.helpers.runif import RunIf
 
 
@@ -24,7 +24,7 @@
 def test_amp_gpus_ddp_fork():
     """Ensure the use of native AMP with `ddp_fork` (or associated alias strategies) does not generate CUDA
     initialization errors."""
-    _ = NativeMixedPrecisionPlugin(precision=16, device="cuda")
+    _ = MixedPrecisionPlugin(precision=16, device="cuda")
     with multiprocessing.get_context("fork").Pool(1) as pool:
         in_bad_fork = pool.apply(torch.cuda._is_in_bad_fork)
     assert not in_bad_fork

tests/tests_pytorch/models/test_hooks.py

@@ -518,14 +518,15 @@ def training_step(self, batch, batch_idx):
         "state_dict": ANY,
         "loops": ANY,
     }
-    if kwargs.get("precision") == 16:
+    using_deepspeed = kwargs.get("strategy") == "deepspeed"
+    if kwargs.get("precision") == 16 and not using_deepspeed:
         saved_ckpt[trainer.precision_plugin.__class__.__qualname__] = ANY
     device = torch.device("cuda:0" if "accelerator" in kwargs and kwargs["accelerator"] == "gpu" else "cpu")
     expected = [
         dict(name="configure_callbacks"),
         dict(name="prepare_data"),
         # DeepSpeed needs the batch size to figure out throughput logging
-        *([dict(name="train_dataloader")] if kwargs.get("strategy") == "deepspeed" else []),
+        *([dict(name="train_dataloader")] if using_deepspeed else []),
         dict(name="Callback.setup", args=(trainer, model), kwargs=dict(stage="fit")),
         dict(name="setup", kwargs=dict(stage="fit")),
         dict(name="configure_sharded_model"),

tests/tests_pytorch/strategies/test_sharded_strategy.py

@@ -10,7 +10,7 @@
 from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.demos.boring_classes import BoringModel
-from pytorch_lightning.plugins import NativeMixedPrecisionPlugin
+from pytorch_lightning.plugins import MixedPrecisionPlugin
 from pytorch_lightning.strategies import DDPShardedStrategy, DDPSpawnShardedStrategy
 from pytorch_lightning.trainer.states import TrainerFn
 from tests_pytorch.helpers.runif import RunIf
@@ -91,7 +91,7 @@ def test_ddp_choice_sharded_amp(strategy, expected):
     """Test to ensure that plugin native amp plugin is correctly chosen when using sharded."""
     trainer = Trainer(fast_dev_run=True, accelerator="gpu", devices=1, precision=16, strategy=strategy)
     assert isinstance(trainer.strategy, expected)
-    assert isinstance(trainer.precision_plugin, NativeMixedPrecisionPlugin)
+    assert isinstance(trainer.precision_plugin, MixedPrecisionPlugin)
 
 
 @RunIf(fairscale=True)

tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py

@@ -432,9 +432,9 @@ def test_validate_precision_type(precision):
 
 
 def test_amp_level_raises_error_with_native():
-    with pytest.deprecated_call(
-        match="Setting `amp_level` inside the `Trainer` is deprecated in v1.8.0"
-    ), pytest.raises(MisconfigurationException, match="O2'` but it's only supported with `amp_backend='apex'`"):
+    with pytest.deprecated_call(match="apex AMP implementation has been deprecated"), pytest.raises(
+        MisconfigurationException, match="O2'` but it's only supported with `amp_backend='apex'`"
+    ):
         _ = Trainer(amp_level="O2", amp_backend="native", precision=16)