skip device id for older pytorch versions

Author: SkafteNicki

src/lightning/fabric/strategies/ddp.py

@@ -41,6 +41,7 @@
     _sync_ddp_if_available,
 )
 from lightning.fabric.utilities.distributed import group as _group
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_3
 from lightning.fabric.utilities.rank_zero import rank_zero_only
 
 _DDP_FORK_ALIASES = (
@@ -212,11 +213,13 @@ def _setup_distributed(self) -> None:
         self._set_world_ranks()
         self._process_group_backend = self._get_process_group_backend()
         assert self.cluster_environment is not None
+        kwargs = {"timeout": self._timeout}
+        if _TORCH_GREATER_EQUAL_2_3:
+            kwargs["device_id"] = self.root_device if self.root_device.type != "cpu" else None
         _init_dist_connection(
             self.cluster_environment,
             self._process_group_backend,
-            timeout=self._timeout,
-            device_id=self.root_device if self.root_device.type != "cpu" else None,
+            **kwargs,
         )
 
     def _get_process_group_backend(self) -> str:

src/lightning/fabric/strategies/fsdp.py

@@ -663,11 +663,13 @@ def _setup_distributed(self) -> None:
         self._set_world_ranks()
         self._process_group_backend = self._get_process_group_backend()
         assert self.cluster_environment is not None
+        kwargs = {"timeout": self._timeout}
+        if _TORCH_GREATER_EQUAL_2_3:
+            kwargs["device_id"] = self.root_device if self.root_device.type != "cpu" else None
         _init_dist_connection(
             self.cluster_environment,
             self._process_group_backend,
-            timeout=self._timeout,
-            device_id=self.root_device if self.root_device.type != "cpu" else None,
+            **kwargs,
         )
 
     def _get_process_group_backend(self) -> str:

src/lightning/fabric/strategies/model_parallel.py

@@ -302,11 +302,13 @@ def _setup_distributed(self) -> None:
         self._set_world_ranks()
         self._process_group_backend = self._get_process_group_backend()
         assert self.cluster_environment is not None
+        kwargs = {"timeout": self._timeout}
+        if _TORCH_GREATER_EQUAL_2_3:
+            kwargs["device_id"] = self.root_device if self.root_device.type != "cpu" else None
         _init_dist_connection(
             self.cluster_environment,
             self._process_group_backend,
-            timeout=self._timeout,
-            device_id=self.root_device if self.root_device.type != "cpu" else None,
+            **kwargs,
         )
 
     def _get_process_group_backend(self) -> str:

src/lightning/pytorch/strategies/ddp.py

@@ -36,7 +36,7 @@
     _sync_ddp_if_available,
 )
 from lightning.fabric.utilities.distributed import group as _group
-from lightning.fabric.utilities.imports import _IS_WINDOWS
+from lightning.fabric.utilities.imports import _IS_WINDOWS, _TORCH_GREATER_EQUAL_2_3
 from lightning.fabric.utilities.optimizer import _optimizers_to_device
 from lightning.fabric.utilities.seed import reset_seed
 from lightning.fabric.utilities.types import ReduceOp
@@ -200,11 +200,13 @@ def setup_distributed(self) -> None:
         self.set_world_ranks()
         self._process_group_backend = self._get_process_group_backend()
         assert self.cluster_environment is not None
+        kwargs = {"timeout": self._timeout}
+        if _TORCH_GREATER_EQUAL_2_3:
+            kwargs["device_id"] = self.root_device if self.root_device.type != "cpu" else None
         _init_dist_connection(
             self.cluster_environment,
             self._process_group_backend,
-            timeout=self._timeout,
-            device_id=self.root_device if self.root_device.type != "cpu" else None,
+            **kwargs,
         )
 
     def _get_process_group_backend(self) -> str:

src/lightning/pytorch/strategies/fsdp.py

@@ -61,7 +61,7 @@
     _sync_ddp_if_available,
 )
 from lightning.fabric.utilities.distributed import group as _group
-from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_2
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_2, _TORCH_GREATER_EQUAL_2_3
 from lightning.fabric.utilities.init import _has_meta_device_parameters_or_buffers
 from lightning.fabric.utilities.load import _lazy_load, _materialize_tensors
 from lightning.fabric.utilities.optimizer import _optimizers_to_device
@@ -260,11 +260,13 @@ def setup_environment(self) -> None:
 
         self._process_group_backend = self._get_process_group_backend()
         assert self.cluster_environment is not None
+        kwargs = {"timeout": self._timeout}
+        if _TORCH_GREATER_EQUAL_2_3:
+            kwargs["device_id"] = self.root_device if self.root_device.type != "cpu" else None
         _init_dist_connection(
             self.cluster_environment,
             self._process_group_backend,
-            timeout=self._timeout,
-            device_id=self.root_device if self.root_device.type != "cpu" else None,
+            **kwargs,
         )
 
         # if 'device_mesh' in the `kwargs` is provided as a tuple, update it into the `DeviceMesh` object here

src/lightning/pytorch/strategies/model_parallel.py

@@ -39,7 +39,7 @@
     _sync_ddp_if_available,
 )
 from lightning.fabric.utilities.distributed import group as _group
-from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_4
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_3, _TORCH_GREATER_EQUAL_2_4
 from lightning.fabric.utilities.init import _materialize_distributed_module
 from lightning.fabric.utilities.load import _METADATA_FILENAME
 from lightning.fabric.utilities.optimizer import _optimizers_to_device
@@ -350,11 +350,13 @@ def _setup_distributed(self) -> None:
         self.set_world_ranks()
         self._process_group_backend = self._get_process_group_backend()
         assert self.cluster_environment is not None
+        kwargs = {"timeout": self._timeout}
+        if _TORCH_GREATER_EQUAL_2_3:
+            kwargs["device_id"] = self.root_device if self.root_device.type != "cpu" else None
         _init_dist_connection(
             self.cluster_environment,
             self._process_group_backend,
-            timeout=self._timeout,
-            device_id=self.root_device if self.root_device.type != "cpu" else None,
+            **kwargs,
         )
 
     def _get_process_group_backend(self) -> str: