set_device before init_process_group (#19184)

Author: carmocca

src/lightning/fabric/accelerators/cuda.py

@@ -358,7 +358,7 @@ def _is_ampere_or_later(device: Optional[torch.device] = None) -> bool:
 
 @lru_cache(1)  # show the warning only ever once
 def _check_cuda_matmul_precision(device: torch.device) -> None:
-    if not _is_ampere_or_later(device):
+    if not torch.cuda.is_available() or not _is_ampere_or_later(device):
         return
     # check that the user hasn't changed the precision already, this works for both `allow_tf32 = True` and
     # `set_float32_matmul_precision`

src/lightning/fabric/strategies/ddp.py

@@ -116,8 +116,8 @@ def _configure_launcher(self) -> None:
 
     @override
     def setup_environment(self) -> None:
-        self._setup_distributed()
         super().setup_environment()
+        self._setup_distributed()
 
     @override
     def setup_module(self, module: Module) -> DistributedDataParallel:

src/lightning/fabric/strategies/deepspeed.py

@@ -604,12 +604,16 @@ def _initialize_engine(
         return deepspeed_engine, deepspeed_optimizer
 
     @override
-    def _setup_distributed(self) -> None:
+    def setup_environment(self) -> None:
         if not isinstance(self.accelerator, CUDAAccelerator):
             raise RuntimeError(
                 f"The DeepSpeed strategy is only supported on CUDA GPUs but `{self.accelerator.__class__.__name__}`"
                 " is used."
             )
+        super().setup_environment()
+
+    @override
+    def _setup_distributed(self) -> None:
         assert self.parallel_devices is not None
         _validate_device_index_selection(self.parallel_devices)
         reset_seed()

src/lightning/fabric/strategies/fsdp.py

@@ -251,8 +251,8 @@ def _configure_launcher(self) -> None:
 
     @override
     def setup_environment(self) -> None:
-        self._setup_distributed()
         super().setup_environment()
+        self._setup_distributed()
 
     @override
     def setup_module_and_optimizers(

src/lightning/pytorch/strategies/ddp.py

@@ -150,8 +150,8 @@ def _configure_launcher(self) -> None:
 
     @override
     def setup_environment(self) -> None:
-        self.setup_distributed()
         super().setup_environment()
+        self.setup_distributed()
 
     @override
     def setup(self, trainer: "pl.Trainer") -> None:

src/lightning/pytorch/strategies/deepspeed.py

@@ -328,12 +328,16 @@ def _load_config(self, config: Optional[Union[_PATH, Dict[str, Any]]]) -> Option
         return config
 
     @override
-    def setup_distributed(self) -> None:
+    def setup_environment(self) -> None:
         if not isinstance(self.accelerator, CUDAAccelerator):
             raise RuntimeError(
                 f"The DeepSpeed strategy is only supported on CUDA GPUs but `{self.accelerator.__class__.__name__}`"
                 " is used."
             )
+        super().setup_environment()
+
+    @override
+    def setup_distributed(self) -> None:
         assert self.parallel_devices is not None
         _validate_device_index_selection(self.parallel_devices)
         reset_seed()

src/lightning/pytorch/strategies/fsdp.py

@@ -248,6 +248,7 @@ def lightning_restore_optimizer(self) -> bool:
 
     @override
     def setup_environment(self) -> None:
+        super().setup_environment()
         log.debug(f"{self.__class__.__name__}: setting up distributed...")
         reset_seed()
 
@@ -257,7 +258,6 @@ def setup_environment(self) -> None:
         self._process_group_backend = self._get_process_group_backend()
         assert self.cluster_environment is not None
         _init_dist_connection(self.cluster_environment, self._process_group_backend, timeout=self._timeout)
-        super().setup_environment()
 
     def _get_process_group_backend(self) -> str:
         return self._process_group_backend or _get_default_process_group_backend_for_device(self.root_device)

tests/tests_fabric/accelerators/test_cuda.py

@@ -89,7 +89,8 @@ def test_force_nvml_based_cuda_check():
 
 @mock.patch("torch.cuda.get_device_capability", return_value=(10, 1))
 @mock.patch("torch.cuda.get_device_name", return_value="Z100")
-def test_tf32_message(_, __, caplog, monkeypatch):
+@mock.patch("torch.cuda.is_available", return_value=True)
+def test_tf32_message(_, __, ___, caplog, monkeypatch):
     # for some reason, caplog doesn't work with our rank_zero_info utilities
     monkeypatch.setattr(lightning.fabric.accelerators.cuda, "rank_zero_info", logging.info)
 

tests/tests_fabric/strategies/test_deepspeed.py

@@ -400,10 +400,12 @@ def test_validate_parallel_devices_indices(device_indices):
     DeepSpeed doesn't support it and needs the index to match to the local rank of the process.
 
     """
+    accelerator = Mock(spec=CUDAAccelerator)
     strategy = DeepSpeedStrategy(
-        accelerator=CUDAAccelerator(), parallel_devices=[torch.device("cuda", i) for i in device_indices]
+        accelerator=accelerator, parallel_devices=[torch.device("cuda", i) for i in device_indices]
     )
     with pytest.raises(
         RuntimeError, match=escape(f"device indices {device_indices!r} don't match the local rank values of processes")
     ):
         strategy.setup_environment()
+    accelerator.setup_device.assert_called_once_with(torch.device("cuda", device_indices[0]))

tests/tests_pytorch/strategies/test_deepspeed.py

@@ -18,7 +18,7 @@
 from re import escape
 from typing import Any, Dict
 from unittest import mock
-from unittest.mock import ANY
+from unittest.mock import ANY, Mock
 
 import pytest
 import torch
@@ -1264,13 +1264,15 @@ def test_validate_parallel_devices_indices(device_indices):
     DeepSpeed doesn't support it and needs the index to match to the local rank of the process.
 
     """
+    accelerator = Mock(spec=CUDAAccelerator)
     strategy = DeepSpeedStrategy(
-        accelerator=CUDAAccelerator(), parallel_devices=[torch.device("cuda", i) for i in device_indices]
+        accelerator=accelerator, parallel_devices=[torch.device("cuda", i) for i in device_indices]
     )
     with pytest.raises(
         RuntimeError, match=escape(f"device indices {device_indices!r} don't match the local rank values of processes")
     ):
         strategy.setup_environment()
+    accelerator.setup_device.assert_called_once_with(torch.device("cuda", device_indices[0]))
 
 
 @RunIf(min_cuda_gpus=2, standalone=True, deepspeed=True, bf16_cuda=True)