Move ipu precision flag check to IPUPrecisionPlugin init (#12148)

Lightning-AI · Mar 5, 2022 · 91052dc · 91052dc
1 parent b5fe056
commit 91052dc
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 9 deletions.
diff --git a/pytorch_lightning/plugins/precision/ipu.py b/pytorch_lightning/plugins/precision/ipu.py
@@ -27,9 +27,20 @@
 
 
 class IPUPrecisionPlugin(PrecisionPlugin):
-    """Precision plugin for IPU integration."""
+    """Precision plugin for IPU integration.
+
+    Raises:
+        ValueError:
+            If the precision is neither 16 nor 32.
+    """
 
     def __init__(self, precision: int) -> None:
+        supported_precision_values = (16, 32)
+        if precision not in supported_precision_values:
+            raise ValueError(
+                f"`Trainer(accelerator='ipu', precision={precision!r})` is not supported."
+                f" `precision` must be one of: {supported_precision_values}."
+            )
         super().__init__()
         self.precision = precision
 

diff --git a/pytorch_lightning/trainer/connectors/accelerator_connector.py b/pytorch_lightning/trainer/connectors/accelerator_connector.py
@@ -685,12 +685,6 @@ def _check_and_init_precision(self) -> PrecisionPlugin:
 
     def _validate_precision_choice(self) -> None:
         """Validate the combination of choices for precision, AMP type, and accelerator."""
-        # TODO: change exception type to ImpactableConfigurationException
-        if isinstance(self.accelerator, IPUAccelerator):
-            if self._precision_flag not in (16, 32):
-                raise MisconfigurationException(
-                    f"`Trainer(accelerator='ipu', precision={self._precision_flag!r})` is not supported."
-                )
         if isinstance(self.accelerator, TPUAccelerator):
             if self._precision_flag == 64:
                 raise MisconfigurationException(

diff --git a/tests/accelerators/test_accelerator_connector.py b/tests/accelerators/test_accelerator_connector.py
@@ -929,9 +929,9 @@ def test_unsupported_ipu_choice(mock_ipu_acc_avail, monkeypatch):
 
     monkeypatch.setattr(imports, "_IPU_AVAILABLE", True)
     monkeypatch.setattr(ipu, "_IPU_AVAILABLE", True)
-    with pytest.raises(MisconfigurationException, match=r"accelerator='ipu', precision='bf16'\)` is not supported"):
+    with pytest.raises(ValueError, match=r"accelerator='ipu', precision='bf16'\)` is not supported"):
         Trainer(accelerator="ipu", precision="bf16")
-    with pytest.raises(MisconfigurationException, match=r"accelerator='ipu', precision=64\)` is not supported"):
+    with pytest.raises(ValueError, match=r"accelerator='ipu', precision=64\)` is not supported"):
         Trainer(accelerator="ipu", precision=64)