fix gpu setup environment

Author: four4fish

docs/source/extensions/accelerators.rst

@@ -26,7 +26,9 @@ One to handle differences from the training routine and one to handle different
     from pytorch_lightning.plugins import NativeMixedPrecisionPlugin, DDPPlugin
 
     accelerator = GPUAccelerator()
-    trainer = Trainer(accelerator=accelerator)
+    precision_plugin=NativeMixedPrecisionPlugin(precision=16, device="cuda")
+    training_type_plugin=DDPPlugin(accelerator=accelerator, precision_plugin=precision_plugin)
+    trainer = Trainer(strategy=training_type_plugin)
 
 
 We expose Accelerators and Plugins mainly for expert users who want to extend Lightning to work with new

pytorch_lightning/accelerators/accelerator.py

@@ -30,7 +30,7 @@ class Accelerator:
     - IPU
     """
 
-    def setup_environment(self) -> None:
+    def setup_environment(self, root_device: torch.device) -> None:
         """Setup any processes or distributed connections.
 
         This is called before the LightningModule/DataModule setup hook which allows the user to access the accelerator

pytorch_lightning/accelerators/cpu.py

@@ -15,29 +15,21 @@
 
 import torch
 
-# import pytorch_lightning as pl
 from pytorch_lightning.accelerators.accelerator import Accelerator
-
-# from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 
 class CPUAccelerator(Accelerator):
     """Accelerator for CPU devices."""
 
-    # @property
-    # def root_device(self):
-    #     return torch.device("cpu")
-
-    # def setup(self, trainer: "pl.Trainer") -> None:
-    #     """
-    #     Raises:
-    #         MisconfigurationException:
-    #             If the selected device is not CPU.
-    #     """
-    #     if "cpu" not in str(self.root_device):
-    #         raise MisconfigurationException(f"Device should be CPU, got {self.root_device} instead.")
-
-    #     return super().setup(trainer)
+    def setup_environment(self, root_device: torch.device) -> None:
+        """
+        Raises:
+            MisconfigurationException:
+                If the selected device is not CPU.
+        """
+        if "cpu" not in str(root_device):
+            raise MisconfigurationException(f"Device should be CPU, got {root_device} instead.")
 
     def get_device_stats(self, device: Union[str, torch.device]) -> Dict[str, Any]:
         """CPU device stats aren't supported yet."""

pytorch_lightning/accelerators/gpu.py

@@ -21,8 +21,7 @@
 
 import pytorch_lightning as pl
 from pytorch_lightning.accelerators.accelerator import Accelerator
-
-# from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_8
 
 _log = logging.getLogger(__name__)
@@ -31,19 +30,18 @@
 class GPUAccelerator(Accelerator):
     """Accelerator for GPU devices."""
 
-    # def setup_environment(self) -> None:
-    #     """
-    #     Raises:
-    #         MisconfigurationException:
-    #             If the selected device is not GPU.
-    #     """
-    #     if "cuda" not in str(self.root_device):
-    #         raise MisconfigurationException(f"Device should be GPU, got {self.root_device} instead")
-    #     torch.cuda.set_device(self.root_device)
+    def setup_environment(self, root_device: torch.device) -> None:
+        """
+        Raises:
+            MisconfigurationException:
+                If the selected device is not GPU.
+        """
+        if "cuda" not in str(root_device):
+            raise MisconfigurationException(f"Device should be GPU, got {self.root_device} instead")
+        torch.cuda.set_device(root_device)
 
     def setup(self, trainer: "pl.Trainer") -> None:
         self.set_nvidia_flags(trainer.local_rank)
-        super().setup(trainer)
         # clear cache before training
         torch.cuda.empty_cache()
 

pytorch_lightning/plugins/training_type/training_type_plugin.py

@@ -93,7 +93,7 @@ def setup_environment(self) -> None:
         This is called before the LightningModule/DataModule setup hook which allows the user to access the accelerator
         environment before setup is complete.
         """
-        self.accelerator.setup_environment()
+        self.accelerator.setup_environment(self.root_device)
 
     def setup_optimizers(self, trainer: "pl.Trainer") -> None:
         """Creates optimizers and schedulers.