Update deepspeed.py

Lightning-AI · Dec 7, 2024 · 7638439 · 7638439
1 parent ad74bb3
commit 7638439
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/lightning/pytorch/strategies/deepspeed.py b/src/lightning/pytorch/strategies/deepspeed.py
@@ -366,7 +366,7 @@ 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