Fix tests

Author: carmocca

src/pytorch_lightning/plugins/precision/apex_amp.py

@@ -32,10 +32,10 @@
 @functools.lru_cache(maxsize=1)
 def _import_amp_without_deprecation() -> ModuleType:
     # hide the warning upstream in favor of our deprecation
-    with warnings.filterwarnings(action="ignore", message="apex.amp is deprecated", category=FutureWarning):
-        from apex import amp
+    warnings.filterwarnings(action="ignore", message="apex.amp is deprecated", category=FutureWarning)
+    from apex import amp
 
-        return amp
+    return amp
 
 
 # TODO: remove in v1.10.0

src/pytorch_lightning/plugins/precision/deepspeed.py

@@ -66,8 +66,8 @@ def __init__(self, precision: Union[str, int], amp_type: str = "native", amp_lev
                     f"`{type(self).__name__}(amp_level={amp_level!r})` is only relevant when using NVIDIA/apex"
                 )
             rank_zero_deprecation(
-                f"Passing `{type(self).__name__}(amp_type=...)` been deprecated in v1.9.0 and will be removed in"
-                " v1.10.0. This argument is no longer necessary."
+                f"Passing `{type(self).__name__}(amp_type={amp_type!r})` been deprecated in v1.9.0 and will be removed"
+                f" in v1.10.0. This argument is no longer necessary."
             )
 
         supported_precision = (PrecisionType.HALF, PrecisionType.FLOAT, PrecisionType.BFLOAT)

src/pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -48,7 +48,7 @@
     FullyShardedNativeMixedPrecisionPlugin,
     HPUPrecisionPlugin,
     IPUPrecisionPlugin,
-    NativeMixedPrecisionPlugin,
+    MixedPrecisionPlugin,
     PLUGIN_INPUT,
     PrecisionPlugin,
     ShardedNativeMixedPrecisionPlugin,
@@ -717,7 +717,7 @@ def _check_and_init_precision(self) -> PrecisionPlugin:
 
         if self._precision_flag in (16, "bf16"):
             rank_zero_info(
-                f"Using 16bit {self._amp_type_flag.value} Automatic Mixed Precision (AMP)"  # type: ignore
+                f"Using 16bit {self._amp_type_flag} Automatic Mixed Precision (AMP)"  # type: ignore
                 if self._precision_flag == 16
                 else "Using bfloat16 Automatic Mixed Precision (AMP)"
             )
@@ -731,7 +731,7 @@ def _check_and_init_precision(self) -> PrecisionPlugin:
                     return FullyShardedNativeNativeMixedPrecisionPlugin(self._precision_flag, device)
                 if isinstance(self.strategy, DDPFullyShardedStrategy):
                     return FullyShardedNativeMixedPrecisionPlugin(self._precision_flag, device)
-                return NativeMixedPrecisionPlugin(self._precision_flag, device)
+                return MixedPrecisionPlugin(self._precision_flag, device)
 
             if self._amp_type_flag == "apex":
                 self._amp_level_flag = self._amp_level_flag or "O2"
@@ -771,7 +771,7 @@ def _validate_precision_choice(self) -> None:
             )
         if self._precision_flag == "bf16" and self._amp_type_flag != "native":
             raise MisconfigurationException(
-                f"You passed `Trainer(amp_type={self._amp_type_flag.value!r}, precision='bf16')` but "  # type: ignore
+                f"You passed `Trainer(amp_type={self._amp_type_flag!r}, precision='bf16')` but "  # type: ignore
                 "it's not supported. Try using `amp_type='native'` instead."
             )
         if self._precision_flag in (16, "bf16") and self._amp_type_flag == "apex":

tests/tests_pytorch/plugins/precision/test_deepspeed_precision.py

@@ -20,7 +20,9 @@
 
 
 def test_invalid_precision_with_deepspeed_precision():
-    with pytest.raises(ValueError, match="is not supported. `precision` must be one of"):
+    with pytest.deprecated_call(match=r"amp_type='native'\)` been deprecated in v1.9.0"), pytest.raises(
+        ValueError, match="is not supported. `precision` must be one of"
+    ):
         DeepSpeedPrecisionPlugin(precision=64, amp_type="native")
 
 

tests/tests_pytorch/plugins/precision/test_native_amp.py

@@ -16,14 +16,14 @@
 import pytest
 from torch.optim import Optimizer
 
-from pytorch_lightning.plugins import NativeMixedPrecisionPlugin
+from pytorch_lightning.plugins import MixedPrecisionPlugin
 from pytorch_lightning.utilities import GradClipAlgorithmType
 
 
 def test_clip_gradients():
     """Test that `.clip_gradients()` is a no-op when clipping is disabled."""
     optimizer = Mock(spec=Optimizer)
-    precision = NativeMixedPrecisionPlugin(precision=16, device="cuda:0", scaler=Mock())
+    precision = MixedPrecisionPlugin(precision=16, device="cuda:0", scaler=Mock())
     precision.clip_grad_by_value = Mock()
     precision.clip_grad_by_norm = Mock()
     precision.clip_gradients(optimizer)
@@ -47,7 +47,7 @@ def test_optimizer_amp_scaling_support_in_step_method():
     gradient clipping (example: fused Adam)."""
 
     optimizer = Mock(_step_supports_amp_scaling=True)
-    precision = NativeMixedPrecisionPlugin(precision=16, device="cuda:0", scaler=Mock())
+    precision = MixedPrecisionPlugin(precision=16, device="cuda:0", scaler=Mock())
 
     with pytest.raises(RuntimeError, match="The current optimizer.*does not allow for gradient clipping"):
         precision.clip_gradients(optimizer, clip_val=1.0)

tests/tests_pytorch/plugins/test_amp_plugins.py

@@ -20,13 +20,13 @@
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.demos.boring_classes import BoringModel
-from pytorch_lightning.plugins import ApexMixedPrecisionPlugin, NativeMixedPrecisionPlugin
+from pytorch_lightning.plugins import ApexMixedPrecisionPlugin, MixedPrecisionPlugin
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests_pytorch.conftest import mock_cuda_count
 from tests_pytorch.helpers.runif import RunIf
 
 
-class MyNativeAMP(NativeMixedPrecisionPlugin):
+class MyNativeAMP(MixedPrecisionPlugin):
     pass
 
 
@@ -52,7 +52,7 @@ class MyApexPlugin(ApexMixedPrecisionPlugin):
 @pytest.mark.parametrize(
     "amp,custom_plugin,plugin_cls",
     [
-        ("native", False, NativeMixedPrecisionPlugin),
+        ("native", False, MixedPrecisionPlugin),
         ("native", True, MyNativeAMP),
         pytest.param("apex", False, ApexMixedPrecisionPlugin, marks=RunIf(amp_apex=True)),
         pytest.param("apex", True, MyApexPlugin, marks=RunIf(amp_apex=True)),
@@ -189,9 +189,7 @@ def configure_optimizers(self):
                 torch.optim.SGD(self.layer2.parameters(), lr=0.1),
             ]
 
-    trainer = Trainer(
-        default_root_dir=tmpdir, accelerator="gpu", devices=1, fast_dev_run=1, amp_backend="native", precision=16
-    )
+    trainer = Trainer(default_root_dir=tmpdir, accelerator="gpu", devices=1, fast_dev_run=1, precision=16)
     model = CustomBoringModel()
     trainer.fit(model)
 
@@ -246,7 +244,7 @@ def test_amp_apex_ddp_spawn_fit(amp_level, tmpdir):
 
 def test_cpu_amp_precision_context_manager(tmpdir):
     """Test to ensure that the context manager correctly is set to CPU + bfloat16."""
-    plugin = NativeMixedPrecisionPlugin("bf16", "cpu")
+    plugin = MixedPrecisionPlugin("bf16", "cpu")
     assert plugin.device == "cpu"
     assert plugin.scaler is None
     context_manager = plugin.autocast_context_manager()
@@ -256,16 +254,20 @@ def test_cpu_amp_precision_context_manager(tmpdir):
 
 
 def test_precision_selection_raises(monkeypatch):
-    with pytest.raises(
+    with pytest.deprecated_call(match=r"amp_backend='apex'\)` argument is deprecated"), pytest.raises(
         MisconfigurationException, match=r"precision=16, amp_type='apex'\)` but apex AMP not supported on CPU"
     ):
         Trainer(amp_backend="apex", precision=16)
 
-    with pytest.raises(MisconfigurationException, match=r"amp_type='apex', precision='bf16'\)` but it's not supported"):
+    with pytest.deprecated_call(match=r"amp_backend='apex'\)` argument is deprecated"), pytest.raises(
+        MisconfigurationException, match=r"amp_type='apex', precision='bf16'\)` but it's not supported"
+    ):
         Trainer(amp_backend="apex", precision="bf16")
 
     mock_cuda_count(monkeypatch, 1)
-    with pytest.raises(MisconfigurationException, match="Sharded plugins are not supported with apex"):
+    with pytest.deprecated_call(match=r"amp_backend='apex'\)` argument is deprecated"), pytest.raises(
+        MisconfigurationException, match="Sharded plugins are not supported with apex"
+    ):
         with mock.patch("lightning_lite.accelerators.cuda.is_cuda_available", return_value=True):
             Trainer(amp_backend="apex", precision=16, accelerator="gpu", devices=1, strategy="ddp_fully_sharded")
 
@@ -274,5 +276,5 @@ def test_precision_selection_raises(monkeypatch):
     monkeypatch.setattr(apex, "_APEX_AVAILABLE", False)
     with mock.patch("lightning_lite.accelerators.cuda.is_cuda_available", return_value=True), pytest.raises(
         MisconfigurationException, match="asked for Apex AMP but `apex` is not installed"
-    ), pytest.deprecated_call(match="apex AMP implementation has been deprecated"):
+    ), pytest.deprecated_call(match=r"amp_backend='apex'\)` argument is deprecated"):
         Trainer(amp_backend="apex", precision=16, accelerator="gpu", devices=1)

tests/tests_pytorch/trainer/optimization/test_manual_optimization.py

@@ -65,19 +65,51 @@ def configure_optimizers(self):
 
 
 @pytest.mark.parametrize(
-    "kwargs",
-    [
-        {},
-        pytest.param(
-            {"accelerator": "gpu", "devices": 1, "precision": 16, "amp_backend": "native"}, marks=RunIf(min_cuda_gpus=1)
-        ),
-        pytest.param(
-            {"accelerator": "gpu", "devices": 1, "precision": 16, "amp_backend": "apex"},
-            marks=RunIf(min_cuda_gpus=1, amp_apex=True),
-        ),
-    ],
+    "kwargs", [{}, pytest.param({"accelerator": "gpu", "devices": 1, "precision": 16}, marks=RunIf(min_cuda_gpus=1))]
 )
 def test_multiple_optimizers_manual_no_return(tmpdir, kwargs):
+    class TestModel(ManualOptModel):
+        def training_step(self, batch, batch_idx):
+            # avoid returning a value
+            super().training_step(batch, batch_idx)
+
+        def training_epoch_end(self, outputs):
+            # outputs is empty as training_step does not return
+            # and it is not automatic optimization
+            assert not outputs
+
+    model = TestModel()
+    model.val_dataloader = None
+
+    limit_train_batches = 2
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        limit_train_batches=limit_train_batches,
+        limit_val_batches=2,
+        max_epochs=1,
+        log_every_n_steps=1,
+        enable_model_summary=False,
+        **kwargs,
+    )
+
+    if kwargs.get("precision") == 16:
+        # mock the scaler instead of the optimizer step because it can be skipped with NaNs
+        scaler_step_patch = mock.patch.object(
+            trainer.precision_plugin.scaler, "step", wraps=trainer.precision_plugin.scaler.step
+        )
+        scaler_step = scaler_step_patch.start()
+
+    with mock.patch.object(Strategy, "backward", wraps=trainer.strategy.backward) as bwd_mock:
+        trainer.fit(model)
+    assert bwd_mock.call_count == limit_train_batches * 3
+
+    if kwargs.get("precision") == 16:
+        scaler_step_patch.stop()
+        assert scaler_step.call_count == len(model.optimizers()) * limit_train_batches
+
+
+@RunIf(min_cuda_gpus=1, amp_apex=True)
+def test_multiple_optimizers_manual_no_return_apex(tmpdir):
     apex_optimizer_patches = []
     apex_optimizer_steps = []
 
@@ -92,8 +124,6 @@ def training_epoch_end(self, outputs):
             assert not outputs
 
         def on_train_start(self):
-            if kwargs.get("amp_backend") != "apex":
-                return
             # extremely ugly. APEX patches all the native torch optimizers on `_initialize` which we call on
             # `ApexMixedPrecisionPlugin.dispatch`. Additionally, their replacement `new_step` functions are locally
             # defined so can't even patch those, thus we need to create the mock after APEX has been initialized
@@ -106,19 +136,15 @@ def on_train_start(self):
                 apex_optimizer_steps.append(patch.start())
 
         def on_train_end(self):
-            if kwargs.get("amp_backend") == "apex":
-                for p in apex_optimizer_patches:
-                    p.stop()
+            for p in apex_optimizer_patches:
+                p.stop()
 
     model = TestModel()
     model.val_dataloader = None
 
     limit_train_batches = 2
-    plugins = []
-    if kwargs.get("amp_backend") == "apex":
-        with pytest.deprecated_call(match="apex AMP implementation has been deprecated"):
-            apex_plugin = ApexMixedPrecisionPlugin(amp_level="O2")
-        plugins.append(apex_plugin)
+    with pytest.deprecated_call(match="apex AMP implementation has been deprecated"):
+        plugins = [ApexMixedPrecisionPlugin(amp_level="O2")]
 
     trainer = Trainer(
         default_root_dir=tmpdir,
@@ -128,25 +154,16 @@ def on_train_end(self):
         log_every_n_steps=1,
         enable_model_summary=False,
         plugins=plugins,
-        **kwargs,
+        accelerator="gpu",
+        devices=1,
+        precision=16,
     )
 
-    if kwargs.get("amp_backend") == "native":
-        # mock the scaler instead of the optimizer step because it can be skipped with NaNs
-        scaler_step_patch = mock.patch.object(
-            trainer.precision_plugin.scaler, "step", wraps=trainer.precision_plugin.scaler.step
-        )
-        scaler_step = scaler_step_patch.start()
-
     with mock.patch.object(Strategy, "backward", wraps=trainer.strategy.backward) as bwd_mock:
         trainer.fit(model)
     assert bwd_mock.call_count == limit_train_batches * 3
 
-    if kwargs.get("amp_backend") == "native":
-        scaler_step_patch.stop()
-        assert scaler_step.call_count == len(model.optimizers()) * limit_train_batches
-    if kwargs.get("amp_backend") == "apex":
-        assert [s.call_count for s in apex_optimizer_steps] == [len(model.optimizers())] * limit_train_batches
+    assert [s.call_count for s in apex_optimizer_steps] == [len(model.optimizers())] * limit_train_batches
 
 
 def test_multiple_optimizers_manual_return(tmpdir):
@@ -396,7 +413,6 @@ def on_train_epoch_end(self, *_, **__):
         limit_test_batches=0,
         limit_val_batches=0,
         precision=16,
-        amp_backend="native",
         accelerator="gpu",
         devices=1,
     )
@@ -480,7 +496,6 @@ def log_grad_norm(self, grad_norm_dict):
         log_every_n_steps=1,
         enable_model_summary=False,
         precision=16,
-        amp_backend="native",
         accelerator="gpu",
         devices=1,
         track_grad_norm=2,