Update deepspeed.py

jedyang97 · web-flow · commit 5872eb911c61 · 2024-12-07T15:44:17.000-05:00
diff --git a/src/lightning/fabric/strategies/deepspeed.py b/src/lightning/fabric/strategies/deepspeed.py
@@ -650,7 +650,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
`@@ -650,7 +650,7 @@ def _init_deepspeed_distributed(self) -> None:`
`650`	`650`	`f"MEMBER: {self.global_rank + 1}/{self.world_size}"`
`651`	`651`	`)`
`652`	`652`	`self._process_group_backend = self._get_process_group_backend()`
`653`		`- deepspeed.init_distributed(self._process_group_backend, distributed_port=self.cluster_environment.main_port)`
	`653`	`+ deepspeed.init_distributed(self._process_group_backend, distributed_port=self.cluster_environment.main_port, timeout=self._timeout)`
`654`	`654`
`655`	`655`	`def _set_node_environment_variables(self) -> None:`
`656`	`656`	`assert self.cluster_environment is not None`