Update deepspeed.py

jedyang97 · web-flow · commit 763843982b4f · 2024-12-07T15:43:34.000-05:00
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

Original file line number	Diff line number	Diff line change
`@@ -366,7 +366,7 @@ def _init_deepspeed_distributed(self) -> None:`
`366`	`366`	`f"MEMBER: {self.global_rank + 1}/{self.world_size}"`
`367`	`367`	`)`
`368`	`368`	`self._process_group_backend = self._get_process_group_backend()`
`369`		`- deepspeed.init_distributed(self._process_group_backend, distributed_port=self.cluster_environment.main_port)`
	`369`	`+ deepspeed.init_distributed(self._process_group_backend, distributed_port=self.cluster_environment.main_port, timeout=self._timeout)`
`370`	`370`
`371`	`371`	`def _set_node_environment_variables(self) -> None:`
`372`	`372`	`assert self.cluster_environment is not None`