support customized ttp and accelerator

Author: four4fish

pytorch_lightning/accelerators/gpu.py

@@ -41,6 +41,7 @@ def setup_environment(self, root_device: torch.device) -> None:
         torch.cuda.set_device(root_device)
 
     def setup(self, trainer: "pl.Trainer") -> None:
+        # TODO refactor input from trainer to local_rank @four4fish
         self.set_nvidia_flags(trainer.local_rank)
         # clear cache before training
         torch.cuda.empty_cache()

pytorch_lightning/accelerators/tpu.py

@@ -15,12 +15,7 @@
 
 import torch
 
-# import pytorch_lightning as pl
 from pytorch_lightning.accelerators.accelerator import Accelerator
-
-# from pytorch_lightning.plugins.precision import TPUPrecisionPlugin
-# from pytorch_lightning.plugins.training_type.single_tpu import SingleTPUPlugin
-# from pytorch_lightning.plugins.training_type.tpu_spawn import TPUSpawnPlugin
 from pytorch_lightning.utilities import _XLA_AVAILABLE
 
 if _XLA_AVAILABLE:
@@ -30,25 +25,6 @@
 class TPUAccelerator(Accelerator):
     """Accelerator for TPU devices."""
 
-    # def setup(self, trainer: "pl.Trainer") -> None:
-    #     """
-    #     Raises:
-    #         ValueError:
-    #             If the precision or training type plugin are unsupported.
-    #     """
-    #     if not isinstance(self.training_type_plugin.precision_plugin, TPUPrecisionPlugin):
-    #         # this configuration should have been avoided in the accelerator connector
-    #         raise ValueError(
-    #             f"The `TPUAccelerator` can only be used with a `TPUPrecisionPlugin`,"
-    #             f" found: {self.training_type_plugin.precision_plugin}."
-    #         )
-    #     if not isinstance(self.training_type_plugin, (SingleTPUPlugin, TPUSpawnPlugin)):
-    #         raise ValueError(
-    #             "The `TPUAccelerator` can only be used with a `SingleTPUPlugin` or `TPUSpawnPlugin,"
-    #             f" found {self.training_type_plugin}."
-    #         )
-    #     return super().setup(trainer)
-
     def get_device_stats(self, device: Union[str, torch.device]) -> Dict[str, Any]:
         """Gets stats for the given TPU device.
 

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -178,6 +178,8 @@ def __init__(
         self.training_type_plugin = self.final_training_type_plugin()
         self.accelerator = self.training_type_plugin.accelerator
 
+        self._check_tpu_mis_config()
+
         # benchmarking
         # TODO: should this be moved to GPU accelerator?
         torch.backends.cudnn.benchmark = self.benchmark
@@ -405,12 +407,19 @@ def final_training_type_plugin(self) -> TrainingTypePlugin:
         # attach checkpoint plugin to the training type plugin
         if self._checkpoint_io is not None:
             self._training_type_plugin.checkpoint_io = self._checkpoint_io
-        precision_plugin = self.precision_plugin
-        if precision_plugin is not None:
-            self._training_type_plugin._precision_plugin = precision_plugin
+        if (
+            (hasattr(self.strategy, "precision_plugin") and self.precision_plugin is None)
+            or not hasattr(self.strategy, "precision_plugin")
+        ):
+            precision_plugin = self.precision_plugin
+            if precision_plugin is not None:
+                self._training_type_plugin._precision_plugin = precision_plugin
         self._training_type_plugin_resolved = True
-
-        self._training_type_plugin.accelerator = self.select_accelerator()
+        if (
+            (hasattr(self.strategy, "accelerator") and self.strategy.accelerator is None)
+            or not hasattr(self.strategy, "accelerator")
+        ):
+            self._training_type_plugin.accelerator = self.select_accelerator()
         return self._training_type_plugin
 
     @property
@@ -1016,3 +1025,18 @@ def _is_slurm_managing_tasks(self) -> bool:
         total_requested_devices = (self.num_gpus or self.num_processes) * self.num_nodes
         num_slurm_tasks = int(os.environ["SLURM_NTASKS"], 0)
         return num_slurm_tasks == total_requested_devices
+
+    def _check_tpu_mis_config(self) -> None:
+        # TODO moved from TPUAccelerator when refactor accelerator. Revisit when refactor
+        # accelerator_connector @four4fish
+        if isinstance(self.accelerator, TPUAccelerator):
+            if not isinstance(self.training_type_plugin.precision_plugin, TPUPrecisionPlugin):
+                raise ValueError(
+                    f"The `TPUAccelerator` can only be used with a `TPUPrecisionPlugin`,"
+                    f" found: {self.training_type_plugin.precision_plugin}."
+                )
+            if not isinstance(self.training_type_plugin, (SingleTPUPlugin, TPUSpawnPlugin)):
+                raise ValueError(
+                    "The `TPUAccelerator` can only be used with a `SingleTPUPlugin` or `TPUSpawnPlugin,"
+                    f" found {self.training_type_plugin}."
+                )

tests/accelerators/test_accelerator_connector.py

@@ -397,7 +397,14 @@ def creates_processes_externally(self) -> bool:
 
 @mock.patch.dict(
     os.environ,
-    {"SLURM_NTASKS": "2", "SLURM_JOB_NAME": "SOME_NAME", "SLURM_NODEID": "0", "LOCAL_RANK": "0", "SLURM_LOCALID": "0"},
+    {
+        "SLURM_NTASKS": "2",
+        "SLURM_JOB_NAME": "SOME_NAME",
+        "SLURM_NODEID": "0",
+        "LOCAL_RANK": "0",
+        "SLURM_PROCID": "0",
+        "SLURM_LOCALID": "0",
+    },
 )
 @mock.patch("torch.cuda.device_count", return_value=0)
 @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True)
@@ -408,28 +415,29 @@ class Accel(Accelerator):
     class Prec(PrecisionPlugin):
         pass
 
-    class TrainTypePlugin(SingleDevicePlugin):
+    class TrainTypePlugin(DDPPlugin):
         pass
 
+    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)
+    assert isinstance(trainer.precision_plugin, Prec)
+    assert trainer._accelerator_connector.training_type_plugin is ttp
+
+    class DistributedPlugin(DDPPlugin):
+        pass
 
-#     ttp = TrainTypePlugin(device=torch.device("cpu"))
-#     accelerator = Accel(training_type_plugin=ttp, precision_plugin=Prec())
-#     trainer = Trainer(accelerator=accelerator, fast_dev_run=True, num_processes=2)
-#     assert isinstance(trainer.accelerator, Accel)
-#     assert isinstance(trainer.training_type_plugin, TrainTypePlugin)
-#     assert isinstance(trainer.precision_plugin, Prec)
-#     assert trainer._accelerator_connector.training_type_plugin is ttp
-
-#     class DistributedPlugin(DDPPlugin):
-#         pass
-
-#     ttp = DistributedPlugin()
-#     accelerator = Accel(training_type_plugin=ttp, precision_plugin=Prec())
-#     trainer = Trainer(accelerator=accelerator, fast_dev_run=True, num_processes=2)
-#     assert isinstance(trainer.accelerator, Accel)
-#     assert isinstance(trainer.training_type_plugin, DistributedPlugin)
-#     assert isinstance(trainer.precision_plugin, Prec)
-#     assert trainer._accelerator_connector.training_type_plugin is ttp
+    ttp = DistributedPlugin(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, DistributedPlugin)
+    assert isinstance(trainer.precision_plugin, Prec)
+    assert trainer._accelerator_connector.training_type_plugin is ttp
 
 
 @mock.patch.dict(

tests/accelerators/test_tpu.py

@@ -288,28 +288,27 @@ def forward(self, x):
 
 
 def test_tpu_invalid_raises():
-    # TODO move TPUAccelerator() and CPUAccelerator() setup() misconfig logic into strategies
     training_type_plugin = TPUSpawnPlugin(accelerator=TPUAccelerator(), precision_plugin=Mock())
-    # with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `TPUPrecisionPlugin"):
-    #     training_type_plugin.setup(Mock())
+    with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `TPUPrecisionPlugin"):
+        Trainer(strategy=training_type_plugin)
 
     training_type_plugin = DDPPlugin(accelerator=TPUAccelerator(), precision_plugin=TPUPrecisionPlugin())
-    # with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUPlugin`):
-    #     training_type_plugin.setup(Mock())
+    with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUPlugin`"):
+        Trainer(strategy=training_type_plugin)
 
 
 def test_tpu_invalid_raises_set_precision_with_strategy():
     accelerator = TPUAccelerator()
     training_type_plugin = TPUSpawnPlugin(accelerator=accelerator, precision_plugin=object())
-    # with pytest.raises(ValueError, match="`TPUAccelerator` can only be used with a `TPUPrecisionPlugin`"):
-    #     training_type_plugin.setup(object())
+    with pytest.raises(ValueError, match="`TPUAccelerator` can only be used with a `TPUPrecisionPlugin`"):
+        Trainer(strategy=training_type_plugin)
 
     accelerator = TPUAccelerator()
     training_type_plugin = DDPPlugin(accelerator=accelerator, precision_plugin=TPUPrecisionPlugin())
-    # with pytest.raises(
-    #     ValueError, match="TPUAccelerator` can only be used with a `SingleTPUPlugin` or `TPUSpawnPlugin"
-    # ):
-    #     training_type_plugin.setup(object())
+    with pytest.raises(
+        ValueError, match="The `TPUAccelerator` can only be used with a `SingleTPUPlugin` or `TPUSpawnPlugin"
+    ):
+        Trainer(strategy=training_type_plugin)
 
 
 @RunIf(tpu=True)