Add timeout to DeepSpeedStrategy

src/lightning/fabric/strategies/deepspeed.py

@@ -97,6 +97,7 @@ def __init__(
         load_full_weights: bool = False,
         precision: Optional[Precision] = None,
         process_group_backend: Optional[str] = None,
+        **kwargs: Any,
     ) -> None:
         """Provides capabilities to run training using the DeepSpeed library, with training optimizations for large
         billion parameter models. `For more information: https://pytorch-
@@ -239,6 +240,7 @@ def __init__(
             cluster_environment=cluster_environment,
             precision=precision,
             process_group_backend=process_group_backend,
+            **kwargs,
         )
         self._backward_sync_control = None  # DeepSpeed handles gradient accumulation internally
 
@@ -648,7 +650,9 @@ def _init_deepspeed_distributed(self) -> None:
                 f"MEMBER: {self.global_rank + 1}/{self.world_size}"
             )
         self._process_group_backend = self._get_process_group_backend()
-        deepspeed.init_distributed(self._process_group_backend, distributed_port=self.cluster_environment.main_port)
+        deepspeed.init_distributed(
+            self._process_group_backend, distributed_port=self.cluster_environment.main_port, timeout=self._timeout
+        )
 
     def _set_node_environment_variables(self) -> None:
         assert self.cluster_environment is not None

src/lightning/pytorch/strategies/deepspeed.py

@@ -119,6 +119,7 @@ def __init__(
         load_full_weights: bool = False,
         precision_plugin: Optional[Precision] = None,
         process_group_backend: Optional[str] = None,
+        **kwargs: Any,
     ) -> None:
         """Provides capabilities to run training using the DeepSpeed library, with training optimizations for large
         billion parameter models. `For more information: https://pytorch-
@@ -263,6 +264,7 @@ def __init__(
             cluster_environment=cluster_environment,
             precision_plugin=precision_plugin,
             process_group_backend=process_group_backend,
+            **kwargs,
         )
 
         self.config = self._load_config(config)
@@ -364,7 +366,9 @@ def _init_deepspeed_distributed(self) -> None:
                 f"MEMBER: {self.global_rank + 1}/{self.world_size}"
             )
         self._process_group_backend = self._get_process_group_backend()
-        deepspeed.init_distributed(self._process_group_backend, distributed_port=self.cluster_environment.main_port)
+        deepspeed.init_distributed(
+            self._process_group_backend, distributed_port=self.cluster_environment.main_port, timeout=self._timeout
+        )
 
     def _set_node_environment_variables(self) -> None:
         assert self.cluster_environment is not None