fix tpu error check

Author: four4fish

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -177,6 +177,7 @@ def __init__(
         self._training_type_plugin_resolved = False
         self.training_type_plugin = self.final_training_type_plugin()
         self.accelerator = self.training_type_plugin.accelerator
+        self.precision_plugin = self.training_type_plugin._precision_plugin
 
         self._check_tpu_mis_config()
 
@@ -391,8 +392,7 @@ def handle_given_plugins(self) -> None:
     def accelerator_types(self) -> List[str]:
         return ["auto"] + list(_AcceleratorType)
 
-    @property
-    def precision_plugin(self) -> PrecisionPlugin:
+    def final_precision_plugin(self) -> PrecisionPlugin:
         if self._precision_plugin is None:
             self._precision_plugin = self.select_precision_plugin()
         return self._precision_plugin
@@ -408,16 +408,28 @@ def final_training_type_plugin(self) -> TrainingTypePlugin:
         if self._checkpoint_io is not None:
             self._training_type_plugin.checkpoint_io = self._checkpoint_io
         if (
-            (hasattr(self.strategy, "precision_plugin") and self.precision_plugin is None)
-            or not hasattr(self.strategy, "precision_plugin")
+            # handle custom strategy with custom precision
+            (
+                isinstance(self.strategy, TrainingTypePlugin) and (
+                    self.strategy.precision_plugin is None
+                    or not isinstance(self.strategy.precision_plugin, PrecisionPlugin)
+                )
+            )
+            or not isinstance(self.strategy, TrainingTypePlugin)
         ):
-            precision_plugin = self.precision_plugin
+            precision_plugin = self.final_precision_plugin()
             if precision_plugin is not None:
                 self._training_type_plugin._precision_plugin = precision_plugin
         self._training_type_plugin_resolved = True
         if (
-            (hasattr(self.strategy, "accelerator") and self.strategy.accelerator is None)
-            or not hasattr(self.strategy, "accelerator")
+            # handle custom strategy with custom accelerator
+            (
+                isinstance(self.strategy, TrainingTypePlugin) and (
+                    self.strategy.accelerator is None
+                    or not isinstance(self.strategy.accelerator, Accelerator)
+                )
+            )
+            or not isinstance(self.strategy, TrainingTypePlugin)
         ):
             self._training_type_plugin.accelerator = self.select_accelerator()
         return self._training_type_plugin
@@ -790,12 +802,6 @@ def resolve_training_type_plugin(self, training_type: TrainingTypePlugin) -> Tra
     def select_accelerator(self) -> Accelerator:
         if isinstance(self.distributed_backend, Accelerator):
             # custom accelerator from user
-            if self._precision_plugin is not None or self._training_type_plugin is not None:
-                # plugins also specified by user
-                rank_zero_warn(
-                    "Specified `Precision` and `TrainingType` plugins will be ignored,"
-                    " since an `Accelerator` instance was provided."
-                )
             return self.distributed_backend
 
         if self.use_gpu:

tests/accelerators/test_accelerator_connector.py

@@ -418,11 +418,7 @@ class Prec(PrecisionPlugin):
     class TrainTypePlugin(DDPPlugin):
         pass
 
-    ttp = TrainTypePlugin(
-        device=torch.device("cpu"),
-        accelerator=Accel(),
-        precision_plugin=Prec()
-    )
+    ttp = TrainTypePlugin(device=torch.device("cpu"), accelerator=Accel(), precision_plugin=Prec())
     trainer = Trainer(strategy=ttp, fast_dev_run=True, num_processes=2)
     assert isinstance(trainer.accelerator, Accel)
     assert isinstance(trainer.training_type_plugin, TrainTypePlugin)
@@ -1038,10 +1034,13 @@ def test_unsupported_tpu_choice(monkeypatch):
     with pytest.raises(MisconfigurationException, match=r"accelerator='tpu', precision=64\)` is not implemented"):
         Trainer(accelerator="tpu", precision=64)
 
-    with pytest.warns(UserWarning, match=r"accelerator='tpu', precision=16\)` but native AMP is not supported"):
-        Trainer(accelerator="tpu", precision=16)
-    with pytest.warns(UserWarning, match=r"accelerator='tpu', precision=16\)` but apex AMP is not supported"):
-        Trainer(accelerator="tpu", precision=16, amp_backend="apex")
+    with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUPlugin`"):
+        with pytest.warns(UserWarning, match=r"accelerator='tpu', precision=16\)` but native AMP is not supported"):
+            Trainer(accelerator="tpu", precision=16)
+
+    with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUPlugin`"):
+        with pytest.warns(UserWarning, match=r"accelerator='tpu', precision=16\)` but apex AMP is not supported"):
+            Trainer(accelerator="tpu", precision=16, amp_backend="apex")
 
 
 def test_unsupported_ipu_choice(monkeypatch):